Minimal miniKanren implementation in Haskell

Main.hs

module Main where

import MiniKanren

-- r = n + 1
peanoSuccO :: Term -> Term -> Goal 
peanoSuccO n r = r === Val (Cons n (Val Nil))

-- r = n - 1
peanoPredO :: Term -> Term -> Goal 
peanoPredO = flip peanoSuccO

-- r = m + n
peanoPlusO :: Term -> Term -> Term -> Goal
peanoPlusO m n r = conde' [
  [n === Val Nil, r === m],
  [fresh (\m' -> [
    fresh (\n' -> [
      peanoPredO n n',
      peanoSuccO m m',
      peanoPlusO m' n' r])])]]

-- builds a peano numeral term out of a positive integer
natToPeano :: Int -> Term
natToPeano 0 = Val Nil
natToPeano n = Val (Cons (natToPeano (n - 1)) (Val Nil))

-- find x such that x + 1 = 3
main :: IO ()
main = do
  print $ run' 1 $ \x -> peanoPlusO x (natToPeano 1) (natToPeano 3)

MiniKanren.hs

module MiniKanren where

import qualified Data.Map as Map

-- # Triangular Substitutions # --

type Variable = Int -- we'll just treat variables as unique id numbers
data Value = Number Int | Boolean Bool | Cons Term Term | Nil | String String
  deriving (Eq, Show)
data Term = Var Variable | Val Value
  deriving (Eq, Show)

type Subst = Map.Map Variable Term

emptySubst :: Subst
emptySubst = Map.empty

-- extSubst a b s |-> s, extended with a->b
extSubst :: Variable -> Term -> Subst -> Subst
extSubst = Map.insert

-- getSubst a s |-> Just b if a->b in s, otherwise Nothing
getSubst :: Variable -> Subst -> Maybe Term
getSubst = Map.lookup

lenSubst :: Subst -> Int
lenSubst = Map.size

-- # The `walk` Function # --

walk :: Term -> Subst -> Term

