Add function to find predicates used by a Prolog program

1 files changed, 18 insertions, 0 deletions

diff --git a/prolog/util.py b/prolog/util.py
index 11a6f94..e69d4d5 100644
--- a/prolog/util.py
+++ b/prolog/util.py
@@ -224,6 +224,24 @@ def map_vars(a, b, tokens, variables):
 
     return [t if t.type != 'VARIABLE' else t.clone(val=mapping[t.val]) for t in b]
 
+# Return a set of predicate names (e.g. conc/3) used in [program].
+def used_predicates(program):
+    predicates = set()
+    def walk(tree):
+        if isinstance(tree, Tree):
+            if tree.label() == 'term' and len(tree) >= 3 and \
+                    isinstance(tree[0], Tree) and tree[0].label() == 'functor':
+                if len(tree) == 3:
+                    predicates.add('{}/0'.format(tree[0][0]))
+                else:
+                    predicates.add('{}/{}'.format(tree[0][0], (len(tree[2])+1)//2))
+            for subtree in tree:
+                walk(subtree)
+    tree = parse(program)
+    if tree is not None:
+        walk(tree)
+    return predicates
+
 # Basic sanity check.
 if __name__ == '__main__':
     var_names = {}