Clean up monkey.prolog.util

1 files changed, 11 insertions, 13 deletions

diff --git a/monkey/edits.py b/monkey/edits.py
index 58843f6..3e0ae08 100644
--- a/monkey/edits.py
+++ b/monkey/edits.py
@@ -170,31 +170,29 @@ def get_edits_from_traces(traces):
         # Update the submissions/queries counters; rename variables first to
         # remove trivial differences.
         for submission in trace_submissions:
-            tokens = tokenize(submission)
-            rename_vars(tokens)
-            code = stringify(tokens)
+            code = stringify(rename_vars(tokenize(submission)))
             submissions[code] += 1
 
         for query in trace_queries:
-            tokens = tokenize(query)
-            rename_vars(tokens)
-            code = stringify(tokens)
+            code = stringify(rename_vars(tokenize(query)))
             queries[code] += 1
 
         # Get edits.
+        done = set()
         for path in get_paths(nodes[0]):
             for i in range(len(path)):
-                start = list(remove_punct(path[i]))
-                var_names = rename_vars(start)
-                start_t = tuple(start)
+                var_names = {}
+                start = remove_punct(path[i])
+                start_t = tuple(rename_vars(start, var_names))
 
                 for j in range(len(path[i+1:])):
+                    var_names_copy = {k: v for k, v in var_names.items()}
                     end = list(remove_punct(path[i+1+j]))
-                    rename_vars(end, var_names)
-                    end_t = tuple(end)
+                    end_t = tuple(rename_vars(end, var_names_copy))
 
-                    if start_t != end_t:
-                        edit = (start_t, end_t)
+                    edit = (start_t, end_t)
+                    if start_t != end_t and edit not in done:
+                        done.add(edit)
                         edits[edit] += 1
                         lines[start_t] += 1