TypeCheckCommon.hs

{-# 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