Unification.hs

-- Unification implementation inspired heavily by Adam Gundry's thesis.
module Unification where

import qualified Data.Map as M
import qualified Data.Set as S

import Control.Monad.Except

import BwdFwd
import FreshNames
import Syntax
import TypeCheckCommon
import Debug

data Extension = Restore | Replace Suffix

restore :: Contextual Extension
restore = return Restore

replace :: Suffix -> Contextual Extension
replace = return . Replace

-- 1) Go through context until the first FlexMVar x:=d appears.
-- 2) Run "f" on "x:=d", resulting in either "Restore" or "Replace-by-ext"
onTop :: (Id -> Decl -> Contextual Extension) -> Contextual ()
onTop f = popEntry >>= focus
  where focus :: Entry -> Contextual ()
        focus e@(FlexMVar x d) =
          do m <- f x d
             case m of
               Replace ext -> modify (<>< entrify ext)
               Restore -> modify (:< e)
        focus e = onTop f >> modify (:< e)

-- unify 2 val tys in current context
unify :: VType Desugared -> VType Desugared -> Contextual ()                    -- corresponding rule in Gundry's thesis:
unify t0 t1 = do logBeginUnify t0 t1
                 unify' t0 t1
                 logEndUnify t0 t1
  where
  unify' (DTTy dt0 ts0 _) (DTTy dt1 ts1 _)                                      -- decompose (modified)
    | dt0 == dt1 && eqLens ts0 ts1 =
      zipWithM_ unifyTyArg ts0 ts1
  unify' v1@(DTTy dt0 ts0 _) v2@(DTTy dt1 ts1 _)                                -- decompose fails
    | not $ eqLens ts0 ts1 =
    throwError $ errorUnifDiffNumberArgs v1 v2
  unify' (SCTy cty0 _)   (SCTy cty1 _)        = unifyCType cty0 cty1
  unify' (RTVar a _)     (RTVar b _) | a == b = return ()
  unify' (IntTy _)       (IntTy _)            = return ()
  unify' (CharTy _)      (CharTy _)           = return ()
  unify' (FloatTy _)     (FloatTy _)          = return ()
  unify' fta@(FTVar a _) ftb@(FTVar b _)      = onTop $ \c d ->
    cmp (a == c) (b == c) d
    where cmp :: Bool -> Bool -> Decl -> Contextual Extension
          cmp True  True  _           = restore                                 -- idle
          cmp True  False Hole        = replace [(a, TyDefn ftb)]               -- define
          cmp False True  Hole        = replace [(b, TyDefn fta)]               -- define
          cmp True  False (TyDefn ty) = unify ty ftb >> restore                 -- subs
          cmp False True  (TyDefn ty) = unify fta ty >> restore                 -- subs
          cmp False False _           = unify fta ftb >> restore                -- skip-ty
          cmp _     _     (AbDefn _)  = error "unification invariant broken"
  unify' (FTVar x a)     ty                   = solve x a [] ty                 -- inst
  unify' ty              (FTVar y b)          = solve y b [] ty                 -- inst
  unify' t               s                    = throwError $ errorUnifTys t s

-- unify 2 eff tys in current context
unifyAb :: Ab Desugared -> Ab Desugared -> Contextual ()
unifyAb ab0@(Ab v0 m0 a) ab1@(Ab v1 m1 b) =
  do logBeginUnifyAb ab0 ab1
     -- expand abs by resolving all flexible variables
     ab0' <- expandAb ab0
     ab1' <- expandAb ab1
     -- then unify
     unifyAb' ab0' ab1'
     logEndUnifyAb ab0 ab1
  where unifyAb' :: Ab Desugared -> Ab Desugared -> Contextual ()
        -- Same eff ty vars leaves nothing to unify but instantiat's m0, m1
        unifyAb' ab0@(Ab v0 m1 _) ab1@(Ab v1 m2 _) | strip v0 == strip v1 =
          catchError (unifyItfMap m1 m2) (unifyAbError ab0 ab1)
        -- Both eff ty vars are flexible
        unifyAb' (Ab (AbFVar x a') m1 a) (Ab (AbFVar y b') m2 b) =
          do -- For same occurrences of interfaces, their instantiat's must coincide
             unifyItfMap (intersectItfMap m1 m2) (intersectItfMap m2 m1)
             v <- AbFVar <$> freshMVar "£" <*> pure (Desugared Generated)
             solveForEVar x a' [] (Ab v (m2 `cutItfMapSuffix` m1) a)  -- TODO: LC: locations assigned correctly?
             solveForEVar y b' [] (Ab v (m1 `cutItfMapSuffix` m2) b)
        -- One eff ty var is flexible, the other one either empty or rigid
        unifyAb' (Ab (AbFVar x a') m1 _) (Ab v m2 b)
          | isItfMapEmpty (cutItfMapSuffix m1 m2) =
            do unifyItfMap (intersectItfMap m1 m2) (intersectItfMap m2 m1)
               solveForEVar x a' [] (Ab v (m2 `cutItfMapSuffix` m1) b)   -- TODO: LC: locations assigned correctly?
        unifyAb' (Ab v m1 a) (Ab (AbFVar y b') m2 _)
          | isItfMapEmpty (cutItfMapSuffix m2 m1) =
            do unifyItfMap (intersectItfMap m1 m2) (intersectItfMap m2 m1)
               solveForEVar y b' [] (Ab v (m1 `cutItfMapSuffix` m2) a)   -- TODO: LC: locations assigned correctly?
        -- In any other case
        unifyAb' ab0 ab1 = throwError $ errorUnifAbs ab0 ab1

