Ignore edits after the first correct version

1 files changed, 7 insertions, 0 deletions

diff --git a/monkey/edits.py b/monkey/edits.py
index bf6560d..8747a6e 100644
--- a/monkey/edits.py
+++ b/monkey/edits.py
@@ -24,6 +24,7 @@ def edit_graph(actions, debug=False):
     leaves = {0: nodes[0]}  # Current leaf node for each line.
     rank = 1                # Rank (order / y-position) for the next node.
     code_next = ''          # Program code after applying the current action.
+    done = False            # Set to True on first correct version.
 
     # Ensure there is a separate action for each inserted/removed character.
     expand(actions)
@@ -33,11 +34,17 @@ def edit_graph(actions, debug=False):
 
         if action.type == 'test':
             submissions.add(code)
+            if action.total == action.passed:
+                done = True
 
         elif action.type == 'solve' or action.type == 'solve_all':
             queries.add(action.query)
 
         elif action.type == 'insert' or action.type == 'remove':
+            # Ignore edits after the first correct version.
+            if done:
+                continue
+
             # Number of the changed line.
             line = code[:action.offset].count('\n')
             # Current leaf node for this line.