Remember correct answers to test queries

1 files changed, 13 insertions, 10 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index ec27e27..92caa0a 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -93,11 +93,14 @@ class PrologEngine(object):
         }
         self.err_flags = PL_Q_NODEBUG|PL_Q_CATCH_EXCEPTION
 
+        # Dictionary of correct answers; keys are (pid, query).
+        self.answers = {}
+
         # Increase memory limits.
         self.call('set_prolog_stack/2', [Term('global'), Term('limit(2*10**9)')])
         self.call('set_prolog_stack/2', [Term('local'), Term('limit(2*10**9)')])
 
-        # Discard messages from swipl library.
+        # Discard messages from the swipl library.
         self.call('assertz/1', [Term('message_hook(_, _, _)')])
 
     def exception(self, qid):
@@ -189,19 +192,19 @@ class PrologEngine(object):
         return self.call('add_import_module/3', [Term(m_user), Term(m_solution), Term('end')])
 
     # Test whether [code] gives the same answer to [query] as the correct
-    # solution.
+    # 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.
         m_user = 'user{}'.format(uid)
         m_solution = 'solution{}'.format(pid)
 
-        # Find the correct answers. Code for [pid] should be loaded beforehand.
-        # TODO Memoize correct answers.
-        answers = []
+        # Find the correct answers if not already known.
         for query in queries:
-            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))
-            answers.append(result[0]['X'])  # TODO maybe we could check all vars?
+            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
         refs = []
@@ -214,7 +217,7 @@ class PrologEngine(object):
 
             for i, query in enumerate(queries):
                 result = self.query(query, m_user)
-                correct &= (len(result) == 1 and result[0]['X'] == answers[i])
+                correct &= (len(result) == 1 and result[0]['X'] == self.answers[(pid, query)])
         except Exception as ex:
             correct = False
         finally: