Simplify get_edits_from_traces

1 files changed, 7 insertions, 8 deletions

diff --git a/monkey/monkey.py b/monkey/monkey.py
index e185630..e90d70a 100755
--- a/monkey/monkey.py
+++ b/monkey/monkey.py
@@ -59,7 +59,7 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
                         if new_end > new_start:
                             new_rules.append((new_start, new_end))
                     new_step = ('remove_line', line_idx, (tuple(line), ()))
-                    new_cost = removes[line_normal] if line_normal in removes.keys() else 0.9
+                    new_cost = removes.get(line_normal, 1.0) * 0.99
 
                     yield (new_lines, new_rules, new_step, new_cost)
 
@@ -74,8 +74,9 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
                     new_rules.append((old_start + (0 if old_start < end else 1),
                                       old_end + (0 if old_end < end else 1)))
                 new_step = ('add_subgoal', end, ((), after_real))
+                new_cost = cost * 0.7
 
-                yield (new_lines, new_rules, new_step, cost)
+                yield (new_lines, new_rules, new_step, new_cost)
 
         # Add a new fact at the end.
         if len(rules) < 2:
@@ -83,8 +84,9 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
                 new_lines = lines + (after,)
                 new_rules = rules + (((len(lines), len(lines)+1)),)
                 new_step = ('add_rule', len(new_lines)-1, (tuple(), tuple(after)))
+                new_cost = cost * 0.7
 
-                yield (new_lines, new_rules, new_step, cost)
+                yield (new_lines, new_rules, new_step, new_cost)
 
     # Return a cleaned-up list of steps:
     #   - multiple line edits in a single line are merged into one
@@ -103,9 +105,6 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
             i += 1
         return new_steps
 
-    # 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
@@ -146,8 +145,8 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
         # Otherwise generate new solutions.
         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 * step_cost
-            if new_path_cost < 0.005:
+            new_path_cost = path_cost * new_cost
+            if new_path_cost < 0.4:
                 continue
             new_path = path + (new_step,)
             todo.push(((tuple(new_lines), tuple(new_rules)), new_path, new_path_cost), -new_path_cost)