Refactor monkey.monkey.fix

1 files changed, 48 insertions, 45 deletions

diff --git a/monkey/monkey.py b/monkey/monkey.py
index 5d26c6e..ef4ac2c 100755
--- a/monkey/monkey.py
+++ b/monkey/monkey.py
@@ -26,49 +26,10 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
     todo = PQueue()  # priority queue of candidate solutions
     done = set()     # set of already-analyzed solutions
 
-    # Add a new candidate solution ([lines]+[rules]) to the priority queue.
-    # This solution is generated by applying [step] with [cost] to [prev] task.
-    def add_task(lines, rules, prev=None, step=None, cost=None):
-        if prev is None:
-            path = ()
-            path_cost = 1.0
-        else:
-            path = tuple(list(prev[1]) + [step])
-            path_cost = prev[2] * cost
-        todo.push(((tuple(lines), tuple(rules)), path, path_cost), -path_cost)
-
-    lines, rules = decompose(code)
-    add_task(lines, rules)
-
     inserts, removes, changes = classify_edits(edits)
-    start_time = time.monotonic()
-    n_tested = 0
-    while True:
-        total_time = time.monotonic() - start_time
-        if total_time > timeout:
-            break
-
-        task = todo.pop()
-        if task == None:
-            break
 
-        (lines, rules), path, path_cost = task
-        code = compose(lines, rules)
-        if code in done:
-            continue
-        done.add(code)
-
-        if debug:
-            print('Cost {:.12f}'.format(path_cost))
-            for line, (before, after) in path:
-                print('line ' + str(line) + ':\t' + stringify(before) + ' → ' + stringify(after))
-
-        # If the code is correct, we are done.
-        if test(name, code + '\n' + aux_code):
-            return code, path, total_time, n_tested
-        n_tested += 1
-
-        # Otherwise generate new solutions.
+    # Generate states that can be reached from (lines, rules) with one edit.
+    def step(lines, rules):
         rule_no = 0
         for start, end in rules:
             rule = lines[start:end]
@@ -90,7 +51,7 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
                         new_lines = lines[:line_idx] + (after_real,) + lines[line_idx+1:]
                         new_step = ((rule_no, line_idx-start), (tuple(line), after_real))
 
-                        add_task(new_lines, rules, prev=task, step=new_step, cost=cost)
+                        yield (new_lines, rules, new_step, cost)
 
                 # if nothing could be done with this line, try removing it
                 # (maybe try removing in any case?)
@@ -105,7 +66,7 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
                     new_step = ((rule_no, line_idx-start), (tuple(line), ()))
                     new_cost = removes[line_normal] if line_normal in removes.keys() else 0.9
 
-                    add_task(new_lines, new_rules, prev=task, step=new_step, cost=new_cost)
+                    yield (new_lines, new_rules, new_step, new_cost)
 
             # try adding a line to this rule… would need to distinguish between
             # head/body lines in transforms
@@ -120,7 +81,7 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
                                       old_end + (0 if old_end < end else 1)))
                 new_step = ((rule_no, end-start), ((), after_real))
 
-                add_task(new_lines, new_rules, prev=task, step=new_step, cost=cost)
+                yield (new_lines, new_rules, new_step, cost)
             rule_no += 1
 
         # try adding a new fact
@@ -130,6 +91,48 @@ def fix(name, code, edits, aux_code='', timeout=30, debug=False):
                 new_rules = rules + (((len(lines), len(lines)+1)),)
                 new_step = ((len(new_rules)-1, 0), (tuple(), tuple(after)))
 
-                add_task(new_lines, new_rules, prev=task, step=new_step, cost=cost)
+                yield (new_lines, new_rules, new_step, cost)
+
+    # Add the initial program to the queue.
+    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
+
+        # Get next candidate.
+        task = todo.pop()
+        if task == None:
+            break
+        (lines, rules), path, path_cost = task
+
+        # If we have already seen this code, skip it.
+        code = compose(lines, rules)
+        if code in done:
+            continue
+        done.add(code)
+
+        if debug:
+            print('Cost {:.12f}'.format(path_cost))
+            for line, (before, after) in path:
+                print('line ' + str(line) + ':\t' + stringify(before) + ' → ' + stringify(after))
+
+        # If the code is correct, we are done.
+        if test(name, code + '\n' + aux_code):
+            return code, path, total_time, n_tested
+        n_tested += 1
+
+        # Otherwise generate new solutions.
+        path = list(path)
+        for new_lines, new_rules, new_step, new_cost in step(lines, rules):
+            new_path = tuple(path + [new_step])
+            new_path_cost = path_cost * new_cost * 0.99
+            if new_path_cost > 0.005:
+                todo.push(((tuple(new_lines), tuple(new_rules)), new_path, new_path_cost), -new_path_cost)
 
     return '', [], total_time, n_tested