Handle unlifted tycons and tuples correctly during vectorisation
[ghc-hetmet.git] / compiler / vectorise / VectMonad.hs
index dc26b4b..eed5a81 100644 (file)
@@ -1,8 +1,9 @@
 module VectMonad (
+  Scope(..),
   VM,
 
-  noV, tryV, maybeV, orElseV, localV, closedV, initV,
-  newLocalVar, newTyVar,
+  noV, tryV, maybeV, orElseV, fixV, localV, closedV, initV,
+  cloneName, newLocalVar, newTyVar,
   
   Builtins(..), paDictTyCon,
   builtin,
@@ -13,8 +14,10 @@ module VectMonad (
   LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
-  lookupTyCon,
-  lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
+  lookupVar, defGlobalVar,
+  lookupTyCon, defTyCon,
+  lookupDataCon, defDataCon,
+  lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
 
   lookupInst, lookupFamInst
 ) where
@@ -25,10 +28,12 @@ import HscTypes
 import CoreSyn
 import Class
 import TyCon
+import DataCon
 import Type
 import Var
 import VarEnv
 import Id
+import OccName
 import Name
 import NameEnv
 
@@ -42,6 +47,10 @@ import Panic
 import Outputable
 import FastString
 
+import Control.Monad ( liftM )
+
+data Scope a b = Global a | Local b
+
 -- ----------------------------------------------------------------------------
 -- Vectorisation monad
 
@@ -101,9 +110,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
@@ -114,6 +123,9 @@ data GlobalEnv = GlobalEnv {
                 -- instances
                 --
                 , global_fam_inst_env :: FamInstEnvs
+
+                -- Hoisted bindings
+                , global_bindings :: [(Var, CoreExpr)]
                 }
 
 data LocalEnv = LocalEnv {
@@ -122,43 +134,48 @@ 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
-
-                 -- Hoisted bindings
-               , local_bindings :: [(Var, CoreExpr)]
                }
               
 
-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) $ vectInfoCCVar   info
+      global_vars          = mapVarEnv  (Var . snd) $ vectInfoVar   info
     , global_exported_vars = emptyVarEnv
-    , global_tycons        = mapNameEnv snd $ vectInfoCCTyCon 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      = []
     }
 
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
+                 , local_tyvars   = []
                  , local_tyvar_pa = emptyVarEnv
-                 , local_bindings = []
                  }
 
 -- FIXME
 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 updVectInfo env tyenv info
   = info {
-      vectInfoCCVar   = global_exported_vars env
-    , vectInfoCCTyCon = tc_env
+      vectInfoVar     = global_exported_vars 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]]
+    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
 
@@ -189,6 +206,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
@@ -234,6 +256,16 @@ getInstEnv = readGEnv global_inst_env
 getFamInstEnv :: VM FamInstEnvs
 getFamInstEnv = readGEnv global_fam_inst_env
 
+cloneName :: (OccName -> OccName) -> Name -> VM Name
+cloneName mk_occ name = liftM make (liftDs newUnique)
+  where
+    occ_name = mk_occ (nameOccName name)
+
+    make u | isExternalName name = mkExternalName u (nameModule name)
+                                                    occ_name
+                                                    (nameSrcSpan name)
+           | otherwise           = mkSystemName u occ_name
+
 newLocalVar :: FastString -> Type -> VM Var
 newLocalVar fs ty
   = do
@@ -246,17 +278,58 @@ newTyVar fs k
       u <- liftDs newUnique
       return $ mkTyVar (mkSysTvName u fs) k
 
+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
+  = do
+      r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
+      case r of
+        Just e  -> return (Local e)
+        Nothing -> liftM Global
+                 $  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 
 
-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.
 --
@@ -328,7 +401,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)