Store test cases in PrologEngine

1 files changed, 48 insertions, 26 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index c859b67..85c5559 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -66,6 +66,14 @@ class Termv(object):
         for i, term in enumerate(terms):
             PL_put_term(self.ref+i, term.ref)
 
+class Problem(object):
+    def __init__(self, name, solution, facts, tests):
+        self.name = name
+        self.solution = solution
+        self.facts = facts
+        self.tests = {t: None for t in tests}
+        self.answers = {}
+
 class PrologEngine(object):
     def __init__(self):
         # Initialize the swipl library.
@@ -103,6 +111,7 @@ class PrologEngine(object):
         self.call('consult/1', [Term(Atom('prolog/lib.pl'))])
 
         # Load the time module (for call_with_time_limit) then disable autoload.
+        self.call('use_module/1', [Term('library(random)')])
         self.call('use_module/1', [Term('library(time)')])
         self.call('set_prolog_flag/2', [Term('autoload'), Term('false')])
 
@@ -113,32 +122,54 @@ class PrologEngine(object):
         # Discard messages from the swipl library.
         self.call('assertz/1', [Term('message_hook(_, _, _)')])
 
-        # Dictionary of correct answers; keys are (pid, query).
-        self.answers = {}
+        # Problem data loaded with load_problem.
+        self.problems = {}
+
+        # The set of already loaded facts.
+        self.facts = set()
 
-    # Loads the correct solution [code] to problem [pid].
-    def load_solution(self, pid, code, facts=None):
-        module = 'solution{}'.format(pid)
+    # Load the [solution] for problem [pid] called [name] and find answers to
+    # [tests]. Also load [facts] in the main module, and import modules for
+    # problems in [depends] into this problem's module.
+    def load_problem(self, pid, name, solution, depends, facts, tests):
+        self.problems[pid] = Problem(name, solution, facts, tests)
+
+        # Load the solution in 'solution<pid>' module.
+        m_problem = 'problem{}'.format(pid)
 
         fid = PL_open_foreign_frame()
         predicates = set()
-        for rule in prolog.util.split(code):
-            self.call('assertz/1', [Term('{}:({})'.format(module, rule))])
-            predicates.add('{}:{}'.format(module, self.predicate_indicator(rule)))
+        for rule in prolog.util.split(solution):
+            self.call('assertz/1', [Term('{}:({})'.format(m_problem, rule))])
+            predicates.add('{}:{}'.format(m_problem, self.predicate_indicator(rule)))
         if facts:
             for rule in prolog.util.split(facts):
-                self.call('assertz/1', [Term(rule)])
-                predicates.add(self.predicate_indicator(rule))
+                if rule not in self.facts:
+                    self.call('assertz/1', [Term(rule)])
+                    predicates.add(self.predicate_indicator(rule))
+                    self.facts.add(rule)
         self.call('compile_predicates/1', [Term([Term(p) for p in predicates])])
+
+        # Import solutions for dependency predicates.
+        for i in depends:
+            m_dependency = 'problem{}'.format(i)
+            self.call('add_import_module/3', [Term(m_problem), Term(m_dependency), Term('end')])
+
+        # Find the correct test answers.
+        for query in tests:
+            result = self.query(query, m_problem)
+            if result is None or len(result) < 1 or 'X' not in result[0]:
+                raise Exception('Error finding correct answer to query "{}"'.format(query))
+            self.problems[pid].tests[query] = result[0]['X']
         PL_discard_foreign_frame(fid)
 
     # Import the correct solution for problem [pid] into module for user [uid].
     def mark_solved(self, uid, pid):
         m_user = 'user{}'.format(uid)
-        m_solution = 'solution{}'.format(pid)
+        m_problem = 'problem{}'.format(pid)
 
         fid = PL_open_foreign_frame()
-        result = self.call('add_import_module/3', [Term(m_user), Term(m_solution), Term('end')])
+        result = self.call('add_import_module/3', [Term(m_user), Term(m_problem), Term('end')])
         PL_discard_foreign_frame(fid)
         return result
 
@@ -196,21 +227,12 @@ class PrologEngine(object):
         return result
 
     # Test whether [code] gives the same answer to [query] as the correct
-    # solution. The solution for problem [pid] should be loaded beforehand.
-    def test(self, uid, pid, code, queries):
-        # Module names for user code and the correct solution.
+    # solution for problem [pid]. The solution should be loaded beforehand.
+    def test(self, uid, pid, code):
+        # Module name for user code.
         m_user = 'user{}'.format(uid)
-        m_solution = 'solution{}'.format(pid)
 
         fid = PL_open_foreign_frame()
-        # Find the correct answers if not already known.
-        for query in queries:
-            if (pid, query) not in self.answers:
-                result = self.query(query, m_solution)
-                if result is None or len(result) < 1 or 'X' not in result[0]:
-                    raise Exception('Error finding correct answer to query "{}"'.format(query))
-                self.answers[(pid, query)] = result[0]['X']  # TODO maybe we could check all vars?
-
         correct = True
         predicates = set()
         try:
@@ -219,9 +241,9 @@ class PrologEngine(object):
                 self.call('assertz/1', [Term('{}:({})'.format(m_user, rule))])
                 predicates.add(self.predicate_indicator(rule))
 
-            for query in queries:
+            for query, answer in sorted(self.problems[pid].tests.items()):
                 result = self.query(query, m_user, n=1)
-                if len(result) != 1 or result[0]['X'] != self.answers[(pid, query)]:
+                if len(result) != 1 or result[0]['X'] != answer:
                     correct = False
                     break