module Vectorise.Env (
Scope(..),
LocalEnv(..),
emptyLocalEnv,
GlobalEnv(..),
initGlobalEnv,
extendImportedVarsEnv,
setFamEnv,
extendFamEnv,
extendTyConsEnv,
extendDataConsEnv,
extendPAFunsEnv,
setPRFunsEnv,
setBoxedTyConsEnv,
modVectInfo
) where
import HscTypes
import InstEnv
import FamInstEnv
import CoreSyn
import Type
import TyCon
import DataCon
import VarEnv
import VarSet
import Var
import NameSet
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_vect_decls :: VarEnv (Type, CoreExpr)
, global_scalar_vars :: VarSet
, global_scalar_tycons :: NameSet
, global_novect_vars :: 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 -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
initGlobalEnv info vectDecls instEnvs famInstEnvs
= GlobalEnv
{ global_vars = mapVarEnv snd $ vectInfoVar info
, global_vect_decls = mkVarEnv vects
, global_scalar_vars = vectInfoScalarVars info `extendVarSetList` scalars
, global_scalar_tycons = vectInfoScalarTyCons info
, global_novect_vars = mkVarSet novects
, 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 = []
}
where
vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
scalars = [var | Vect var Nothing <- vectDecls]
novects = [var | NoVect var <- vectDecls]
extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
extendImportedVarsEnv ps genv
= genv { global_vars = extendVarEnvList (global_vars genv) ps }
setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
setFamEnv l_fam_inst genv
= genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
where (g_fam_inst, _) = global_fam_inst_env genv
extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
extendFamEnv new genv
= genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
where (g_fam_inst, l_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 }
modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
modVectInfo 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
, vectInfoScalarVars = global_scalar_vars env `minusVarSet` vectInfoScalarVars info
, vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
}
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]]