bergamot/src/Language/Bergamot/Unifier.purs

module Language.Bergamot.Unifier where

import Prelude

import Language.Bergamot.Rules (Metavariable, ProofTree(..), Rule(..), instantiate)
import Language.Bergamot.Syntax (Expr(..), IntVar)
import Language.Bergamot.Latex

import Control.Plus (class Plus, empty)
import Control.Apply (lift2)
import Control.Alt (class Alt, alt, (<|>))
import Control.Alternative (class Alternative)
import Control.Lazy (class Lazy, defer)
import Control.MonadPlus (class MonadPlus)
import Control.Monad.Reader.Class (class MonadReader, class MonadAsk, local, ask, asks)
import Control.Monad.Logic.Class (class MonadLogic, msplit, interleave)
import Control.Monad.Unify.Class (class MonadUnify, unify, reify)
import Control.Monad.Reader.Trans (ReaderT(..), runReaderT)
import Control.Monad.Unify.Trans (UnifyT(..), runUnifyT)
import Control.Monad.Logic.Trans (SFKT(..), runSFKTOnce, unSFKT)
import Control.Monad.State.Trans (StateT(..), runStateT)
import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM)
import Data.Traversable (traverse, oneOf)
import Data.Foldable (oneOfMap)
import Data.Tuple (fst)
import Data.List (List(..), (:), reverse)
import Data.Newtype (class Newtype, un, over2)
import Data.Maybe (Maybe)
import Data.Bifunctor (rmap)

type UnifierEnv = { rules :: Array (Rule Metavariable), fuel :: Int }
newtype Unifier a = MkUnifier (UnifyT IntVar Expr (ReaderT UnifierEnv (SFKT Maybe)) a)

derive instance Newtype (Unifier a) _
derive newtype instance Functor Unifier
derive newtype instance Apply Unifier
derive newtype instance Applicative Unifier
derive newtype instance Alt Unifier
derive newtype instance Plus Unifier
derive newtype instance Alternative Unifier
derive newtype instance Bind Unifier
derive newtype instance Monad Unifier
derive newtype instance MonadPlus Unifier
derive newtype instance MonadUnify IntVar Expr Unifier
derive newtype instance MonadRec Unifier

-- >:(
instance MonadAsk UnifierEnv Unifier where
    ask = MkUnifier $ MkUnifyT ask

-- >:(
instance MonadReader UnifierEnv Unifier where
    local f m = MkUnifier $ MkUnifyT $ StateT $ \s -> local f (runStateT (un MkUnifyT $ un MkUnifier m) s)

-- >:(
instance MonadLogic Unifier where
    msplit m = MkUnifier $ MkUnifyT $ map (map (rmap (MkUnifier <<< MkUnifyT))) $ msplit $ un MkUnifyT $ un MkUnifier m
    interleave = over2 MkUnifier (over2 MkUnifyT interleave)

instance Lazy (Unifier a) where
    defer f = MkUnifier $ MkUnifyT $ StateT $ \s -> ReaderT $ \r -> SFKT (\sk fk -> unSFKT (runReaderT (runStateT (un MkUnifyT $ un MkUnifier $ f unit) s) r) sk fk)

spend :: forall a. Unifier a -> Unifier a
spend m = do
    n <- asks _.fuel
    if n > 0 then local (_ { fuel = n - 1}) m else empty

ints :: forall m. MonadLogic m => Lazy (m Int) => m Int
ints = impl 0
    where
      impl n = pure n <|> defer (\_ -> impl $ n+1)

matchBuiltin :: Expr IntVar -> Unifier (ProofTree IntVar)
matchBuiltin e@(Atom "isInt" (t : Nil)) =
    case t of
        IntLit i -> pure $ isIntProof i
        Var _ -> do
           i <- ints
           _ <- unify t (IntLit i)
           pure $ isIntProof i
        _ -> empty
    where
        isIntRule i = MkRule { head: Atom "isInt" (IntLit i : Nil), tail: Nil }
        isIntProof i = MkProofTree { claim: e, rule: isIntRule i, witnesses: Nil }
matchBuiltin _ = empty

type StackElement =
    { done :: List (ProofTree IntVar)
    , todo :: List (Expr IntVar)
    , claim :: Expr IntVar
    , rule :: Rule Metavariable
    }

type Stack = List StackElement

type Acc =
    { stack :: Stack
    , rules :: Array (Rule Metavariable)
    }

rule :: Expr IntVar -> Rule Metavariable -> Unifier StackElement
rule e r = do
    MkRule {head, tail} <- instantiate r
    _ <- unify e head
    pure $ { done: Nil, todo: tail, claim: e, rule: r }

rules :: Expr IntVar -> Array (Rule Metavariable) -> Unifier StackElement
rules e rs = oneOfMap (rule e) rs

step :: Acc -> Unifier (Step Acc (ProofTree IntVar))
step {stack: Nil} = empty
step acc@{stack: {done, todo: Nil, claim, rule: r} : xs} =
    pure $ case xs of
        Nil -> Done tree
        Cons se xs' -> Loop acc {stack = se { done = tree : se.done } : xs'}
    where
        tree = MkProofTree { claim, rule: r, witnesses: reverse done }
step acc@{stack: se@{todo: (e:es)} : xs} =
    do
        e' <- reify e
        interleave (builtin e') (given e')
    where
        builtin e' = do
            t <- matchBuiltin e'
            pure $ Loop acc { stack = se { todo = es, done = t : se.done } : xs }
        given e' = do
            se' <- rules e' acc.rules
            pure $ Loop acc { stack = se' : se { todo = es } : xs }


-- Note: maybe it's the list / rule operations that are the problem, rather
-- than the stack itself? In particular, could the oneOf be the issue?

match' :: Array (Rule Metavariable) -> Expr IntVar -> Unifier (ProofTree IntVar)
match' rs e = interleave (matchBuiltin e) do
    firstElem <- rules e rs
    tailRecM step { rules: rs, stack: firstElem : Nil }

match :: Array (Rule Metavariable) -> Expr IntVar -> Unifier (ProofTree IntVar)
match rs e = interleave (reify e >>= matchBuiltin) $ oneOfMap (matchSingle e) rs
    where
      matchSingle e' r = spend $ do
         MkRule {head, tail} <- instantiate r
         _ <- unify e' head
         witnesses <- traverse (match rs) tail
         pure $ MkProofTree { claim: e, rule: r, witnesses: witnesses }

reifyProofTree :: ProofTree IntVar -> Unifier (ProofTree IntVar)
reifyProofTree (MkProofTree {claim, rule: r, witnesses}) = do
    claim' <- reify claim
    witnesses' <- traverse reifyProofTree witnesses
    pure $ MkProofTree $ { claim: claim', rule: r, witnesses: witnesses' }

query :: Expr Metavariable -> Unifier (ProofTree IntVar)
query e = (join $ lift2 match' (asks _.rules) (instantiate e)) >>= reifyProofTree

runUnifier :: forall a. Array (Rule Metavariable) -> Unifier a -> Maybe a
runUnifier rs m = runSFKTOnce (fst <$> (runReaderT (runUnifyT $ un MkUnifier m) { rules: rs, fuel: 10 }))