Embedded Logic Programming Monad.

Monad.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Monad where

-- https://gup.ub.gu.se/file/207634

import Control.Monad.Fail
import Control.Monad.State
import Control.Monad.Trans.Maybe
import qualified Data.IntMap as IM

data Term
    = Term String [Term]
    | Var Int
    deriving Show

data Binding = Binding
    { next     :: Int
    , bindings :: IM.IntMap Term
    }

newtype Logic a
    = Logic { unLogic :: StateT Binding [] a }
    deriving (Applicative, Functor, Monad, MonadFail)

instance Semigroup (Logic a) where
    (Logic p) <> (Logic q) =
        Logic $ StateT $ \s -> runStateT p s ++ runStateT q s

instance MonadState Binding Logic where
    get = Logic get
    put = Logic . put
    state = Logic . state

fresh :: Logic Term
fresh = do
    v <- next <$> get
    modify (\s -> s { next = v + 1 })
    return $ Var v

lookupVar :: Int -> Logic (Maybe Term)
lookupVar v = gets (IM.lookup v . bindings)

bind :: Int -> Term -> Logic ()
bind v t = modify (\s -> s { bindings = IM.insert v t (bindings s) })

occurs :: Int -> Term -> Bool
occurs v (Term t ts) = any (occurs v) ts
occurs v (Var v')    = v == v'

apply :: Term -> MaybeT Logic Term
apply (Term t ts) = do
    ts <- traverse apply ts
    return $ Term t ts
apply (Var v) = do
    t <- lift $ lookupVar v
    case t of
        Just t -> do
            guard $ not (occurs v t)
            t <- apply t
            lift $ bind v t
            return t
        Nothing -> return $ Var v

unify :: Term -> Term -> MaybeT Logic Term
unify t t' = do
    t  <- apply t
    t' <- apply t'
    case (t, t') of
        (Var v, Var v')
            | v == v' -> return $ Var v
        (Var v, t) -> do
            lift $ bind v t
            return $ Var v
        (t, Var v) -> do
            lift $ bind v t
            return $ Var v
        (t, t') -> match t t'

zipExact :: [a] -> [b] -> MaybeT Logic [(a, b)]
zipExact [] []         = return []
zipExact (x:xs) (y:ys) = ((x, y) :) <$> zipExact xs ys
zipExact _ _           = mzero

match :: Term -> Term -> MaybeT Logic Term
match (Term t ts) (Term t' ts') = do
    guard (t == t')
    ts <- zipExact ts ts'
    ts <- traverse (uncurry unify) ts
    return $ Term t ts
match t t' = unify t t'

(===) :: Term -> Term -> Logic ()
t === t' = do
    Just _ <- runMaybeT $ unify t t'
    return ()

true :: Logic ()
true = return ()

false :: Logic a
false = Logic $ StateT $ \s -> []

edge :: Term -> Term -> Logic ()
edge x y =
    foldr (<>) false [x === Term a [] >> y === Term b [] | (a, b) <- graph]

graph :: [(String, String)]
graph =
    [ ("a", "b")
    , ("a", "d")
    , ("b", "c")
    , ("b", "d")
    , ("c", "d")
    , ("c", "e")
    , ("d", "e")
    ]

atom :: String -> Term
atom t = Term t []

cons :: Term -> Term -> Term
cons x xs = Term "cons" [x, xs]

nil :: Term
nil = Term "nil" []

path :: Term -> Term -> Term -> Logic ()
path x y query = do
        x === y
        query === cons x nil
    <> do
        query' <- fresh
        query === cons x query'
        z <- fresh
        edge x z
        path z y query'

solve :: (Term -> Logic ()) -> [Term]
solve query =
    (`evalStateT` Binding 0 IM.empty) . unLogic $ do
        var <- fresh
        query var
        Just t <- runMaybeT $ apply var
        return t

-- putStrLn . unlines . map show $ solve (path (atom "a") (atom "e"))