Merge pull request #59 from brandonwillard/efficient-mk-stream-ops

Efficient miniKanren stream functions

Author: brandonwillard

setup.cfg

@@ -13,4 +13,6 @@ filterwarnings =
 
 [coverage:report]
 exclude_lines =
-    pragma: no cover
+    pragma: no cover
+
+    raise NotImplementedError

symbolic_pymc/constraints.py

@@ -0,0 +1,184 @@
+import weakref
+
+from abc import ABC, abstractmethod
+from types import MappingProxyType
+from collections import OrderedDict
+
+from unification import unify, reify, Var
+from unification.core import _unify, _reify
+
+
+class KanrenConstraintStore(ABC):
+    """A class that enforces constraints between logic variables in a miniKanren state."""
+
+    __slots__ = ("mappings",)
+
+    def __init__(self, mappings=None):
+        # Mappings between logic variables and their constraint values
+        # (e.g. the values against which they cannot be unified).
+        self.mappings = mappings if mappings is not None else dict()
+        # TODO: We can't use this until `Var` is a factory returning unique
+        # objects/references for a given `Var.token` value.
+        # self.mappings = weakref.WeakKeyDictionary(mappings)
+
+    @abstractmethod
+    def pre_check(self, state, key=None, value=None):
+        """Check a key-value pair before they're added to a KanrenState."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def post_check(self, new_state, key=None, value=None, old_state=None):
+        """Check a key-value pair after they're added to a KanrenState."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def update(self, *args, **kwargs):
+        """Add a new constraint."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def constraints_str(self, var):
+        """Print the constraints on a logic variable."""
+        raise NotImplementedError()
+
+
+class KanrenState(dict):
+    """A miniKanren state that holds unifications of logic variables and upholds constraints on logic variables."""
+
+    __slots__ = ("constraints", "__weakref__")
+
+    def __init__(self, *s, constraints=None):
+        super().__init__(*s)
+        self.constraints = OrderedDict(constraints or [])
+
+    def pre_checks(self, key, value):
+        return all(cstore.pre_check(self, key, value) for cstore in self.constraints.values())
+
+    def post_checks(self, new_state, key, value):
+        return all(
+            cstore.post_check(new_state, key, value, old_state=self)
+            for cstore in self.constraints.values()
+        )
+
+    def add_constraint(self, constraint):
+        assert isinstance(constraint, KanrenConstraintStore)
+        self.constraints[type(constraint)] = constraint
+
+    def __eq__(self, other):
+        if isinstance(other, KanrenState):
+            return super().__eq__(other)
+
+        # When comparing with a plain dict, disregard the constraints.
+        if isinstance(other, dict):
+            return super().__eq__(other)
+        return False
+
+    def __repr__(self):
+        return f"KanrenState({super().__repr__()}, {self.constraints})"
+
+
+class Disequality(KanrenConstraintStore):
+    """A disequality constraint (i.e. two things do not unify)."""
+
+    def post_check(self, new_state, key=None, value=None, old_state=None):
+        # This implementation follows-up every addition to a `KanrenState` with
+        # a consistency check against all the disequality constraints.  It's
+        # not particularly scalable, but it works for now.
+        return not any(
+            any(new_state == unify(lvar, val, new_state) for val in vals)
+            for lvar, vals in self.mappings.items()
+        )
+
+    def pre_check(self, state, key=None, value=None):
+        return True
+
+    def update(self, key, value):
+        # In this case, logic variables are mapped to a set of values against
+        # which they cannot unify.
+        if key not in self.mappings:
+            self.mappings[key] = {value}
+        else:
+            self.mappings[key].add(value)
+
+    def constraints_str(self, var):
+        if var in self.mappings:
+            return f"=/= {self.mappings[var]}"
+        else:
+            return ""
+
+    def __repr__(self):
+        return ",".join([f"{k} =/= {v}" for k, v in self.mappings.items()])
+
+
+def unify_KanrenState(u, v, S):
+    if S.pre_checks(u, v):
+        s = unify(u, v, MappingProxyType(S))
+        if s is not False and S.post_checks(s, u, v):
+            return KanrenState(s, constraints=S.constraints)
+
+    return False
+
+
+unify.add((object, object, KanrenState), unify_KanrenState)
+unify.add(
+    (object, object, MappingProxyType),
+    lambda u, v, d: unify.dispatch(type(u), type(v), dict)(u, v, d),
+)
+_unify.add(
+    (object, object, MappingProxyType),
+    lambda u, v, d: _unify.dispatch(type(u), type(v), dict)(u, v, d),
+)
+
+
+class ConstrainedVar(Var):
+    """A logic variable that tracks its own constraints.
+
+    Currently, this is only for display/reification purposes.
+
+    """
+
+    __slots__ = ("_id", "token", "S", "var")
+
+    def __new__(cls, var, S):
+        obj = super().__new__(cls, var.token)
+        obj.S = weakref.ref(S)
+        obj.var = weakref.ref(var)
+        return obj
+
+    def __repr__(self):
+        var = self.var()
+        S = self.S()
+        if var is not None and S is not None:
+            u_constraints = ",".join([c.constraints_str(var) for c in S.constraints.values()])
+            return f"{var}: {{{u_constraints}}}"
+
+
+def reify_KanrenState(u, S):
+    u_res = reify(u, MappingProxyType(S))
+    if isinstance(u_res, Var):
+        return ConstrainedVar(u_res, S)
+    else:
+        return u_res
+
+
+_reify.add((tuple(p[0] for p in _reify.ordering if p[1] == dict), KanrenState), reify_KanrenState)
+_reify.add((object, MappingProxyType), lambda u, s: _reify.dispatch(type(u), dict)(u, s))
+
+
+def neq(u, v):
+    """Construct a disequality goal."""
+
+    def neq_goal(S):
+        if not isinstance(S, KanrenState):
+            S = KanrenState(S)
+
+        diseq_constraint = S.constraints.setdefault(Disequality, Disequality())
+
+        diseq_constraint.update(u, v)
+
+        if diseq_constraint.post_check(S):
+            return iter([S])
+        else:
+            return iter([])
+
+    return neq_goal

symbolic_pymc/relations/__init__.py

@@ -1,7 +1,12 @@
-from unification import isvar
+from itertools import tee, chain
+from functools import reduce
+
+from toolz import interleave
+
+from unification import isvar, reify
 
 from kanren import eq
-from kanren.core import lallgreedy
+from kanren.core import goaleval
 from kanren.facts import Relation
 from kanren.goals import goalify
 from kanren.term import term, operator, arguments
@@ -20,13 +25,85 @@
 concat = goalify(lambda *args: "".join(args))
 
 
+def ldisj_seq(goals):
+    """Produce a goal that returns the appended state stream from all successful goal arguments.
+
+    In other words, it behaves like logical disjunction/OR for goals.
+    """
+
+    def ldisj_seq_goal(S):
+        nonlocal goals
+
+        goals, _goals = tee(goals)
+
+        yield from interleave(goaleval(g)(S) for g in _goals)
+
+    return ldisj_seq_goal
+
+
+def lconj_seq(goals):
+    """Produce a goal that returns the appended state stream in which all goals are necessarily successful.
+
+    In other words, it behaves like logical conjunction/AND for goals.
+    """
+
+    def lconj_seq_goal(S):
+        nonlocal goals
+
+        goals, _goals = tee(goals)
+
+        g0 = next(iter(_goals), None)
+
+        if g0 is None:
+            return
+
+        z0 = goaleval(g0)(S)
+
+        yield from reduce(lambda z, g: chain.from_iterable(map(goaleval(g), z)), _goals, z0)
+
+    return lconj_seq_goal
+
+
+def ldisj(*goals):
+    return ldisj_seq(goals)
+
+
+def lconj(*goals):
+    return lconj_seq(goals)
+
+
+def conde(*goals):
+    return ldisj_seq(lconj_seq(g) for g in goals)
+
+
+lall = lconj
+lany = ldisj
+
+
 def buildo(op, args, obj):
-    if not isvar(obj):
-        if not isvar(args):
-            args = etuplize(args, shallow=True)
-        oop, oargs = operator(obj), arguments(obj)
-        return lallgreedy(eq(op, oop), eq(args, oargs))
-    elif isvar(args) or isvar(op):
-        return conso(op, args, obj)
-    else:
-        return eq(obj, term(op, args))
+    """Construct a goal that relates an object and its rand + rators decomposition.
+
+    This version uses etuples.
+
+    """
+
+    def buildo_goal(S):
+        nonlocal op, args, obj
+
+        op, args, obj = reify((op, args, obj), S)
+
+        if not isvar(obj):
+
+            if not isvar(args):
+                args = etuplize(args, shallow=True)
+
+            oop, oargs = operator(obj), arguments(obj)
+
+            yield from lall(eq(op, oop), eq(args, oargs))(S)
+
+        elif isvar(args) or isvar(op):
+            yield from conso(op, args, obj)(S)
+        else:
+            yield from eq(obj, term(op, args))(S)
+
+    return buildo_goal

symbolic_pymc/relations/graph.py

@@ -7,9 +7,10 @@
 
 from kanren import eq
 from cons.core import ConsPair, ConsNull
-from kanren.core import conde, lall
 from kanren.goals import conso, fail
 
+from . import conde, lall
+
 from ..etuple import etuplize, ExpressionTuple
 
 
@@ -217,7 +218,7 @@ def graph_applyo(
         def preprocess_graph(x):
             return x
 
-    def _gapplyo(s):
+    def graph_applyo_goal(s):
 
         nonlocal in_graph, out_graph
 
@@ -256,4 +257,4 @@ def _gapplyo(s):
         g = goaleval(g)
         yield from g(s)
 
-    return _gapplyo
+    return graph_applyo_goal

symbolic_pymc/tensorflow/unify.py

@@ -8,19 +8,20 @@
 from ..meta import metatize
 from ..unify import unify_MetaSymbol
 from ..etuple import etuplize
+from ..constraints import KanrenState
 
 tf_class_abstractions = tuple(c.base for c in TFlowMetaSymbol.base_subclasses())
 
 _unify.add(
-    (TFlowMetaSymbol, tf_class_abstractions, dict),
+    (TFlowMetaSymbol, tf_class_abstractions, (KanrenState, dict)),
     lambda u, v, s: unify_MetaSymbol(u, metatize(v), s),
 )
 _unify.add(
-    (tf_class_abstractions, TFlowMetaSymbol, dict),
+    (tf_class_abstractions, TFlowMetaSymbol, (KanrenState, dict)),
     lambda u, v, s: unify_MetaSymbol(metatize(u), v, s),
 )
 _unify.add(
-    (tf_class_abstractions, tf_class_abstractions, dict),
+    (tf_class_abstractions, tf_class_abstractions, (KanrenState, dict)),
     lambda u, v, s: unify_MetaSymbol(metatize(u), metatize(v), s),
 )
 
@@ -30,7 +31,7 @@ def _reify_TFlowClasses(o, s):
     return reify(meta_obj, s)
 
 
-_reify.add((tf_class_abstractions, dict), _reify_TFlowClasses)
+_reify.add((tf_class_abstractions, (KanrenState, dict)), _reify_TFlowClasses)
 
 operator.add((tf.Tensor,), lambda x: operator(metatize(x)))