unifyTyArg :: TyArg Desugared -> TyArg Desugared -> Contextual ()
unifyTyArg (VArg t0 _) (VArg t1 _) = unify t0 t1
unifyTyArg (EArg ab0 _) (EArg ab1 _) = unifyAb ab0 ab1
unifyTyArg arg0 arg1 = throwError $ errorUnifTyArgs arg0 arg1

unifyAbError :: Ab Desugared -> Ab Desugared -> String -> Contextual ()
unifyAbError ab0 ab1 _ = throwError $ errorUnifAbs ab0 ab1

-- Given two interface maps, check that each instantiation has a unifiable
-- counterpart in the other ability and unify.
unifyItfMap :: ItfMap Desugared -> ItfMap Desugared -> Contextual ()
unifyItfMap m0@(ItfMap m0' a) m1@(ItfMap m1' b) = do
  mapM_ (unifyItfMap' m1) (M.toList m0')
  mapM_ (unifyItfMap' m0) (M.toList m1')
  where unifyItfMap' :: ItfMap Desugared -> (Id, Bwd [TyArg Desugared]) -> Contextual ()
        unifyItfMap' (ItfMap m _) (x, insts) = case M.lookup x m of
          Nothing -> throwError $ errorUnifItfMaps m0 m1
          Just insts' -> if length insts /= length insts' then throwError $ errorUnifItfMaps m0 m1
                         else zipWithM_ (zipWithM_ unifyTyArg) (bwd2fwd insts) (bwd2fwd insts')

unifyCType :: CType Desugared -> CType Desugared -> Contextual ()
unifyCType (CType xs p0 _) (CType ys p1 _) =
  zipWithM_ unifyPort xs ys >> unifyPeg p0 p1

unifyPeg :: Peg Desugared -> Peg Desugared -> Contextual ()
unifyPeg (Peg ab0 ty0 _) (Peg ab1 ty1 _) = unifyAb ab0 ab1 >> unify ty0 ty1

-- Two ports unify if
-- 1) their adjustments coincide (first normalise both)
-- 2) their types unify
unifyPort :: Port Desugared -> Port Desugared -> Contextual ()
unifyPort (Port adjs1 ty1 _) (Port adjs2 ty2 _) =
  do let (insts1, adps1) = adjsNormalForm adjs1
     let (insts2, adps2) = adjsNormalForm adjs2
     unifyItfMap (ItfMap insts1 (Desugared Generated))  -- TODO: LC: fix this
                 (ItfMap insts2 (Desugared Generated))
     if adps1 == adps2 then
       unify ty1 ty2
     else throwError $ "adaptors not the same"

-- unify a meta variable "x" with a type "ty"
solve :: Id -> Desugared -> Suffix -> VType Desugared -> Contextual ()
solve x a ext ty = do logBeginSolve x ext ty
                      solve'
                      logEndSolve x ext ty
  where
  solve' = onTop $ \y d -> do
    logSolveStep (x == y, S.member y (fmv ty), d)
    case (x == y, S.member y (fmv ty), d) of
      (_,     _,     TyDefn bty) -> modify (<>< entrify ext) >>                 -- inst-subs (val ty var)
                                    unify (subst bty y (FTVar x a)) (subst bty y ty) >>
                                    restore
      (_,     _,     AbDefn ab) ->  modify (<>< entrify ext) >>                 -- inst-subs (eff ty var)
                                    unify (substEVar ab y (FTVar x a)) (substEVar ab y ty) >>
                                    restore
      (True,  True,  Hole) -> throwError errorUnifSolveOccurCheck
      (True,  False, Hole) -> replace (ext ++ [(x, TyDefn ty)])                 -- inst-define
      (False, True,  Hole) -> solve x a ((y, d):ext) ty >> replace []           -- inst-depend
      (False, False, Hole) -> solve x a ext ty >> restore                       -- inst-skip-ty

