mini-mini Kanren

Logic

{-# LANGUAGE RecordWildCards #-}
module Unify where
import Control.Monad.Logic
import qualified Data.Map as M

type Id = Integer
type Atom = String

data Term = Var Id
          | Atom Atom
          | Integer Integer
          | Pair Term Term
          deriving Show

type Sol = M.Map Id Term

-- | Substitute all bound variables in a term giving the canonical
-- term in an environment. All terms in an environment must be
-- canonized otherwise the world will burn.  And I'll be very sad.
canonize :: Sol -> Term -> Term
canonize sol t = case t of
  Atom a -> Atom a
  Integer i -> Integer i
  Pair l r -> canonize sol l `Pair` canonize sol r
  Var i -> maybe (Var i) id $ M.lookup i sol

-- | Extend an environment with a given term.
extend :: Id -> Term -> Sol -> Sol
extend i t sol = M.insert i (canonize sol t) sol

-- | Unification cannot need not backtrack so this will either
-- universally succeed or failure.
unify :: Term -> Term -> Sol -> Maybe Sol
unify l r sol= case (l, r) of
  (Atom a, Atom a') | a == a' -> Just sol
  (Integer i, Integer j) | i == j -> Just sol
  (Pair h t, Pair h' t') -> unify h h' sol >>= unify t t'
  (Var i, t) -> Just (extend i t sol)
  (t, Var i) -> Just (extend i t sol)
  _ -> Nothing

data State = State { sol  :: Sol
                   , vars :: [Term] }

type Predicate = State -> Logic State

-- | Equating two terms will attempt to unify them and backtrack if
-- this is impossible.
equate :: Term -> Term -> Predicate
equate l r s@State {..} =
  case unify (canonize sol l) (canonize sol r) sol of
   Just sol' -> return s{sol = sol'}
   Nothing   -> mzero

-- | Generate a fresh (not rigid) term to use for our program.
fresh :: (Term -> Predicate) -> Predicate
fresh withTerm (State sol (t : ts)) = withTerm t (State sol ts)

-- | Conjunction
conj :: Predicate -> Predicate -> Predicate
conj p1 p2 s = p1 s >>= p2

-- | Disjunction
disconj :: Predicate -> Predicate -> Predicate
disconj p1 p2 s = p1 s `interleave` p2 s

-- | Lots of disjunction, a logical switch statement.
conde :: [Predicate] -> Predicate
conde = foldr disconj (const mzero)

-- | Lots of conjuction, basically our toplevel program combinator.
program :: [Predicate] -> Predicate
program = foldr conj return

-- | Run a program and find all solutions for the parametrized term.
run :: (Term -> Predicate) -> [Term]
run mkProg = map answer (observeAll prog)
  where prog = mkProg (Var 0) $ State M.empty (map Var [1..])
        answer State{..} = canonize sol (Var 0)

test :: Term -> Predicate
test t = fresh $ \t' -> equate (Pair t' t) (Pair (Atom "foo") t')