Remove mapping from tycons to PA dictionaries from vect environment

[ghc-hetmet.git] / compiler / vectorise / VectMonad.hs

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index fee294f..2b7a7ca 100644 (file)
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -2,7 +2,7 @@ module VectMonad (
   Scope(..),
   VM,
 
-  noV, tryV, maybeV, orElseV, localV, closedV, initV,
+  noV, tryV, maybeV, orElseV, fixV, localV, closedV, initV,
   cloneName, newLocalVar, newTyVar,
   
   Builtins(..), paDictTyCon,
@@ -16,7 +16,7 @@ module VectMonad (
 
   defGlobalVar, lookupVar,
   lookupTyCon,
-  lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
+  lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
 
   lookupInst, lookupFamInst
 ) where
@@ -27,6 +27,7 @@ import HscTypes
 import CoreSyn
 import Class
 import TyCon
+import DataCon
 import Type
 import Var
 import VarEnv
@@ -108,9 +109,9 @@ data GlobalEnv = GlobalEnv {
                   --
                 , global_tycons :: NameEnv TyCon
 
-                  -- Mapping from TyCons to their PA dictionaries
+                  -- Mapping from DataCons to their vectorised versions
                   --
-                , global_tycon_pa :: NameEnv CoreExpr
+                , global_datacons :: NameEnv DataCon
 
                 -- External package inst-env & home-package inst-env for class
                 -- instances
@@ -132,6 +133,10 @@ data LocalEnv = LocalEnv {
                  --
                  local_vars :: VarEnv (CoreExpr, CoreExpr)
 
+                 -- In-scope type variables
+                 --
+               , local_tyvars :: [TyVar]
+
                  -- Mapping from tyvars to their PA dictionaries
                , local_tyvar_pa :: VarEnv CoreExpr
                }
@@ -143,7 +148,7 @@ initGlobalEnv info instEnvs famInstEnvs
       global_vars          = mapVarEnv  (Var . snd) $ vectInfoVar   info
     , global_exported_vars = emptyVarEnv
     , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
-    , global_tycon_pa      = emptyNameEnv
+    , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
     , global_inst_env      = instEnvs
     , global_fam_inst_env  = famInstEnvs
     , global_bindings      = []
@@ -151,6 +156,7 @@ initGlobalEnv info instEnvs famInstEnvs
 
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
+                 , local_tyvars   = []
                  , local_tyvar_pa = emptyVarEnv
                  }
 
@@ -158,8 +164,9 @@ emptyLocalEnv = LocalEnv {
 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 updVectInfo env tyenv info
   = info {
-      vectInfoVar   = global_exported_vars env
-    , vectInfoTyCon = tc_env
+      vectInfoVar     = global_exported_vars env
+    , vectInfoTyCon   = tc_env
+    , vectInfoDataCon = dc_env
     }
   where
     tc_env = mkNameEnv [(tc_name, (tc,tc'))
@@ -167,6 +174,11 @@ updVectInfo env tyenv info
                , let tc_name = tyConName tc
                , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
 
+    dc_env = mkNameEnv [(dc_name, (dc,dc'))
+               | dc <- typeEnvDataCons tyenv
+               , let dc_name = dataConName dc
+               , Just dc' <- [lookupNameEnv (global_datacons env) dc_name]]
+
 data VResult a = Yes GlobalEnv LocalEnv a | No
 
 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
@@ -196,6 +208,11 @@ maybeV p = maybe noV return =<< p
 orElseV :: VM a -> VM a -> VM a
 orElseV p q = maybe q return =<< tryV p
 
+fixV :: (a -> VM a) -> VM a
+fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
+  where
+    unYes (Yes _ _ x) = x
+
 localV :: VM a -> VM a
 localV p = do
              env <- readLEnv id
@@ -263,8 +280,14 @@ newTyVar fs k
       u <- liftDs newUnique
       return $ mkTyVar (mkSysTvName u fs) k
 
-defGlobalVar :: Var -> CoreExpr -> VM ()
-defGlobalVar v e = updGEnv $ \env -> env { global_vars = extendVarEnv (global_vars env) v e }
+defGlobalVar :: Var -> Var -> VM ()
+defGlobalVar v v' = updGEnv $ \env ->
+  env { global_vars = extendVarEnv (global_vars env) v (Var v')
+      , global_exported_vars = upd (global_exported_vars env)
+      }
+  where
+    upd env | isExportedId v = extendVarEnv env v (v, v')
+            | otherwise      = env
 
 lookupVar :: Var -> VM (Scope CoreExpr (CoreExpr, CoreExpr))
 lookupVar v
@@ -281,11 +304,20 @@ lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName
 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv 
 
-extendTyVarPA :: Var -> CoreExpr -> VM ()
-extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
+defLocalTyVar :: TyVar -> VM ()
+defLocalTyVar tv = updLEnv $ \env ->
+  env { local_tyvars   = tv : local_tyvars env
+      , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
+      }
+
+defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
+defLocalTyVarWithPA tv pa = updLEnv $ \env ->
+  env { local_tyvars   = tv : local_tyvars env
+      , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
+      }
 
-deleteTyVarPA :: Var -> VM ()
-deleteTyVarPA tv = updLEnv $ \env -> env { local_tyvar_pa = delVarEnv (local_tyvar_pa env) tv }
+localTyVars :: VM [TyVar]
+localTyVars = readLEnv (reverse . local_tyvars)
 
 -- Look up the dfun of a class instance.
 --