Unification.hs

{-# LANGUAGE FlexibleContexts, OverloadedStrings #-}
module Unification (unifier, fresh, UnifyM, runUnifyM, UnifyError) where
import qualified Miso.String as MS
import Control.Monad
import Control.Monad.State
import Control.Monad.Trans
import Control.Monad.Except
import qualified Data.Map.Strict     as M
import Data.Maybe
import Data.Foldable
import qualified Data.Set            as S

import Terms

-- almost all of this is a direct port of Tobias Nipkow's pattern unification implementation
-- in standard ML.
type UnifyError = MS.MisoString

gen = do x <- get; put (x + 1); pure x

lams (x:xs) t = Lam (M "x") (lams xs t)
lams [] t = t

hnf xs f ss = lams xs (applyApTelescope f ss)

occ f sS (MetaVar g) = (f == g)  || (case lookup g sS of
                                      Just s -> occ f sS s
                                      Nothing -> False)
occ f sS (Ap s t) = occ f sS s || occ f sS t
occ f sS (Lam _ t) = occ f sS t
occ f sS _ = False


mapbnd f = let mpb d (LocalVar i) = LocalVar (if i<d then i else f(i-d)+d)
               mpb d (Lam n t) = Lam n (mpb (d+1) t)
               mpb d (Ap u1 u2) = Ap (mpb d u1) (mpb d u2)
               mpb d t = t
            in mpb 0

incr = mapbnd (+1)

red (Lam _ s) (LocalVar i : bs) js = red s bs (i:js)
red s         bs                js = applyApTelescope (mapbnd (js !!) s) bs;

devar sS s | (MetaVar f, is) <- peelApTelescope s
           , Just t <- lookup f sS = devar sS (red t is [])
devar sS s = s

pos :: (a -> Bool) -> [a] -> [Term]
pos p (x:xs) = if p x then LocalVar(length xs) : pos p xs else pos p xs
pos p [] = []

posM p (x:xs) = do c <- p x; if c then (LocalVar(length xs):) <$>  posM p xs else posM p xs
posM p [] = pure []
idx (b:bs) b' = if b==b' then LocalVar(length bs) else idx bs b'
idx [] _ = LocalVar(-10000)


fresh = MetaVar <$> lift gen

proj sS s = case peelApTelescope (devar sS s) of
        (Lam _ t,_) -> proj sS t
        (Const _ _,ss) -> foldlM proj sS ss
        (LocalVar i,ss) | i >= 0 -> foldlM proj sS ss
                        | otherwise -> throwError "Unification Failure i < 0"
        (MetaVar f,bs) -> do
           var <- fresh
           bs' <- posM ( \ t -> case t of (LocalVar i) -> pure (i >= 0)
                                          otherthing  -> throwError "Non-pattern equation" ) bs
           pure ((f , hnf bs var bs' ):sS)


flexflex1 f ym zn sS
    | ym == zn =  pure $ sS
    | otherwise = do
       var <- fresh
       pure ((f, hnf ym var (pos (uncurry (==)) (zip ym zn))) : sS)


subset :: (Eq a) => [a] -> [a] -> Bool
subset as bs = all (`elem` bs) as

intersection :: (Eq a) => [a] -> [a] -> [a]
intersection xs ys = filter (`elem` ys) xs

flexflex2 f im g jn sS
    | im `subset` jn = pure $ ((g, lam' jn (MetaVar f) im) : sS )
    | jn `subset` im = pure $ ((f, lam' im (MetaVar g) jn) : sS)
    | otherwise = do
         let kl = im `intersection` jn
         h <- fresh
         pure ((f, lam' im h kl ) : (g, lam' jn h kl ) : sS)
  where
    lam' im g jn = hnf im g (map (idx im) jn)


flexflex f ym g zn sS
  | f == g = flexflex1 f ym zn sS
  | otherwise = flexflex2 f ym g zn sS

flexrigid f im t sS
  | occ f sS t = throwError "Unification failure (occurs check)"
  | otherwise  = let u = mapbnd (\i -> let (LocalVar n) = idx im (LocalVar i) in n) t
                  in proj((f,lams im u):sS) u

unif sS (s,t) = case (devar sS s,devar sS t) of
        (Lam _ s, Lam _ t) -> unif sS (s,t)
        (Lam _ s,t) -> unif sS (s, Ap (incr t) (LocalVar 0))
        (s,Lam _ t) -> unif sS (Ap (incr s) (LocalVar 0), t)
        (s,t) -> cases sS (s,t)

cases sS (s,t) = case (peelApTelescope s,peelApTelescope t) of
        ((MetaVar f,ym),(MetaVar g,zn)) -> flexflex f ym g zn sS
        ((MetaVar f,ym),_) -> flexrigid f ym t sS
        (_,(MetaVar f,ym)) -> flexrigid f ym s sS
        ((a,sm),(b,tn)) -> rigidrigid a sm b tn sS

rigidrigid a ss b ts sS
  | a == b, length ss == length ts  = foldlM unif sS (zip ss ts)
  | otherwise = throwError (MS.toMisoString (concat ["Unification Error (", show a, ", ", show b, ")"])) --"Unification Error (rigid, rigid)"

type Gen = State Int

type UnifyM = ExceptT MS.MisoString (State Int)



unifier :: Term -> Term -> UnifyM Subst
unifier t1 t2 = fromUnifier <$> unif [] (t1,t2)

runUnifyM = runState . runExceptT