content: {-# LANGUAGE FlexibleInstances,StandaloneDeriving,
               MultiParamTypeClasses,GeneralizedNewtypeDeriving #-}
  module TypeCheckCommon where
  
  import qualified Data.Map.Strict as M
  import qualified Data.Set as S
  
  import Control.Monad
  import Control.Monad.Except
  import Control.Monad.Identity
  import Control.Monad.Fail (MonadFail (..))
  import Control.Monad.State hiding (modify)
  
  import BwdFwd
  import FreshNames
  import Syntax
  
  newtype Contextual a = Contextual
                         { unCtx :: StateT TCState (FreshMT (Except String)) a}
  --                                       TCState is carried along
  --                                                Fresh variables via 'fresh'
  --                                                         Exceptions possible
  --                                (Except String) o FreshMT o (StateT TCState) a
  
  type IdCmdInfoMap = M.Map Id      (Id,      [TyArg Desugared], [TyArg Desugared], [VType Desugared], VType Desugared)
  --                        cmd-id   itf-id   itf-ty-vars        cmd-ty-vars        cmd-arg-tys        cmd-ret-ty
  --                                          (as MkRTVar's or
  --                                              MkAbRVar's)
  
  type CtrInfoMap = M.Map Id       (Id,    [TyArg Desugared], [VType Desugared])
  --                      ctr-id    dt-id  dt-ty-vars         cmd-arg-tys
  --                                       (as MkRTVar's or
  --                                           MkAbRVar's)
  
  data TCState = MkTCState
    { ctx :: Context
    , amb :: Ab Desugared     -- current ambient
    , cmdMap :: IdCmdInfoMap  -- cmd-id -> (itf-id, itf-ty-vars, cmd-arg-tys, cmd-ret-ty)
    , ctrMap :: CtrInfoMap    -- ctr-id -> (dt-id, dt-ty-vars, cmd-arg-tys)
    , ms :: [Int]         -- each entry represents one type checking phase.
                              --            counts the marks in that phase
    }
  
  deriving instance Functor Contextual
  deriving instance Applicative Contextual
  deriving instance Monad Contextual
  deriving instance MonadState TCState Contextual
  deriving instance MonadError String Contextual
  deriving instance GenFresh Contextual
  
  instance MonadFail Contextual where
    fail = throwError
  
  data Entry = FlexMVar Id Decl
             | TermVar (Operator Typed) (VType Desugared)
             | Mark
             deriving (Show)
  data Decl = Hole
            | TyDefn (VType Desugared)
            | AbDefn (Ab Desugared)
            deriving (Show)
  type Context = Bwd Entry
  type TermBinding = (Operator Typed, VType Desugared)
  type Suffix = [(Id, Decl)]
  
  -- push fresh meta variable on context (corresponds to "freshMeta" in Gundry's thesis)
  freshMVar :: Id -> Contextual Id
  freshMVar x = do n <- fresh
                   let s = trimVar x ++ "$f" ++ show n
                   modify (:< FlexMVar s Hole)
                   return s
  
  -- free meta variables (MkFTVar's, MkAbFVar's)
  fmv :: VType Desugared -> S.Set Id
  fmv (DTTy _ ts _) = foldMap fmvTyArg ts
  fmv (SCTy cty _) = fmvCType cty
  fmv (FTVar x _) = S.singleton x
  fmv (RTVar x _) = S.empty
  fmv (StringTy _) = S.empty
  fmv (IntTy _) = S.empty
  fmv (CharTy _) = S.empty
  
  fmvAb :: Ab Desugared -> S.Set Id
  fmvAb (Ab v (ItfMap m _) _) = S.union (fmvAbMod v) (foldMap (foldMap (foldMap fmvTyArg)) (M.elems m))
  
  fmvTyArg :: TyArg Desugared -> S.Set Id
  fmvTyArg (VArg t _) = fmv t
  fmvTyArg (EArg ab _) = fmvAb ab
  
  fmvAbMod :: AbMod Desugared -> S.Set Id
  fmvAbMod (EmpAb _) = S.empty
  fmvAbMod (AbRVar _ _) = S.empty
  fmvAbMod (AbFVar x _) = S.singleton x
  
  fmvAdj :: Adjustment Desugared -> S.Set Id
  fmvAdj (ConsAdj x ts _) = foldMap fmvTyArg ts
  
  fmvCType :: CType Desugared -> S.Set Id
  fmvCType (CType ps q _) = S.union (foldMap fmvPort ps) (fmvPeg q)
  
  fmvPeg :: Peg Desugared -> S.Set Id
  fmvPeg (Peg ab ty _) = S.union (fmvAb ab) (fmv ty)
  
  fmvPort :: Port Desugared -> S.Set Id
  fmvPort (Port adjs ty _) = S.union (foldMap fmvAdj adjs) (fmv ty)
  
  entrify :: Suffix -> [Entry]
  entrify = map $ uncurry FlexMVar
  
  getContext :: Contextual Context
  getContext = do s <- get
                  return (ctx s)
  
  putContext :: Context -> Contextual ()
  putContext x = do s <- get
                    put $ s { ctx = x }
  
  getAmbient :: Contextual (Ab Desugared)
  getAmbient = do s <- get
                  return (amb s)
  
  putAmbient :: Ab Desugared -> Contextual ()
  putAmbient ab = do s <- get
                     put $ s { amb = ab }
  
  -- ms:  [...]   ->    [0, ...]
  pushMarkCtx :: Contextual ()
  pushMarkCtx = do s <- get
                   put $ s { ms = 0 : ms s }
  
  -- ctx: [...]   ->    [..., Mark]
  -- ms:  h:ts    ->    (succ h):ts
  addMark :: Contextual ()
  addMark = do modify (:< Mark)
               s <- get
               let h = head (ms s)
                   ts = tail (ms s)
               put $ s { ms = succ h : ts }
  
  -- ctx: [..., E, Mark_n, ..., Mark_1, ...]   ->   [..., E]
  -- ms:                                n:ts   ->   ts
  purgeMarks :: Contextual ()
  purgeMarks = do s <- get
                  let n = head (ms s)
                  put $ s { ctx = skim n (ctx s), ms = tail (ms s) }
    where -- delete everything up to (and including) the recent n Mark's
          skim :: Int -> Context -> Context
          skim 0 es = es
          skim n (es :< Mark) = skim (n-1) es
          skim n (es :< _) = skim n es
  
  getCmd :: Id -> Contextual (Id, [TyArg Desugared], [TyArg Desugared], [VType Desugared], VType Desugared)
  getCmd cmd = get >>= \s -> case M.lookup cmd (cmdMap s) of
    Nothing -> error $ "invariant broken: " ++ show cmd ++ " not a command"
    Just (itf, qs, rs, xs, y) -> return (itf, qs, rs, xs, y)
  
  getCmdTyVars :: Id -> Contextual [Id]
  getCmdTyVars cmd = do
    (_, _, cmdTyVars, _, _) <- getCmd cmd
    return $ map filterTyVar cmdTyVars
    where
      filterTyVar :: TyArg Desugared -> Id
      filterTyVar (VArg (RTVar x _) _) = x
      filterTyVar (EArg (Ab (AbRVar x _) _ _) _) = x
      -- TODO: LC: This is a terrible invariant: refactor the way in which commands
      --           (and constructors) are represented in the Contextual. Works for
      --           now though.
  
  -- called "popL" in Gundry's thesis
  popEntry :: Contextual Entry
  popEntry = do es :< e <- getContext
                putContext es
                return e
  
  -- Alter the context by applying the provided function
  modify :: (Context -> Context) -> Contextual ()
  modify f = do ctx <- getContext
                putContext $ f ctx
  
  -- Initialise the TCState, return the input program
  -- Type variables of data type / interface definitions are converted into
  -- rigid ty var args (which can later be made flexible for unification)
  initContextual :: Prog Desugared -> Contextual (Prog Desugared)
  initContextual (MkProg ttms) =
    do mapM_ f [d | DataTm d _ <- ttms] -- init ctrMap
       mapM_ g [i | ItfTm i _ <- ttms]   -- init cmdMap
       mapM_ h [d | DefTm d _ <- ttms]    -- init ctx with [hdr: hdr-type]
       return (MkProg ttms)
    where -- data dt p_1 ... p_n = ctr_1 x_11 ... x_1m
          --                     | ...
          -- For each ctr_i add to ctrMap: ctr_i -> (dt, dt-ty-vars, xs_i)
          f :: DataT Desugared -> Contextual ()
          f (DT dt ps ctrs _) =
            let ps' = map tyVar2rigTyVarArg ps in
              mapM_ (\(Ctr ctr xs _) -> addCtr dt ctr ps' xs) ctrs
  
          -- interface itf p_1 .. p_m =
          --   cmd_1 q_11 ... q_1n: x_11 -> ... -> q_1l -> y_1
          -- | ...
          -- For each cmd_i add to ctrMap: cmd_i -> (itf, itf-ty-vars, cmd-ty-vars, xs_i, y_i)
          g :: Itf Desugared -> Contextual ()
          g (Itf itf ps cmds _) =
            let ps' = map tyVar2rigTyVarArg ps in
              mapM_ (\(Cmd cmd qs xs y _) -> addCmd cmd itf ps' (map tyVar2rigTyVarArg qs) xs y) cmds
  
          -- init context for each handler id of type ty with "id := ty"
          h :: MHDef Desugared -> Contextual ()
          h (Def id ty _ a) = modify (:< TermVar (Poly id a') (SCTy ty a))
            where a' = refToTyped a
  
  initTCState :: TCState
  initTCState = MkTCState BEmp (Ab (EmpAb a) (ItfMap M.empty a) a) M.empty M.empty []
    where a = Desugared Generated
  
  -- Only to be used for initialising the contextual monad
  addCmd :: Id -> Id -> [TyArg Desugared] -> [TyArg Desugared] -> [VType Desugared] -> VType Desugared -> Contextual ()
  addCmd cmd itf ps qs xs q = get >>= \s ->
    put $ s { cmdMap = M.insert cmd (itf, ps, qs, xs, q) (cmdMap s) }
  
  addCtr :: Id -> Id -> [TyArg Desugared] -> [VType Desugared] -> Contextual ()
  addCtr dt     ctr     ts         xs         = get >>= \s ->
  --     dt-id  ctr-id  type-args  value-args
    put $ s { ctrMap = M.insert ctr (dt,ts,xs) (ctrMap s) }
  
  -- Given eff var "x", find assigned ability in context and return "Just" it
  -- if not found, return "Nothing"
  findAbVar :: AbMod Desugared -> Contextual (Maybe (Ab Desugared))
  findAbVar (EmpAb _) = return Nothing
  findAbVar (AbRVar _ _) = return Nothing
  findAbVar (AbFVar x _) = getContext >>= find'
    where find' BEmp = return Nothing
          find' (es :< FlexMVar y (AbDefn ab)) | x == y = return $ Just ab
          find' (es :< _) = find' es
  
  -- Search the context and return true if the identifier has a definition
  isMVarDefined :: Id -> Contextual Bool
  isMVarDefined x = getContext >>= find'
    where find' BEmp = return False
          find' (es :< FlexMVar y dec) | x == y = case dec of
            Hole -> return False
            _    -> return True
          find' (es :< _) = find' es
  
  -- Given ability [e | ...], "findAbVar" and thereby substitute the eff vars e
  -- so long until it is completely unfolded.
  expandAb :: Ab Desugared -> Contextual (Ab Desugared)
  expandAb (Ab v m a) =  do
    ab <- findAbVar v
    case ab of
      Just (Ab v' m' a') -> expandAb $ Ab v' (m' `plusItfMap` m) a
      Nothing ->            return $ Ab v m a
  
  refToTyped :: Desugared -> Typed
  refToTyped (Desugared loc) = Typed loc
  