Add each edit at most once per solution

1 files changed, 4 insertions, 4 deletions

diff --git a/monkey/edits.py b/monkey/edits.py
index 1863d73..f27871b 100644
--- a/monkey/edits.py
+++ b/monkey/edits.py
@@ -27,11 +27,11 @@ def get_edits_from_trace(trace, test, id):
 
     # For each observed edit, store a list of features (with repeats) of ASTs
     # where they were observed.
-    edits = collections.defaultdict(list)
+    edits = collections.defaultdict(set)
     def add_edit(path, start, end, tree):
         if start == end:
             return
-        edits[(path, start, end)].append(id)
+        edits[(path, start, end)].add(id)
 
     # Parse trace actions and ensure there is a separate action for each
     # inserted/removed character.
@@ -127,7 +127,7 @@ def get_edits_from_trace(trace, test, id):
 def get_edits_from_solutions(solutions, test):
     # For each observed edit, store a list of features (with repeats) of ASTs
     # where they were observed.
-    submissions = collections.Counter()
+    submissions = collections.defaultdict(set)
     queries = collections.Counter()
     edits = collections.defaultdict(list)
 
@@ -143,7 +143,7 @@ def get_edits_from_solutions(solutions, test):
         # Update submission/query counters (use normalized variables).
         for code, correct in trace_submissions:
             code = stringify(rename_vars(tokenize(code)))
-            submissions[(code, correct)] += 1
+            submissions[(code, correct)].add(uid)
         for query in trace_queries:
             code = stringify(rename_vars(tokenize(query)))
             queries[code] += 1