Collect hoisted vectorised functions

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

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index 56189f6..dc26b4b 100644 (file)
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -1,7 +1,7 @@
 module VectMonad (
   VM,
 
-  noV, tryV, maybeV, orElseV, localV, initV,
+  noV, tryV, maybeV, orElseV, localV, closedV, initV,
   newLocalVar, newTyVar,
   
   Builtins(..), paDictTyCon,
@@ -13,7 +13,8 @@ module VectMonad (
   LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
-  lookupTyCon, extendTyVarPA, deleteTyVarPA,
+  lookupTyCon,
+  lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
 
   lookupInst, lookupFamInst
 ) where
@@ -123,6 +124,9 @@ data LocalEnv = LocalEnv {
 
                  -- Mapping from tyvars to their PA dictionaries
                , local_tyvar_pa :: VarEnv CoreExpr
+
+                 -- Hoisted bindings
+               , local_bindings :: [(Var, CoreExpr)]
                }
               
 
@@ -140,6 +144,7 @@ initGlobalEnv info instEnvs famInstEnvs
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
                  , local_tyvar_pa = emptyVarEnv
+                 , local_bindings = []
                  }
 
 -- FIXME
@@ -191,6 +196,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) }
 
@@ -236,6 +249,9 @@ newTyVar fs k
 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 }
 
@@ -262,9 +278,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