Improve monkey.fix.postprocess

1 files changed, 17 insertions, 12 deletions

diff --git a/monkey/monkey.py b/monkey/monkey.py
index 080b317..7155c23 100755
--- a/monkey/monkey.py
+++ b/monkey/monkey.py
@@ -107,18 +107,23 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
     #   - check if any lines can be removed from the program
     def postprocess(steps):
         new_steps = []
-        i = 0
-        while i < len(steps):
-            # Check for successive edits on the same line.
-            if i < len(steps)-1 and \
-               steps[i][0] in ('add_rule', 'add_subgoal', 'change_line') and \
-               steps[i+1][0] == 'change_line' and \
-               steps[i][1] == steps[i+1][1] and steps[i][2][1] == steps[i+1][2][0]:
-                new_steps.append(('change_line', steps[i][1], (steps[i][2][0], steps[i+1][2][1])))
-                i += 1
-            else:
-                new_steps.append(steps[i])
-            i += 1
+        for step in steps:
+            if new_steps:
+                prev = new_steps[-1]
+
+                if prev[1] == step[1] and \
+                   prev[0] in ('add_rule', 'add_subgoal', 'change_line') and \
+                   step[0] == 'change_line' and \
+                   normalized(prev[2][1]) == normalized(step[2][0]):
+                    new_steps[-1] = (prev[0], prev[1], (prev[2][0], step[2][1]))
+                    continue
+
+                if prev[0] == 'add_subgoal' and step[0] == 'remove_line' and \
+                   prev[1]+1 == step[1]:
+                    new_steps[-1] = ('change_line', prev[1], (step[2][0], prev[2][1]))
+                    continue
+
+            new_steps.append(step)
         return new_steps
 
     # Main loop: best-first search through generated programs.