Clean up prolog.engine

Use findnsols/4 to allow limiting the number of solutions found by
prolog.engine.query. Add a basic test case to prolog.engine.

1 files changed, 91 insertions, 81 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index f9cc6ce..d18b08b 100755
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -1,44 +1,55 @@
 #!/usr/bin/python3
 
-import re
-
 from prolog.core import *
 from prolog.util import *
 
 class Atom(object):
     __slots__ = 'ref'
 
-    def __init__(self, src=None, ref=None):
-        if ref != None:
+    def __init__(self, val=None, ref=None):
+        if ref is not None:
             self.ref = ref
-        else:
-            self.ref = PL_new_atom(bytes(src, encoding=encoding))
+            return
+        self.ref = PL_new_atom(bytes(val, encoding=encoding))
 
 class Term(object):
     __slots__ = 'ref'
 
-    def __init__(self, src=None, ref=None):
-        if ref != None:
+    # Initialize from term reference [ref] if given, otherwise construct a new
+    # term from [val] and possibly [args].
+    def __init__(self, val=None, args=None, ref=None):
+        if ref is not None:
             self.ref = ref
-        else:
-            self.ref = PL_new_term_ref()
-            if isinstance(src, str):
-                if not PL_chars_to_term(bytes(src, encoding=encoding), self.ref):
+            return
+        self.ref = PL_new_term_ref()
+        if isinstance(val, str):
+            if args is not None:
+                # Explicitly constructed compound term with name [val] and arguments [args].
+                name = PL_new_atom(bytes(val, encoding=encoding))
+                PL_cons_functor_v(self.ref, PL_new_functor(name, len(args)), Termv(args).ref)
+            else:
+                # Parse term from [val].
+                if not PL_chars_to_term(bytes(val, encoding=encoding), self.ref):
                     raise ValueError("invalid compound term")
-            elif isinstance(src, int):
-                PL_put_integer(self.ref, src)
-            elif isinstance(src, float):
-                PL_put_float(self.ref, src)
-            elif isinstance(src, list):
-                PL_put_nil(self.ref)
-                for t in src:
-                    PL_cons_list(self.ref, t.ref, self.ref)
-            elif isinstance(src, tuple) and len(src) == 2:
-                name = PL_new_atom(bytes(src[0], encoding=encoding))
-                args = src[1]
-                PL_cons_functor_v(self.ref, PL_new_functor(name, len(args)), args.ref)
-            elif isinstance(src, Atom):
-                PL_put_atom(self.ref, src.ref)
+        elif isinstance(val, int):
+            PL_put_integer(self.ref, val)
+        elif isinstance(val, float):
+            PL_put_float(self.ref, val)
+        elif isinstance(val, list):
+            PL_put_nil(self.ref)
+            for t in val:
+                PL_cons_list(self.ref, t.ref, self.ref)
+        elif isinstance(val, Atom):
+            PL_put_atom(self.ref, val.ref)
+
+    def __iter__(self):
+        if not PL_is_list(self.ref):
+            raise TypeError('term is not a list')
+        ref = self.ref
+        head, tail = Term(), Term()
+        while PL_get_list(ref, head.ref, tail.ref):
+            yield head
+            ref = tail.ref
 
     def __str__(self):
         ptr = c_char_p()
@@ -46,29 +57,12 @@ class Term(object):
             return str(ptr.value, encoding=encoding)
 
 class Termv(object):
-    __slots__ = 'ref', 'size'
+    __slots__ = 'ref'
 
     def __init__(self, terms):
-        self.size = len(terms)
-        self.ref = PL_new_term_refs(self.size)
-        for i in range(len(terms)):
-            PL_put_term(self.ref+i, terms[i].ref)
-
-    def __len__(self):
-        return self.size
-
-    def __getitem__(self, i):
-        if not isinstance(i, int) or len(self) == 0 or i < 0 or i >= len(self):
-            raise TypeError
-        return Term(ref=self.ref+i)
-
-    def __iter__(self):
-        for i in range(len(self)):
-            yield Term(ref=self.ref+i)
-        return
-
-    def __str__(self):
-        return '[' + ','.join([str(term) for term in self]) + ']'
+        self.ref = PL_new_term_refs(len(terms))
+        for i, term in enumerate(terms):
+            PL_put_term(self.ref+i, term.ref)
 
 class PrologEngine(object):
     def __init__(self):
