module DDC.Type.Env
( Env(..)
, SuperEnv
, KindEnv
, TypeEnv
, empty
, singleton
, extend
, extends
, union
, unions
, fromList
, fromTypeMap
, depth
, member
, memberBind
, lookup
, lookupName
, setPrimFun
, isPrim
, wrapTForalls
, lift)
where
import DDC.Type.Exp
import DDC.Type.Transform.LiftT
import Data.Maybe
import Data.Map (Map)
import Prelude hiding (lookup)
import qualified Data.Map.Strict as Map
import qualified Prelude as P
import Control.Monad
data Env n
= Env
{
envMap :: !(Map n (Type n))
, envStack :: ![Type n]
, envStackLength :: !Int
, envPrimFun :: !(n -> Maybe (Type n)) }
type SuperEnv n = Env n
type KindEnv n = Env n
type TypeEnv n = Env n
empty :: Env n
empty = Env
{ envMap = Map.empty
, envStack = []
, envStackLength = 0
, envPrimFun = \_ -> Nothing }
singleton :: Ord n => Bind n -> Env n
singleton b
= extend b empty
extend :: Ord n => Bind n -> Env n -> Env n
extend bb env
= case bb of
BName n k -> env { envMap = Map.insert n k (envMap env) }
BAnon k -> env { envStack = k : envStack env
, envStackLength = envStackLength env + 1 }
BNone{} -> env
extends :: Ord n => [Bind n] -> Env n -> Env n
extends bs env
= foldl (flip extend) env bs
setPrimFun :: (n -> Maybe (Type n)) -> Env n -> Env n
setPrimFun f env
= env { envPrimFun = f }
isPrim :: Env n -> n -> Bool
isPrim env n
= isJust $ envPrimFun env n
fromList :: Ord n => [Bind n] -> Env n
fromList bs
= foldr extend empty bs
fromTypeMap :: Map n (Type n) -> Env n
fromTypeMap m
= empty { envMap = m}
union :: Ord n => Env n -> Env n -> Env n
union env1 env2
= Env
{ envMap = envMap env1 `Map.union` envMap env2
, envStack = envStack env2 ++ envStack env1
, envStackLength = envStackLength env2 + envStackLength env1
, envPrimFun = \n -> envPrimFun env2 n `mplus` envPrimFun env1 n }
unions :: Ord n => [Env n] -> Env n
unions envs
= foldr union empty envs
member :: Ord n => Bound n -> Env n -> Bool
member uu env
= isJust $ lookup uu env
memberBind :: Ord n => Bind n -> Env n -> Bool
memberBind uu env
= case uu of
BName n _ -> Map.member n (envMap env)
_ -> False
lookup :: Ord n => Bound n -> Env n -> Maybe (Type n)
lookup uu env
= case uu of
UName n
-> Map.lookup n (envMap env)
`mplus` envPrimFun env n
UIx i -> P.lookup i (zip [0..] (envStack env))
UPrim n _ -> envPrimFun env n
lookupName :: Ord n => n -> Env n -> Maybe (Type n)
lookupName n env
= Map.lookup n (envMap env)
depth :: Env n -> Int
depth env = envStackLength env
lift :: Ord n => Int -> Env n -> Env n
lift n env
= Env
{ envMap = Map.map (liftT n) (envMap env)
, envStack = map (liftT n) (envStack env)
, envStackLength = envStackLength env
, envPrimFun = envPrimFun env }
wrapTForalls :: Ord n => Env n -> Type n -> Type n
wrapTForalls env tBody
= let bsNamed = [BName b t | (b, t) <- Map.toList $ envMap env ]
bsAnon = [BAnon t | t <- envStack env]
tInner = foldr TForall tBody (reverse bsAnon)
in foldr TForall tInner bsNamed