module Vectorise.Env (
Scope(..),
LocalEnv(..),
emptyLocalEnv,
GlobalEnv(..),
initGlobalEnv,
extendImportedVarsEnv,
extendScalars,
setFamInstEnv,
extendTyConsEnv,
extendDataConsEnv,
extendPAFunsEnv,
setPRFunsEnv,
setBoxedTyConsEnv,
updVectInfo
) where
import HscTypes
import InstEnv
import FamInstEnv
import CoreSyn
import TyCon
import DataCon
import VarEnv
import VarSet
import Var
import Name
import NameEnv
import FastString
data Scope a b
= Global a
| Local b
data LocalEnv
= LocalEnv {
local_vars :: VarEnv (Var, Var)
, local_tyvars :: [TyVar]
, local_tyvar_pa :: VarEnv CoreExpr
, local_bind_name :: FastString
}
emptyLocalEnv :: LocalEnv
emptyLocalEnv = LocalEnv {
local_vars = emptyVarEnv
, local_tyvars = []
, local_tyvar_pa = emptyVarEnv
, local_bind_name = fsLit "fn"
}
data GlobalEnv
= GlobalEnv {
global_vars :: VarEnv Var
, global_scalars :: VarSet
, global_exported_vars :: VarEnv (Var, Var)
, global_tycons :: NameEnv TyCon
, global_datacons :: NameEnv DataCon
, global_pa_funs :: NameEnv Var
, global_pr_funs :: NameEnv Var
, global_boxed_tycons :: NameEnv TyCon
, global_inst_env :: (InstEnv, InstEnv)
, global_fam_inst_env :: FamInstEnvs
, global_bindings :: [(Var, CoreExpr)]
}
initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
initGlobalEnv info instEnvs famInstEnvs
= GlobalEnv
{ global_vars = mapVarEnv snd $ vectInfoVar info
, global_scalars = emptyVarSet
, global_exported_vars = emptyVarEnv
, global_tycons = mapNameEnv snd $ vectInfoTyCon info
, global_datacons = mapNameEnv snd $ vectInfoDataCon info
, global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
, global_pr_funs = emptyNameEnv
, global_boxed_tycons = emptyNameEnv
, global_inst_env = instEnvs
, global_fam_inst_env = famInstEnvs
, global_bindings = []
}
extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
extendImportedVarsEnv ps genv
= genv { global_vars = extendVarEnvList (global_vars genv) ps }
extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
extendScalars vs genv
= genv { global_scalars = extendVarSetList (global_scalars genv) vs }
setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
setFamInstEnv l_fam_inst genv
= genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
where (g_fam_inst, _) = global_fam_inst_env genv
extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
extendTyConsEnv ps genv
= genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
extendDataConsEnv ps genv
= genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
extendPAFunsEnv ps genv
= genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
setPRFunsEnv ps genv
= genv { global_pr_funs = mkNameEnv ps }
setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
setBoxedTyConsEnv ps genv
= genv { global_boxed_tycons = mkNameEnv ps }
updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
updVectInfo env tyenv info
= info
{ vectInfoVar = global_exported_vars env
, vectInfoTyCon = mk_env typeEnvTyCons global_tycons
, vectInfoDataCon = mk_env typeEnvDataCons global_datacons
, vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
}
where
mk_env from_tyenv from_env
= mkNameEnv [(name, (from,to))
| from <- from_tyenv tyenv
, let name = getName from
, Just to <- [lookupNameEnv (from_env env) name]]