X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectMonad.hs;h=09e2d2ff3ae54dc62d876a98e3b9a2e5c2dbfabe;hb=24901afd71ec4776b2949f38c87103eb2cda2985;hp=b8f51a2195ba6261b396d50e55e9279b1ef53afd;hpb=d726d04e001bdc5b0a91eecc41aa56123068e362;p=ghc-hetmet.git

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index b8f51a2..09e2d2f 100644
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -9,13 +9,15 @@ module VectMonad (
   builtin,
 
   GlobalEnv(..),
+  setInstEnvs,
   readGEnv, setGEnv, updGEnv,
 
   LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
-  defGlobalVar, lookupVar,
-  lookupTyCon,
+  lookupVar, defGlobalVar,
+  lookupTyCon, defTyCon,
+  lookupDataCon, defDataCon,
   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
 
   lookupInst, lookupFamInst
@@ -109,10 +111,6 @@ data GlobalEnv = GlobalEnv {
                   --
                 , global_tycons :: NameEnv TyCon
 
-                  -- Mapping from TyCons to their PA dictionaries
-                  --
-                , global_tycon_pa :: NameEnv CoreExpr
-
                   -- Mapping from DataCons to their vectorised versions
                   --
                 , global_datacons :: NameEnv DataCon
@@ -146,19 +144,27 @@ data LocalEnv = LocalEnv {
                }
               
 
-initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
-initGlobalEnv info instEnvs famInstEnvs
+initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> Builtins -> GlobalEnv
+initGlobalEnv info instEnvs famInstEnvs bi
   = GlobalEnv {
       global_vars          = mapVarEnv  (Var . snd) $ vectInfoVar   info
     , global_exported_vars = emptyVarEnv
-    , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
-    , global_tycon_pa      = emptyNameEnv
+    , global_tycons        = extendNameEnv (mapNameEnv snd (vectInfoTyCon info))
+                                           (tyConName funTyCon) (closureTyCon bi)
+                              
     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
     , global_inst_env      = instEnvs
     , global_fam_inst_env  = famInstEnvs
     , global_bindings      = []
     }
 
+setInstEnvs :: InstEnv -> FamInstEnv -> GlobalEnv -> GlobalEnv
+setInstEnvs l_inst l_fam_inst genv
+  | (g_inst,     _) <- global_inst_env genv
+  , (g_fam_inst, _) <- global_fam_inst_env genv
+  = genv { global_inst_env     = (g_inst, l_inst)
+         , global_fam_inst_env = (g_fam_inst, l_fam_inst) }
+
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
                  , local_tyvars   = []
@@ -170,19 +176,14 @@ updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 updVectInfo env tyenv info
   = info {
       vectInfoVar     = global_exported_vars env
-    , vectInfoTyCon   = tc_env
-    , vectInfoDataCon = dc_env
+    , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
+    , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
     }
   where
-    tc_env = mkNameEnv [(tc_name, (tc,tc'))
-               | tc <- typeEnvTyCons tyenv
-               , 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]]
+    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]]
 
 data VResult a = Yes GlobalEnv LocalEnv a | No
 
@@ -304,7 +305,21 @@ lookupVar v
                  $  maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 
 lookupTyCon :: TyCon -> VM (Maybe TyCon)
-lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
+lookupTyCon tc
+  | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
+
+  | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
+
+defTyCon :: TyCon -> TyCon -> VM ()
+defTyCon tc tc' = updGEnv $ \env ->
+  env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
+
+lookupDataCon :: DataCon -> VM (Maybe DataCon)
+lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
+
+defDataCon :: DataCon -> DataCon -> VM ()
+defDataCon dc dc' = updGEnv $ \env ->
+  env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 
 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv 
@@ -394,7 +409,7 @@ initV hsc_env guts info p
     go instEnvs famInstEnvs = 
       do
         builtins <- initBuiltins
-        r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs) 
+        r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs builtins) 
                    emptyLocalEnv
         case r of
           Yes genv _ x -> return $ Just (new_info genv, x)