Add support for name cloning to vectorisation monad

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

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index a658253..041928d 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
 
@@ -124,15 +131,18 @@ data LocalEnv = LocalEnv {
 
                  -- 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
   = 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
@@ -141,14 +151,15 @@ initGlobalEnv info instEnvs famInstEnvs
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
                  , 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 = 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,18 @@ 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 }
+
+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)