monkey.edits: only add solutions to predicates that are actually used when testing

1 files changed, 14 insertions, 9 deletions

diff --git a/monkey/edits.py b/monkey/edits.py
index df13940..f087074 100644
--- a/monkey/edits.py
+++ b/monkey/edits.py
@@ -193,9 +193,10 @@ def classify_edits(edits):
 # Extract edits and other data from existing traces for each problem.
 if __name__ == '__main__':
     import pickle
-    from db.models import CodeqUser, Problem, Solution
-    import server.problems
-    from termcolor import colored
+    from db.models import Problem, Solution
+    from db.util import make_identifier
+    from prolog.util import used_predicates
+    from server.problems import get_facts, load_problem, solutions_for_problems
 
     # Ignore traces from these users.
     ignored_users = [
@@ -215,13 +216,17 @@ if __name__ == '__main__':
             continue
 
         # Testing function.
-        group_module = server.problems.load_group(problem.language, problem.group, 'common')
-        problem_module = server.problems.load_problem(problem.language, problem.group, problem.identifier, 'common')
-        solved_problems = [p for p in CodeqUser.solved_problems(1, problem.language)
-                             if p[0] in group_module.allowed_groups and p != (problem.group, problem.identifier)]
-        aux_code = ('\n' + server.problems.solutions_for_problems('prolog', solved_problems) +
-                    '\n' + server.problems.get_facts('prolog', problem_module))
+        problem_module = load_problem(problem.language, problem.group, problem.identifier, 'common')
+        other_problems = [p for p in Problem.filter_language(problem.language)
+                            if p.identifier != problem.identifier]
+        facts = get_facts(problem.language, problem_module)
         def test(code):
+            # Find solutions to other problems that are used by this program.
+            used_predicate_identifiers = {make_identifier(name) for name in used_predicates(code)}
+            dependencies = sorted([p[2:] for p in other_problems
+                                         if p.identifier in used_predicate_identifiers])
+
+            aux_code = '\n' + solutions_for_problems('prolog', dependencies) + '\n' + facts
             correct, hints = problem_module.test(code, aux_code)
             return correct