solveForEVar :: Id -> Desugared -> Suffix -> Ab Desugared -> Contextual ()
solveForEVar x a ext ab = onTop $ \y d ->
  case (x == y, S.member y (fmvAb ab), d) of
    (_, _, AbDefn ab') ->
      let vab = Ab (AbFVar x a) (ItfMap M.empty a) a in
      modify (<>< entrify ext) >>
      unifyAb (substEVarAb ab' y vab) (substEVarAb ab' y ab) >>
      restore
    (True, True, Hole) -> throwError errorUnifSolveOccurCheck
    (True, False, Hole) -> replace (ext ++ [(x, AbDefn ab)])
    (False, True, _) -> solveForEVar x a ((y, d):ext) ab >> replace []
    (False, False, _) -> solveForEVar x a ext ab >> restore
    (_, _, TyDefn _) ->
      error "solveForEVar invariant broken: reached impossible case"

{- Substitute "ty" for "x" in X -}

subst :: VType Desugared -> Id -> VType Desugared -> VType Desugared
subst ty x (DTTy dt ts a) = DTTy dt (map (substTyArg ty x) ts) a
subst ty x (SCTy cty a) = SCTy (substCType ty x cty) a
subst ty x (FTVar y a) | x == y = ty
subst _ _ ty = ty

substAb :: VType Desugared -> Id -> Ab Desugared -> Ab Desugared
substAb ty x (Ab v (ItfMap m a') a) = Ab v (ItfMap m' a') a
  where m' = fmap (fmap (map (substTyArg ty x))) m

substTyArg :: VType Desugared -> Id -> TyArg Desugared -> TyArg Desugared
substTyArg ty x (VArg t a) = VArg (subst ty x t) a
substTyArg ty x (EArg ab a) = EArg (substAb ty x ab) a

substAdj :: VType Desugared -> Id -> Adjustment Desugared -> Adjustment Desugared
substAdj ty x (ConsAdj y ts a) = ConsAdj y (map (substTyArg ty x) ts) a
substAdj ty x (AdaptorAdj adp a) = AdaptorAdj adp a

substCType :: VType Desugared -> Id -> CType Desugared -> CType Desugared
substCType ty x (CType ps peg a) =
  CType (map (substPort ty x) ps) (substPeg ty x peg) a

substPeg :: VType Desugared -> Id -> Peg Desugared -> Peg Desugared
substPeg ty x (Peg ab pty a) = Peg (substAb ty x ab) (subst ty x pty) a

substPort :: VType Desugared -> Id -> Port Desugared -> Port Desugared
substPort ty x (Port adjs pty a) = Port (map (substAdj ty x) adjs) (subst ty x pty) a

substEVar :: Ab Desugared -> Id -> VType Desugared -> VType Desugared
substEVar ab x (DTTy dt ts a) = DTTy dt (map (substEVarTyArg ab x) ts) a
substEVar ab x (SCTy cty a) = SCTy (substEVarCType ab x cty) a
substEVar _ _ ty = ty

-- substitute "ab" for "x" in AB
substEVarAb :: Ab Desugared -> Id -> Ab Desugared -> Ab Desugared
substEVarAb ab@(Ab v m' _) x (Ab (AbFVar y a) (ItfMap m ann') ann) | x == y =
  Ab v m'' ann
  where m'' = plusItfMap m' (ItfMap (fmap (fmap (map (substEVarTyArg ab x))) m) ann')
substEVarAb ab x (Ab v (ItfMap m a') a) = Ab v (ItfMap (M.map (fmap (map (substEVarTyArg ab x))) m) a') a

substEVarTyArg :: Ab Desugared -> Id -> TyArg Desugared -> TyArg Desugared
substEVarTyArg ab x (VArg t a)  = VArg (substEVar ab x t) a
substEVarTyArg ab x (EArg ab' a) = EArg (substEVarAb ab x ab') a

substEVarAdj :: Ab Desugared -> Id -> Adjustment Desugared -> Adjustment Desugared
substEVarAdj ab x (ConsAdj y ts a) = ConsAdj y (map (substEVarTyArg ab x) ts) a
substEvarAdj ab x (AdaptorAdj adp a) = AdaptorAdj adp a

substEVarCType :: Ab Desugared -> Id -> CType Desugared -> CType Desugared
substEVarCType ab x (CType ps peg a) =
  CType (map (substEVarPort ab x) ps) (substEVarPeg ab x peg) a

substEVarPeg :: Ab Desugared -> Id -> Peg Desugared -> Peg Desugared
substEVarPeg ab' x (Peg ab pty a) =
  Peg (substEVarAb ab' x ab) (substEVar ab' x pty) a

substEVarPort :: Ab Desugared -> Id -> Port Desugared -> Port Desugared
substEVarPort ab x (Port adjs pty a) =
  Port (map (substEVarAdj ab x) adjs) (substEVar ab x pty) a

{- Helpers -}

eqLens :: [a] -> [a] -> Bool
eqLens xs ys = length xs == length ys