Clean up monkey.fix.step

2 files changed, 49 insertions, 29 deletions

diff --git a/monkey/edits.py b/monkey/edits.py
index 480af5f..bf6560d 100644
--- a/monkey/edits.py
+++ b/monkey/edits.py
@@ -194,6 +194,11 @@ def get_edits_from_traces(traces):
                 start = normalize(path[i-1], var_names)
                 end = normalize(path[i], var_names)
 
+                # Disallow edits that insert a whole rule (a → … :- …).
+                # TODO improve edit_graph to handle this.
+                if 'FROM' in [t.type for t in end[:-1]]:
+                    continue
+
                 # This should always succeed but check anyway.
                 if start != end:
                     edit = (start, end)
diff --git a/monkey/monkey.py b/monkey/monkey.py
index e90d70a..0e622bf 100755
--- a/monkey/monkey.py
+++ b/monkey/monkey.py
@@ -31,19 +31,37 @@ 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)
+
+            # Apply some edits for each line in this rule.
             for line_idx in range(start, end):
                 if line_idx < start_line:
                     continue
                 line = lines[line_idx]
                 line_normal = tuple(rename_vars(line))
 
-                seen = False
-                # Apply each edit that matches this line.
+                # Apply edits whose left-hand side matches this line.
+                seen = False  # has there been such an edit?
                 for (before, after), cost in changes.items():
                     if line_normal == before:
                         seen = True
                         mapping = map_vars(before, after, line, rule_tokens)
-                        after_real = tuple([t if t.type != 'VARIABLE' else Token('VARIABLE', mapping[t.val]) for t in after])
+                        after_real = tuple([t if t.type != 'VARIABLE'
+                                              else Token('VARIABLE', mapping[t.val])
+                                              for t in after])
                         new_lines = lines[:line_idx] + (after_real,) + lines[line_idx+1:]
                         new_step = ('change_line', line_idx, (tuple(line), after_real))
 
@@ -63,30 +81,26 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
 
                     yield (new_lines, new_rules, new_step, new_cost)
 
-            # Add a line at the end of the current rule.
-            for after, cost in inserts.items():
-                mapping = map_vars([], after, [], rule_tokens)
-                after_real = [t if t.type != 'VARIABLE' else Token('VARIABLE', mapping[t.val]) for t in after]
-                after_real = tuple(after_real)
-                new_lines = lines[:end] + (after_real,) + lines[end:]
-                new_rules = []
-                for old_start, old_end in rules:
-                    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, new_cost)
-
-        # Add a new fact at the end.
-        if len(rules) < 2:
-            for after, cost in inserts.items():
-                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, new_cost)
+                # Add a line after this one.
+                for after, cost in inserts.items():
+                    # Don't try to insert a head into the body.
+                    if after[-1].type == 'FROM':
+                        continue
+                    mapping = map_vars([], after, [], rule_tokens)
+                    after_real = tuple([t if t.type != 'VARIABLE'
+                                          else Token('VARIABLE', mapping[t.val])
+                                          for t in after])
+
+                    idx = line_idx+1
+                    new_lines = lines[:idx] + (after_real,) + lines[idx:]
+                    new_rules = []
+                    for old_start, old_end in rules:
+                        new_rules.append((old_start + (0 if old_start < idx else 1),
+                                          old_end + (0 if old_end < idx else 1)))
+                    new_step = ('add_subgoal', idx, ((), after_real))
+                    new_cost = cost * 0.4
+
+                    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
@@ -96,7 +110,8 @@ def fix(name, code, edits, program_lines, aux_code='', timeout=30, debug=False):
         i = 0
         while i < len(steps):
             if i < len(steps)-1 and \
-               steps[i][0] == 'change_line' and steps[i+1][0] == 'change_line' and \
+               steps[i][0] in ('add_rule', 'add_subgoal', 'change_line') and \
+               steps[i+1][0] == 'change_line' and \
                steps[i][1] == steps[i+1][1] and steps[i][2][1] == steps[i+1][2][0]:
                 new_steps.append(('change_line', steps[i][1], (steps[i][2][0], steps[i+1][2][1])))
                 i += 1
@@ -146,7 +161,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.4:
+            if new_path_cost < 0.05:
                 continue
             new_path = path + (new_step,)
             todo.push(((tuple(new_lines), tuple(new_rules)), new_path, new_path_cost), -new_path_cost)