Correctly find first token in current part in step

1 files changed, 7 insertions, 9 deletions

diff --git a/monkey/monkey.py b/monkey/monkey.py
index 8cf2163..7c8141c 100755
--- a/monkey/monkey.py
+++ b/monkey/monkey.py
@@ -26,7 +26,6 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
         # last step.
         start = path[-1][4] if path else 0
 
-        first = 0  # first token in the current part
         variables = []
         for i, token in enumerate(program):
             # Get variable names in the current rule.
@@ -39,13 +38,14 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
                     elif t.type == 'PERIOD':
                         break
 
+            # Remember the first token in the current part.
+            if i == 0 or program[i-1].stop:
+                first = i
+
             # Skip already modified parts of the program.
             if i < start:
                 continue
 
-            if i == 0 or program[i-1].stop:
-                first = i
-
             # Add a new fact at beginning or after another rule.
             if i == 0 or token.type == 'PERIOD':
                 new_start = i+1 if token.type == 'PERIOD' else 0
@@ -115,9 +115,6 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
                                 new_after[j] = t.clone(rule=t.rule-1)
                             new_program = program[:first] + tuple(new_after)
                             new_step = ('remove_rule', first, part, (), first)
-                            new_cost = removes.get(part_normal, 1.0) * 0.99
-                            yield (new_program, new_step, new_cost)
-
                         else:  # part is subgoal, remove part
                             new_after = list(program[last-1:])
                             for j, t in enumerate(new_after):
@@ -126,8 +123,9 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
                                 new_after[j] = t.clone(part=t.part-1)
                             new_program = program[:first-1] + tuple(new_after)
                             new_step = ('remove_part', first-1, (program[first-1],) + part, (), first-1)
-                            new_cost = removes.get(part_normal, 1.0) * 0.99
-                            yield (new_program, new_step, new_cost)
+
+                        new_cost = removes.get(part_normal, 1.0) * 0.99
+                        yield (new_program, new_step, new_cost)
 
 
                 # Insert a new part (goal) after this part.