Improve adding a rule in monkey.fix.step

1 files changed, 16 insertions, 16 deletions

diff --git a/monkey/monkey.py b/monkey/monkey.py
index c037193..3491186 100755
--- a/monkey/monkey.py
+++ b/monkey/monkey.py
@@ -31,20 +31,20 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
             rule_lines = lines[start:end]
             rule_tokens = [t for line in rule_lines for t in line]
 
-            # Add a new fact before this rule.
-            if len(rules) < 3:
-                for after, cost in inserts.items():
-                    new_lines = lines[:start] + (after,) + lines[start:]
-                    new_rules = []
-                    for old_start, old_end in rules:
-                        new_rules.append((old_start + (0 if old_start < start else 1),
-                                          old_end + (0 if old_end < start else 1)))
-                        if old_end == start:
-                            new_rules.append((start, start+1))
-                    new_step = ('add_rule', len(new_lines)-1, (tuple(), after))
-                    new_cost = cost * 0.3
-
-                    yield (new_lines, new_rules, new_step, new_cost)
+            # Add a new rule (fact) before this rule.
+            for after, cost in inserts.items():
+                new_lines = lines[:start] + (after,) + lines[start:]
+                new_rules = []
+                for old_start, old_end in rules:
+                    if old_start == start:
+                        new_rules.append((start, start+1))
+                    new_rules.append((old_start + (0 if old_start < start else 1),
+                                      old_end + (0 if old_end < start else 1)))
+                new_step = ('add_rule', start, (tuple(), after))
+                # Decrease probability as we add more rules.
+                new_cost = cost * math.pow(0.3, len(rules))
+
+                yield (new_lines, new_rules, new_step, new_cost)
 
             # Apply some edits for each line in this rule.
             for line_idx in range(start, end):
@@ -155,7 +155,7 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
         if debug:
             print('Cost {:.12f}'.format(path_cost))
             for step_type, line, (before, after) in path:
-                print('line {}: {} {} → {}'.format(line, step_type,  stringify(before), stringify(after)))
+                print('line {}: {} {} → {}'.format(line, step_type, stringify(before), stringify(after)))
 
         # If the code is correct, we are done.
         if test(name, code + '\n' + aux_code):
@@ -167,7 +167,7 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
         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
-            if new_path_cost < 0.05:
+            if new_path_cost < 0.01:
                 continue
             new_path = path + (new_step,)
             todo.push(((tuple(new_lines), tuple(new_rules)), new_path, new_path_cost), -new_path_cost)