MicroKanren (μKanren) in Haskell

MicroKanren.hs

import Control.Monad

type Var = Integer
type Subst = [(Var, Term)]
type State = (Subst, Integer)
type Program = State -> KList State

data Term = Atom String | Pair Term Term | Var Var deriving Show

-- Apply a substitution to the top level of a term
walk (Var v) s = case lookup v s of Nothing -> Var v
                                    Just us -> walk us s
walk u s = u

extS x v s = (x, v) : s

-- Try to unify two terms under a substitution;
-- return an extended subst if it succeeds
unify :: Term -> Term -> Subst -> Maybe Subst
unify u v s = un (walk u s) (walk v s)
  where un (Var x1) (Var x2) | x1 == x2 = return s
        un (Var x1) v = return $ extS x1 v s
        un u (Var x2) = return $ extS x2 u s
        un (Pair u1 u2) (Pair v1 v2) =
          do s' <- unify u1 v1 s
             unify u2 v2 s'
        un (Atom a1) (Atom a2) | a1 == a2 = return s
        un _ _  = mzero

-- MicroKanren program formers
zzz :: Program -> Program
zzz g = \sc -> delay (g sc)

equiv :: Term -> Term -> Program
equiv u v = \(s, c) -> case unify u v s of
  Nothing -> mzero
  Just s' -> return (s', c)

callFresh :: (Term -> Program) -> Program
callFresh f = \(s, c) -> f (Var c) (s, c+1)

disj :: Program -> Program -> Program
disj g1 g2 = \sc -> mplus (g1 sc) (g2 sc)

conj :: Program -> Program -> Program
conj g1 g2 = \sc -> g1 sc >>= g2

-- I had originally thought that since Haskell was lazy, we didn't
-- need to worry about any of the inverse-eta-delay stuff that the
-- Scheme version does, but that isn't right. We still need some way
-- to force switching when we recur.
-- It is not very burdensome for us, though; we don't actually need
-- to eta-expand, just need to add some sort of marker.
data KList a = Nil | Cons a (KList a) | Delay (KList a) deriving Show
delay = Delay

-- Hm. Is there any reason to preserve the delays?
instance Monad KList where
  return x = Cons x Nil
  Nil >>= f = mzero
  x `Cons` xs >>= f = f x `mplus` (xs >>= f)
  Delay xs >>= f = Delay (xs >>= f)
instance MonadPlus KList where
  mzero = Nil
  Nil `mplus` xs = xs
  (x `Cons` xs) `mplus` ys = x `Cons` (ys `mplus` xs) -- swapped per sect. 6
  Delay xs `mplus` ys = Delay (ys `mplus` xs)

klistToList Nil = []
klistToList (x `Cons` xs) = x : klistToList xs
klistToList (Delay xs) = klistToList xs

-------------- Test cases
empty_state = ([], 0)

five = callFresh (\x -> equiv x (Atom "5"))
fives_ x = disj (equiv x (Atom "5")) (zzz $ fives_ x)
fives = callFresh fives_

fivesRev_ x = disj (zzz $ fivesRev_ x) (equiv x (Atom "5"))
fivesRev = callFresh fivesRev_

a_and_b = conj
          (callFresh (\a -> equiv a (Atom "7")))
          (callFresh (\b -> disj (equiv b (Atom "5")) (equiv b (Atom "6"))))

runTest p = klistToList (p empty_state)

Note that KList is essentially the Stream monad (updated version here), and Program is the StateT transformation of it. This allows several things to be simplified, notably:

type Program = StateT State Stream
zzz = mapStateT Stream.suspended
callFresh = (fresh >>=)
disj = (<|>)
conj = (>>)
runTest p = Stream.observeAll (execStateT p ([], 0))

which uses the helper function:

fresh :: Program Term
fresh = do
  (s, c) <- get
  put (s, c+1)
  return (Var c)

A side benefit is that it also allows programs to be written a little more haskell-y:

a_and_b = do
  a <- fresh
  equiv a (Atom "7")
  b <- fresh
  equiv b (Atom "5") <|> equiv b (Atom "6")

Note that KList is essentially the Stream monad (updated version here), and Program is the StateT transformation of it. This allows several things to be simplified, notably:

type Program = StateT State Stream
zzz = mapStateT Stream.suspended
callFresh = (fresh >>=)
disj = (<|>)
conj = (>>)
runTest p = Stream.observeAll (execStateT p ([], 0))

which uses the helper function:

fresh :: Program Term
fresh = do
  (s, c) <- get
  put (s, c+1)
  return (Var c)

A side benefit is that it also allows programs to be written a little more haskell-y:

a_and_b = do
  a <- fresh
  equiv a (Atom "7")
  b <- fresh
  equiv b (Atom "5") <|> equiv b (Atom "6")

Only if you are fine with depth first search. It's very intentionally not the same as the stream monad, because the latter would do DFS and gets stuck in subgoals while a BFS implementation guarentees progress on all subgoals. They are very closely related though, and the better way to do it imho is to generalize to priority streams using a mergable heap type like pairing heaps.

@saolof Perhaps you're thinking of the Stream monad from the stream-fusion package (which unfortunately shares the same module namespace)? The stream-monad Stream deliberately avoids the pitfalls of DFS, see here for details