Simplify prolog.util.map_vars

Return a single mapping instead of generating all possible permutations.
Only the first solution was used anyway, with little effect on found
solutions.

1 files changed, 9 insertions, 19 deletions

diff --git a/prolog/util.py b/prolog/util.py
index 85589b5..b7805db 100644
--- a/prolog/util.py
+++ b/prolog/util.py
@@ -122,27 +122,17 @@ def map_vars(before, after, line, rule):
                 new_index += 1
             mapping[formal_name] = actual_name
 
-    remaining_formal = set([t[1] for t in after if t[0] == 'VARIABLE' and t[1] not in mapping.keys()])
-    remaining_actual = set([t[1] for t in rule if t[0] == 'VARIABLE' and t[1] != '_' and t[1] not in mapping.values()])
+    remaining_formal = [t[1] for t in after if t[0] == 'VARIABLE' and t[1] not in mapping.keys()]
+    remaining_actual = [t[1] for t in rule if t[0] == 'VARIABLE' and t[1] != '_' and t[1] not in mapping.values()]
 
-    #for i in range(0, len(remaining_formal)-len(remaining_actual)+2):
-    for var in remaining_formal:
-        remaining_actual.add('New'+str(new_index))
+    while len(remaining_actual) < len(remaining_formal):
+        remaining_actual.append('New'+str(new_index))
         new_index += 1
 
-    # cthulhu-inspired horrors continue...
-    # get all possible mappings of remaining transform var. names to actual names
-    def match(a, b):
-        for bc in itertools.combinations(b, min(len(a), len(b))):
-            for bp in itertools.permutations(bc):
-                for ac in itertools.combinations(a, len(bc)):
-                    yield {ac[i]: bp[i] for i in range(len(bc))}
+    for i, formal_name in enumerate(remaining_formal):
+        mapping[formal_name] = remaining_actual[i]
 
-    for more_mapping in match(remaining_formal, remaining_actual):
-        # copy base mapping dict and add remaining vars
-        mapping_all = {k: v for k, v in mapping.items()}
-        mapping_all.update(more_mapping)
-        yield mapping_all
+    return mapping
 
 # Basic sanity check.
 if __name__ == '__main__':
@@ -153,5 +143,5 @@ if __name__ == '__main__':
     var_names = rename_vars(before)
     rename_vars(after, var_names)
 
-    for mapping in map_vars(before, after, line, rule):
-        print(mapping)
+    mapping = map_vars(before, after, line, rule)
+    print(mapping)