{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts #-}

-----------------------------------------------------------------------------
-- Bi-directional type checking for a language with:
--   polymorphism, existentials, recursive types, and mutable references
--
-- Loosely based off of Pfennings 2004 assignment:
--   https://www.cs.cmu.edu/~fp/courses/15312-f04/assignments/asst6/index.html
--
-- Phillip Mates, 2014
-----------------------------------------------------------------------------

module Typechecker where

import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Error
import Control.Applicative

import MSyntax


-----------------------------------------------------------------------------
-- Bi-directional type checking
-----------------------------------------------------------------------------
data TypingCons = TypingCons {
    tyEnv :: MTypeEnv
  , termEnv :: MTermEnv
  , heapEnv :: MHeapEnv
} deriving (Show)

-- monad that has type-checking environments,
-- state for generating unique symbols (not needed), and error handling
newtype TyCheck a = TyCheck {
  runTyCheck :: ReaderT TypingCons (StateT Int (Either String)) a
} deriving
  ( Monad, Applicative, Functor, MonadError String
  , MonadReader TypingCons, MonadState Int)

runTypeChecker :: TyCheck a -> TypingCons -> Int -> Either String (a, Int)
runTypeChecker tych envs n =
  runStateT (runReaderT (runTyCheck tych) envs) n

-- synthesize type for a term
synth :: MExpr -> TyCheck MType
synth (MAnno e τ) = do check e τ >> return τ
synth (MLoc l) = do
  heapVars <- asks heapEnv
  lookupM (MLoc l) heapVars "Location not in heap environment"
synth (MVar x) = do
  termVars <- asks termEnv
  lookupM (MVar x) termVars "Variable not in term environment"
synth (MProj i (MTuple vs)) = do
  τs <- mapM synth vs
  return $ τs !! i
synth (MUnfold v) = do
  ματ <- synth v
  (α, τ) <- coerceMuTyp ματ
  return $ typeSubst ματ α τ
synth (MAssign l v) = do
  τref <- synth l
  τ <- coerceRef τref
  check v τ
  return MUnitTyp
synth (MRaise τ e) = do
  _ <- check e MExceptionTyp
  return τ
synth (MApp vf τ0s vas) = do
  sΔ <- asks tyEnv
  errorGuard (and $ map (wfType sΔ) τ0s)
             ((show τ0s) ++ " not well-formed under " ++ (show sΔ))
  -- synthesize type of function
  τf <- synth vf
  -- make sure it has an arrow shape
  (αs, τas, τr) <- coerceArrow τf
  -- calculate the argument type: τas[τ0s/αs]
  let τasTypeSubbed = map (typeSubstList τ0s αs) τas
  -- check that all the arguments have this type
  _ <- mapM (\(v,τ) -> check v τ) $ zip vas τasTypeSubbed
  -- τr[τ0s/αs]
  return $ typeSubstList τ0s αs τr
synth e = throwError $ "No synthesis case for " ++ (show e)

-- check a term has a certain type
check :: MExpr -> MType -> TyCheck ()
check (MInt _) MIntTyp = return ()
check MUnit MUnitTyp = return ()
check (MNew v) (MRefTyp τ) = check v τ
check (MDeref l) τ = check l (MRefTyp τ)
check (MBinOp _ v1 v2) MIntTyp = do
  check v1 MIntTyp
  check v2 MIntTyp
check (MIf0 v e1 e2) τ = do
  check e1 τ
  check e2 τ
  check v MIntTyp
check (MFunc αs xs e) (MArrowTyp αs' τs τ)
  | map snd xs == τs && αs == αs' = do
    local updateEnvs (check e τ)
    where
      -- extend the type and term environments with list of alphas and x:τ's
      updateEnvs env = env { tyEnv = (tyEnv env) ++ αs
                           , termEnv = (termEnv env) ++ xs}
check (MFold ματ v) (MMuTyp α τ) = check v $ typeSubst ματ α τ
check (MPack τ' v _) (MExistsTyp α τ) = do
  sΔ <- asks tyEnv
  errorGuard (wfType sΔ τ') ((show τ') ++ " not well-formed under " ++ (show sΔ))
  check v $ typeSubst τ' α τ
check (MTuple vs) (MTupleTyp τs) = do
  _ <- mapM (uncurry check) $ zip vs τs
  return ()
check (MUnpack α x v e) τ = do
  exty <- synth v
  (_, τ') <- coerceExists exty
  sΔ <- asks tyEnv
  errorGuard (wfType sΔ τ) ((show τ) ++ " not well-formed under " ++ (show sΔ))
  let updateEnvs env = env { tyEnv = (tyEnv env) ++ [α]
                           , termEnv = (termEnv env) ++ [(x,τ')]}
  local updateEnvs (check e τ)
check (MLet (MVar x) e1 e2) τ2 = do
  τ1 <- synth e1
  -- extend the type and term environments with list of alphas and x:τ's
  let updateEnvs env = env {termEnv = (termEnv env) ++ [((MVar x),τ1)]}
  local updateEnvs (check e2 τ2)
check (MHandle e1 (MVar x) e2) τ2 = do
  τ1 <- synth e1
  -- extend the type and term environments with list of alphas and x:τ's
  let updateEnvs env = env {termEnv = (termEnv env) ++ [((MVar x),MExceptionTyp)]}
  if τ1 == τ2
    then local updateEnvs (check e2 τ2)
    else fail "Type of handler doesn't match the type of the body"
check e MExceptionTyp = check e MIntTyp
check e τ = do
  σ <- synth e
  if σ == τ
    then return ()
    else fail "Types failed to unify between checking and synthesis"

typeCheckProgram :: MExpr -> MType -> Either String Bool
typeCheckProgram e τ =
  case runTypeChecker (check e τ) (TypingCons [] [] []) 0 of
    Left msg -> Left msg
    Right _ -> Right True

-----------------------------------------------------------------------------
-- Type substitution
-----------------------------------------------------------------------------
-- Type level substition: τ -> α -> τ' -> τ'[τ/α]
typeSubst :: MType -> MType -> MType -> MType
typeSubst τ (MVarTyp α) (MVarTyp β) | α == β = τ
                                    | otherwise = MVarTyp β
typeSubst _ (MVarTyp _) MUnitTyp = MUnitTyp
typeSubst _ (MVarTyp _) MIntTyp  = MIntTyp
typeSubst _ (MVarTyp _) MExceptionTyp = MExceptionTyp
typeSubst τ (MVarTyp α) (MArrowTyp αs τs τ') =
  MArrowTyp αs τs' τ''
  where τs' = map (typeSubst τ (MVarTyp α)) τs
        τ'' = typeSubst τ (MVarTyp α) τ'
typeSubst τ (MVarTyp α) (MExistsTyp β τ') = MExistsTyp β (typeSubst τ (MVarTyp α) τ')
typeSubst τ (MVarTyp α) (MMuTyp β τ')     = MMuTyp β (typeSubst τ (MVarTyp α) τ')
typeSubst τ (MVarTyp α) (MRefTyp τ')      = MRefTyp $ typeSubst τ (MVarTyp α) τ'
typeSubst τ (MVarTyp α) (MTupleTyp τs)    = MTupleTyp $ map (typeSubst τ (MVarTyp α)) τs
typeSubst _ _ τ' = τ'

-- Type-level substition over lists: [τ] -> [α] -> τ' -> τ'[τ/α]
typeSubstList :: [MType] -> [MType] -> MType -> MType
typeSubstList τs αs τ = foldr (uncurry typeSubst) τ (zip τs αs)

-----------------------------------------------------------------------------
-- Well-formed types and environments
-----------------------------------------------------------------------------
-- Type is well-formed under type environment
wfType :: MTypeEnv -> MType -> Bool
wfType sΔ (MVarTyp α) = elem (MVarTyp α) sΔ
wfType _ MUnitTyp = True
wfType _ MIntTyp = True
wfType _ MExceptionTyp = True
wfType sΔ (MArrowTyp αs τs τ') =
  and $ (wfType (αs ++ sΔ) τ') :
        map (wfType (αs ++ sΔ)) τs
wfType sΔ (MExistsTyp α τ) = wfType (α : sΔ) τ
wfType sΔ (MMuTyp α τ) = wfType (α : sΔ) τ
wfType sΔ (MRefTyp τ) = wfType sΔ τ
wfType sΔ (MTupleTyp τs) = and $ map (wfType sΔ) τs

-- Term environment has well-formed types
wfTermEnv :: MTypeEnv -> MTermEnv -> Bool
wfTermEnv _ [] = True
wfTermEnv sΔ ((_,τ):sΓ) =
  and [ wfTermEnv sΔ sΓ
      , wfType sΔ τ]

-- Heap environment has well-formed types
wfHeapTypes :: MHeapEnv -> Bool
wfHeapTypes sΨ = and $ map (wfType []) $ map snd sΨ

-----------------------------------------------------------------------------
-- Utils
-----------------------------------------------------------------------------
-- lift the lookup function to MonadError
lookupM :: (MonadError String m, Eq a) => a -> [(a, b)] -> String -> m b
lookupM x pairs msg =
  case lookup x pairs of
    Just τ -> return τ
    Nothing -> throwError msg

-- Raise error message if test is false
errorGuard :: (MonadError String m) => Bool -> String -> m ()
errorGuard True _ = return ()
errorGuard False msg = throwError msg

-- Raise error message if type is not of the expected form
coerceMuTyp :: (MonadError String m) => MType -> m (MType, MType)
coerceMuTyp (MMuTyp α τ) =  return (α, τ)
coerceMuTyp _ = throwError "type synthesized for fold was not a μα.τ type"

coerceExists :: (MonadError String m) => MType -> m (MType, MType)
coerceExists (MExistsTyp α τ) = return (α, τ)
coerceExists _ = throwError "type synthesized for pack was not an existential type"

coerceArrow :: (MonadError String m) => MType -> m ([MType], [MType], MType)
coerceArrow (MArrowTyp αs τas τr) = return (αs, τas, τr)
coerceArrow _ = throwError "type synthesized for function position was not an arrow type"

coerceRef :: (MonadError String m) => MType -> m MType
coerceRef (MRefTyp τ) = return τ
coerceRef _ = throwError "type synthesized for location position was not an ref type"
