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

2 files changed, 20 insertions, 9 deletions

diff --git a/db/models.py b/db/models.py
index 2363802..eaf67d7 100644
--- a/db/models.py
+++ b/db/models.py
@@ -62,6 +62,12 @@ class Problem(collections.namedtuple('Problem', ['id', 'language', 'group', 'ide
         kwargs = {'p.'+k: v for k, v in kwargs.items()}
         return _general_filter(kwargs, Problem, Problem.__sql_prefix, order=Problem.__sql_order)
 
+    # get a list of problems with the given language identifier
+    @staticmethod
+    def filter_language(language):
+        kwargs = {'l.identifier': language}
+        return _general_filter(kwargs, Problem, Problem.__sql_prefix, order=Problem.__sql_order)
+
 # known as Attempt in the original code
 class Solution(collections.namedtuple('Solution', ['id', 'done', 'content', 'problem_id', 'codeq_user_id', 'trace'])):
     __sql_prefix = 'select id, done, content, problem_id, codeq_user_id, trace from solution'
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