-- for a variable, look up the next step in the substitution
walk t@(Var v) s = case getSubst v s of
  Nothing   -> t         -- if we get nothing, then the variable is free
  (Just t') -> walk t' s -- otherwise, walk again with the result

-- for a value, just return the value
walk t _ = t

-- # Avoiding Circular Substitutions # --

checkCycle :: Variable -> Term -> Subst -> Bool
checkCycle v t s = case t' of
  (Var v') -> v == v' -- walk diverges if t's chain ends at x
  (Val a)  -> case a of -- reify diverges if a pair/list contains itself
                (Cons x y) -> checkCycle v x s || checkCycle v y s
                _          -> False
  where
    t' = walk t s

extSubstChecked :: Variable -> Term -> Subst -> Maybe Subst
extSubstChecked v t s = if checkCycle v t s
                        then Nothing -- fail if there's a cycle
                        else Just (extSubst v t s)

-- # Implementing Unification # --

unify :: Term -> Term -> Subst -> Maybe Subst
unify a b s = unify' (walk a s) (walk b s) s -- find ends of chains
  where
    unify' :: Term -> Term -> Subst -> Maybe Subst
    unify' a b s
      | a == b = Just s -- a already equals b, so nothing to unify
      | (Var v) <- a = extSubstChecked v b s -- a is var; add a->b
      | (Var v) <- b = extSubstChecked v a s -- b is var; add b->a
      | (Val (Cons ha ta)) <- a, -- a and b are pairs/lists; go elementwise
        (Val (Cons hb tb)) <- b = case unify ha hb s of
                                    Nothing   -> Nothing
                                    (Just s') -> unify ta tb s'
      | otherwise = Nothing -- if none of the above, unification failed

-- # Implementing Reification # --

reify :: Term -> Subst -> Term
reify t s = deepWalk t' (reifySubst t' emptySubst)
  where
    t' = deepWalk t s

reifySubst :: Term -> Subst -> Subst
reifySubst (Var v) s = extSubst v (Var (lenSubst s)) s -- new free var
reifySubst (Val (Cons x y)) s = reifySubst x (reifySubst y s)
reifySubst _ s = s

-- # The `deepWalk` Function --

deepWalk :: Term -> Subst -> Term
deepWalk t s = case t' of
  (Val (Cons x y)) -> Val (Cons (deepWalk x s) (deepWalk y s)) -- recurse
  _                -> t' -- walk hit a dead end
  where
    t' = walk t s -- walk does most of the work already

-- # Behind the Streams # --

type Cont a = () -> Stream a -- Continuation of a stream

data Stream a = Empty             -- Empty stream
              | Unit a            -- Singleton stream
              | Choice a (Cont a) -- A value + the rest of the stream
              | Inc (Cont a)      -- The rest of some stream

-- # Implementing `==` With Streams # --

-- goal takes a substitution and produces all satisfying extensions thereof
type Goal = Subst -> Stream Subst

-- goal asserting that two things are the same
(===) :: Term -> Term -> Goal
(===) a b = \s -> case unify a b s of
  (Just s') -> Unit s' -- unification succeeded
  Nothing   -> Empty -- unification failed

-- # Behind the Interleaving Streams # --

-- `conde` just folds over a list of goals using (+)
conde :: [Goal] -> Goal
conde []       = \s -> Empty
conde (g : gs) = \s -> mplus (g s) (\_ -> (conde gs) s)

-- The (+) operator, which interleaves two streams
mplus :: Stream a -> Cont a -> Stream a
mplus s f = case s of
  Empty         -> f ()
  (Unit a)      -> Choice a f
  (Inc f')      -> Inc (\_ -> mplus (f ()) f') -- interleaved!
  (Choice a f') -> Choice a (\_ -> mplus (f ()) f') -- interleaved!

-- # Sequencing Streams # --

seqGoals :: [Goal] -> Goal -- folds over a list of goals with bind
seqGoals []       = \s -> Unit s
seqGoals (g : gs) = \s -> bind (g s) (seqGoals gs)

bind :: Stream a -> (a -> Stream a) -> Stream a -- the >>= operator
bind s g = case s of
  Empty        -> Empty
  (Unit a)     -> g a
  (Inc f)      -> Inc (\_ -> bind (f ()) g)
  (Choice a f) -> mplus (g a) (\_ -> bind (f ()) g)

-- # Taking Values From Streams # --

streamTake :: Int -> Stream a -> [a]
streamTake 0 _            = [] -- take 0, get empty list
streamTake _ Empty        = [] -- nothing to take
streamTake _ (Unit a)     = [a] -- can't take more than one
streamTake n (Inc f)      = streamTake n (f ()) -- immediately continue
streamTake n (Choice a f) = a : streamTake (n - 1) (f ()) -- took one more

streamConsume :: Stream a -> [a]
streamConsume Empty        = []
streamConsume (Unit a)     = [a]
streamConsume (Inc f)      = streamConsume (f ())
streamConsume (Choice a f) = a : streamConsume (f ())

run :: Int -> Term -> Goal -> [Term]
run k t g = (\s -> reify t s) <$> streamTake k (g emptySubst)

runAll :: Term -> Goal -> [Term]
runAll t g = (\s -> reify t s) <$> streamConsume (g emptySubst)

-- # Other Stuff # --

-- conde with conjunction
conde' :: [[Goal]] -> Goal
conde' ggs = conde (seqGoals <$> ggs)

-- hacky way to allocate unique variables:
-- we store the current highest variable id in a special key in the subst
varCountSentinel :: Variable
varCountSentinel = -1

allocVar :: Subst -> (Variable, Subst)
allocVar s = case getSubst varCountSentinel s of
  (Just (Val (Number n))) -> (n, extSubst varCountSentinel (Val (Number (n + 1))) s)
  Nothing -> (1, extSubst varCountSentinel (Val (Number 2)) s)
  _ -> error "Var count is bad!"

fresh :: (Term -> [Goal]) -> Goal
fresh f s = seqGoals (f (Var v)) s'
  where
    (v, s') = allocVar s

-- versions of run and run* that give you a term to work with
run' :: Int -> (Term -> Goal) -> [Term]
run' k f = run k t (f t)
  where
    t = Var 0

runAll' :: (Term -> Goal) -> [Term]
runAll' f = runAll t (f t)
  where
    t = Var 0