Minor cleanups in monkey.monkey

1 files changed, 18 insertions, 16 deletions

diff --git a/monkey/monkey.py b/monkey/monkey.py
index cd44bac..cae42ae 100755
--- a/monkey/monkey.py
+++ b/monkey/monkey.py
@@ -14,14 +14,10 @@ from .util import PQueue, Token
 # Return (solution, edits, time spent, #programs checked). If no solution is
 # found within [timeout] seconds, solution='' and edits=[].
 def fix(name, code, edits, aux_code='', timeout=30, debug=False):
-    todo = PQueue()  # priority queue of candidate solutions
-    done = set()     # programs we have already visited
-
-    inserts, removes, changes = classify_edits(edits)
-
     # Generate states that can be reached from the given program with one edit.
     # Program code is given as a list of [lines], where each line is a list of
     # tokens. Rule ranges are given in [rules] (see prolog.util.decompose).
+    inserts, removes, changes = classify_edits(edits)
     def step(lines, rules, prev=None):
         # Apply edits in order from top to bottom; skip lines with index lower
         # than last step.
@@ -106,18 +102,22 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
             i += 1
         return new_steps
 
-    # Add the initial program to the queue.
+    # Prevents useless edits with cost=1.0.
+    step_cost = 0.99
+
+    # Main loop: best-first search through generated programs.
+    todo = PQueue()  # priority queue of candidate solutions
+    done = set()     # programs we have already visited
+
+    # Each program gets a task with the sequence of edits that generated it and
+    # the associated cost. First candidate with cost 1 is the initial program.
     lines, rules = decompose(code)
     todo.push(((tuple(lines), tuple(rules)), (), 1.0), -1.0)
 
-    start_time = time.monotonic()
     n_tested = 0
-    while True:
-        # Have we been at it for too long?
-        total_time = time.monotonic() - start_time
-        if total_time > timeout:
-            break
-
+    start_time = time.monotonic()
+    total_time = 0
+    while total_time < timeout:
         # Get next candidate.
         task = todo.pop()
         if task == None:
@@ -130,6 +130,7 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
             continue
         done.add(code)
 
+        # Print some info about the current task.
         if debug:
             print('Cost {:.12f}'.format(path_cost))
             for step_type, line, (before, after) in path:
@@ -142,15 +143,16 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
         n_tested += 1
 
         # Otherwise generate new solutions.
-        path = list(path)
         prev_step = path[-1] if path else None
         for new_lines, new_rules, new_step, new_cost in step(lines, rules, prev_step):
-            new_path_cost = path_cost * new_cost * 0.99
+            new_path_cost = path_cost * new_cost * step_cost
             if new_path_cost < 0.005:
                 continue
-            new_path = tuple(path + [new_step])
+            new_path = path + (new_step,)
             todo.push(((tuple(new_lines), tuple(new_rules)), new_path, new_path_cost), -new_path_cost)
 
+        total_time = time.monotonic() - start_time
+
     return '', [], total_time, n_tested
 
 # Return a list of character ranges modified by the sequence [edits].