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 ab77037..041928d 100644 (file)
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -1,10 +1,11 @@
 module VectMonad (
+  Scope(..),
   VM,
 
-  noV, tryV, maybeV, orElseV, localV, initV,
-  newLocalVar, newTyVar,
+  noV, tryV, maybeV, orElseV, localV, closedV, initV,
+  cloneName, newLocalVar, newTyVar,
   
-  Builtins(..),
+  Builtins(..), paDictTyCon,
   builtin,
 
   GlobalEnv(..),
@@ -13,7 +14,9 @@ module VectMonad (
   LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
-  lookupTyCon, extendTyVarPA,
+  defGlobalVar, lookupVar,
+  lookupTyCon,
+  lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
 
   lookupInst, lookupFamInst
 ) where
@@ -28,6 +31,7 @@ import Type
 import Var
 import VarEnv
 import Id
+import OccName
 import Name
 import NameEnv
 
@@ -41,46 +45,50 @@ import Panic
 import Outputable
 import FastString
 
+import Control.Monad ( liftM )
+
+data Scope a b = Global a | Local b
+
 -- ----------------------------------------------------------------------------
 -- Vectorisation monad
 
 data Builtins = Builtins {
                   parrayTyCon      :: TyCon
-                , paTyCon          :: TyCon
+                , paClass          :: Class
                 , closureTyCon     :: TyCon
                 , mkClosureVar     :: Var
                 , applyClosureVar  :: Var
                 , mkClosurePVar    :: Var
                 , applyClosurePVar :: Var
-                , closurePAVar     :: Var
                 , lengthPAVar      :: Var
                 , replicatePAVar   :: Var
                 }
 
+paDictTyCon :: Builtins -> TyCon
+paDictTyCon = classTyCon . paClass
+
 initBuiltins :: DsM Builtins
 initBuiltins
   = do
       parrayTyCon  <- dsLookupTyCon parrayTyConName
-      paTyCon      <- dsLookupTyCon paTyConName
+      paClass      <- dsLookupClass paClassName
       closureTyCon <- dsLookupTyCon closureTyConName
 
       mkClosureVar     <- dsLookupGlobalId mkClosureName
       applyClosureVar  <- dsLookupGlobalId applyClosureName
       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
       applyClosurePVar <- dsLookupGlobalId applyClosurePName
-      closurePAVar     <- dsLookupGlobalId closurePAName
       lengthPAVar      <- dsLookupGlobalId lengthPAName
       replicatePAVar   <- dsLookupGlobalId replicatePAName
 
       return $ Builtins {
                  parrayTyCon      = parrayTyCon
-               , paTyCon          = paTyCon
+               , paClass          = paClass
                , closureTyCon     = closureTyCon
                , mkClosureVar     = mkClosureVar
                , applyClosureVar  = applyClosureVar
                , mkClosurePVar    = mkClosurePVar
                , applyClosurePVar = applyClosurePVar
-               , closurePAVar     = closurePAVar
                , lengthPAVar      = lengthPAVar
                , replicatePAVar   = replicatePAVar
                }
@@ -123,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
@@ -140,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'))
@@ -191,6 +203,14 @@ localV p = do
              setLEnv env
              return x
 
+closedV :: VM a -> VM a
+closedV p = do
+              env <- readLEnv id
+              setLEnv emptyLocalEnv
+              x <- p
+              setLEnv env
+              return x
+
 liftDs :: DsM a -> VM a
 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 
@@ -221,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
@@ -233,12 +263,30 @@ 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)
 
+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 }
 
+deleteTyVarPA :: Var -> VM ()
+deleteTyVarPA tv = updLEnv $ \env -> env { local_tyvar_pa = delVarEnv (local_tyvar_pa env) tv }
+
 -- Look up the dfun of a class instance.
 --
 -- The match must be unique - ie, match exactly one instance - but the 
@@ -259,9 +307,7 @@ lookupInst cls tys
              where
                inst_tys'  = [ty | Right ty <- inst_tys]
                noFlexiVar = all isRight inst_tys
-          _other                  -> 
-             pprPanic "VectMonad.lookupInst: not found: " 
-                      (ppr $ mkTyConApp (classTyCon cls) tys)
+          _other         -> noV
        }
   where
     isRight (Left  _) = False