@@ -91,6 +85,7 @@ class PrologEngine(object):
             'call_with_depth_limit': PL_predicate(b'call_with_depth_limit', 3, None),
             'call_with_time_limit': PL_predicate(b'call_with_time_limit', 2, None),
             'consult': PL_predicate(b'consult', 1, None),
+            'findnsols': PL_predicate(b'findnsols', 4, None),
             'message_to_string': PL_predicate(b'message_to_string', 2, None),
             'read_term_from_atom': PL_predicate(b'read_term_from_atom', 3, None),
             'safe_goal': PL_predicate(b'safe_goal', 1, None),
@@ -147,55 +142,70 @@ class PrologEngine(object):
             self.call('erase', [ref])
         del self.refs[index]
 
-    def query(self, q, module=None):
-        if module != None:
-            q = module + ':(' + q + ')'
+    # Get up to [n] solutions to query [q].
+    def query(self, q, module=None, n=1):
+        if module is not None:
+            q = '{}:({})'.format(module, q)
 
         fid = PL_open_foreign_frame()
-        goal = Term()
-        var_list = Term()
-        options = Term([Term(('variable_names', Termv([var_list])))])
 
-        # parse term
+        # Parse term and store variable names.
+        var_names = Term()
+        options = Term([Term('variable_names', [var_names])])
+        goal = Term()
         if not self.call('read_term_from_atom', [Term(Atom(q)), goal, options]):
             sys.stderr.write('Warning: Could not read term from {}\n'.format(q))
             return None
-        # check if goal is safe with currently loaded rules
+
+        # Check if goal is safe with currently loaded rules.
         if not self.call('safe_goal', [goal]):
-            sys.stderr.write('Warning: Unsafe goal: {}\n'.format(q))
+            sys.stderr.write('Warning: Unsafe goal: {}\n'.format(goal))
             return None
 
-        solutions = []
-        #depth_limit = Term(('call_with_depth_limit', Termv([goal, Term(50), Term('_')])))
+        solutions = Term()
+        goal_aux = Term('findnsols', [Term(n), goal, goal, solutions])
         qid = PL_open_query(None, self.err_flags, self.p['call_with_time_limit'],
-                            Termv([Term(0.01), goal]).ref)
-        while len(solutions) < 10:
-            if PL_next_solution(qid):
-                # get variable values in this solution as strings
+                            Termv([Term(0.01), goal_aux]).ref)
+
+        result = []
+        if PL_next_solution(qid):
+            for solution in solutions:
+                fid_solution = PL_open_foreign_frame()
+                PL_unify(goal.ref, solution.ref)
                 variables = {}
-                head, tail = Term(), Term()
-                while PL_get_list(var_list.ref, head.ref, tail.ref):
+                for var in var_names:
                     name, value = Term(), Term()
-                    PL_get_arg(1, head.ref, name.ref)
-                    PL_get_arg(2, head.ref, value.ref)
+                    PL_get_arg(1, var.ref, name.ref)
+                    PL_get_arg(2, var.ref, value.ref)
                     variables[str(name)] = str(value)
-                    var_list = tail
-                solutions += [variables]
-            else:
-                ex = PL_exception(qid)
+                result.append(variables)
+                PL_rewind_foreign_frame(fid_solution)
+        else:
+            ex = PL_exception(qid)
+            if ex != None:
+                ex = Term(ref=PL_exception(qid))
                 if ex != None:
-                    ex = Term(ref=PL_exception(qid))
-                    if ex != None:
-                        Msg = Term('Msg')
-                        if self.call('message_to_string', [ex, Msg]):
-                            sys.stderr.write('Error: ' + str(Msg) + '\n')
-                        else:
-                            sys.stderr.write('Unknown error\n')
-                # uncomment this if we get segfaults
-                #PL_clear_exception()
-                break
+                    Msg = Term()
+                    if self.call('message_to_string', [ex, Msg]):
+                        sys.stderr.write('Error: ' + str(Msg) + '\n')
+                    else:
+                        sys.stderr.write('Unknown error\n')
+                PL_clear_exception()
 
         PL_close_query(qid)
         PL_discard_foreign_frame(fid)
 
-        return solutions
+        return result
+
+# Basic sanity check.
+if __name__ == '__main__':
+    prolog = PrologEngine()
+    prolog.load('a(1). a(2). a(3).', 'test')
+    n = 1
+    while True:
+        solution = prolog.query('a(X), a(Y), X\=Y', 'test', n=n)
+        if solution is None or len(solution) != n:
+            print('{} solutions found!'.format(n-1))
+            break
+        print(solution[-1])
+        n += 1