Vectorisation of top-level bindings

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

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index a658253..68966a1 100644 (file)
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -1,8 +1,9 @@
 module VectMonad (
+  Scope(..),
   VM,
 
   noV, tryV, maybeV, orElseV, localV, closedV, initV,
-  newLocalVar, newTyVar,
+  cloneName, newLocalVar, newTyVar,
   
   Builtins(..), paDictTyCon,
   builtin,
@@ -13,6 +14,7 @@ module VectMonad (
   LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
+  defGlobalVar, lookupVar,
   lookupTyCon,
   lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
 
@@ -29,6 +31,7 @@ import Type
 import Var
 import VarEnv
 import Id
+import OccName
 import Name
 import NameEnv
 
@@ -42,6 +45,10 @@ import Panic
 import Outputable
 import FastString
 
+import Control.Monad ( liftM )
+
+data Scope a b = Global a | Local b
+
 -- ----------------------------------------------------------------------------
 -- Vectorisation monad
 
@@ -114,6 +121,9 @@ data GlobalEnv = GlobalEnv {
                 -- instances
                 --
                 , global_fam_inst_env :: FamInstEnvs
+
+                -- Hoisted bindings
+                , global_bindings :: [(Var, CoreExpr)]
                 }
 
 data LocalEnv = LocalEnv {
@@ -130,12 +140,13 @@ data LocalEnv = LocalEnv {
 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 initGlobalEnv info instEnvs famInstEnvs
   = 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_tycons        = mapNameEnv snd $ vectInfoTyCon info
     , global_tycon_pa      = emptyNameEnv
     , global_inst_env      = instEnvs
     , global_fam_inst_env  = famInstEnvs
+    , global_bindings      = []
     }
 
 emptyLocalEnv = LocalEnv {
@@ -147,8 +158,8 @@ emptyLocalEnv = LocalEnv {
 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 updVectInfo env tyenv info
   = info {
-      vectInfoCCVar   = global_exported_vars env
-    , vectInfoCCTyCon = tc_env
+      vectInfoVar   = global_exported_vars env
+    , vectInfoTyCon = tc_env
     }
   where
     tc_env = mkNameEnv [(tc_name, (tc,tc'))
@@ -230,6 +241,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
@@ -242,6 +263,24 @@ 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)