Put vectorisation monad into a separate file

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
new file mode 100644 (file)
index 0000000..46204b0
--- /dev/null
+++ b/compiler/vectorise/VectMonad.hs
@@ -0,0 +1,233 @@
+module VectMonad (
+  VM,
+
+  noV, tryV, maybeV, orElseV, localV, initV,
+  newLocalVar, newTyVar,
+  
+  Builtins(..),
+  builtin,
+
+  GlobalEnv(..),
+  readGEnv, setGEnv, updGEnv,
+
+  LocalEnv(..),
+  readLEnv, setLEnv, updLEnv,
+
+  lookupTyCon, extendTyVarPA
+) where
+
+#include "HsVersions.h"
+
+import HscTypes
+import CoreSyn
+import TyCon
+import Type
+import Var
+import VarEnv
+import Id
+import Name
+import NameEnv
+
+import DsMonad
+import PrelNames
+
+import FastString
+
+-- ----------------------------------------------------------------------------
+-- Vectorisation monad
+
+data Builtins = Builtins {
+                  parrayTyCon      :: TyCon
+                , paTyCon          :: TyCon
+                , closureTyCon     :: TyCon
+                , mkClosureVar     :: Var
+                , applyClosureVar  :: Var
+                , mkClosurePVar    :: Var
+                , applyClosurePVar :: Var
+                , closurePAVar     :: Var
+                , lengthPAVar      :: Var
+                , replicatePAVar   :: Var
+                }
+
+initBuiltins :: DsM Builtins
+initBuiltins
+  = do
+      parrayTyCon  <- dsLookupTyCon parrayTyConName
+      paTyCon      <- dsLookupTyCon paTyConName
+      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
+               , closureTyCon     = closureTyCon
+               , mkClosureVar     = mkClosureVar
+               , applyClosureVar  = applyClosureVar
+               , mkClosurePVar    = mkClosurePVar
+               , applyClosurePVar = applyClosurePVar
+               , closurePAVar     = closurePAVar
+               , lengthPAVar      = lengthPAVar
+               , replicatePAVar   = replicatePAVar
+               }
+
+data GlobalEnv = GlobalEnv {
+                  -- Mapping from global variables to their vectorised versions.
+                  -- 
+                  global_vars :: VarEnv CoreExpr
+
+                  -- Exported variables which have a vectorised version
+                  --
+                , global_exported_vars :: VarEnv (Var, Var)
+
+                  -- Mapping from TyCons to their vectorised versions.
+                  -- TyCons which do not have to be vectorised are mapped to
+                  -- themselves.
+                  --
+                , global_tycons :: NameEnv TyCon
+
+                  -- Mapping from TyCons to their PA dictionaries
+                  --
+                , global_tycon_pa :: NameEnv CoreExpr
+                }
+
+data LocalEnv = LocalEnv {
+                 -- Mapping from local variables to their vectorised and
+                 -- lifted versions
+                 --
+                 local_vars :: VarEnv (CoreExpr, CoreExpr)
+
+                 -- Mapping from tyvars to their PA dictionaries
+               , local_tyvar_pa :: VarEnv CoreExpr
+               }
+              
+
+initGlobalEnv :: VectInfo -> GlobalEnv
+initGlobalEnv info
+  = GlobalEnv {
+      global_vars          = mapVarEnv  (Var . snd) $ vectInfoCCVar   info
+    , global_exported_vars = emptyVarEnv
+    , global_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
+    , global_tycon_pa      = emptyNameEnv
+    }
+
+emptyLocalEnv = LocalEnv {
+                   local_vars     = emptyVarEnv
+                 , local_tyvar_pa = emptyVarEnv
+                 }
+
+-- FIXME
+updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
+updVectInfo env tyenv info
+  = info {
+      vectInfoCCVar   = global_exported_vars env
+    , vectInfoCCTyCon = tc_env
+    }
+  where
+    tc_env = mkNameEnv [(tc_name, (tc,tc'))
+               | tc <- typeEnvTyCons tyenv
+               , let tc_name = tyConName tc
+               , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
+
+data VResult a = Yes GlobalEnv LocalEnv a | No
+
+newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
+
+instance Monad VM where
+  return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
+  VM p >>= f = VM $ \bi genv lenv -> do
+                                      r <- p bi genv lenv
+                                      case r of
+                                        Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
+                                        No                -> return No
+
+noV :: VM a
+noV = VM $ \_ _ _ -> return No
+
+tryV :: VM a -> VM (Maybe a)
+tryV (VM p) = VM $ \bi genv lenv ->
+  do
+    r <- p bi genv lenv
+    case r of
+      Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
+      No                -> return (Yes genv  lenv  Nothing)
+
+maybeV :: VM (Maybe a) -> VM a
+maybeV p = maybe noV return =<< p
+
+orElseV :: VM a -> VM a -> VM a
+orElseV p q = maybe q return =<< tryV p
+
+localV :: VM a -> VM a
+localV p = do
+             env <- readLEnv id
+             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) }
+
+builtin :: (Builtins -> a) -> VM a
+builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
+
+readGEnv :: (GlobalEnv -> a) -> VM a
+readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
+
+setGEnv :: GlobalEnv -> VM ()
+setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
+
+updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
+updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
+
+readLEnv :: (LocalEnv -> a) -> VM a
+readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
+
+setLEnv :: LocalEnv -> VM ()
+setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
+
+updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
+updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
+
+newLocalVar :: FastString -> Type -> VM Var
+newLocalVar fs ty
+  = do
+      u <- liftDs newUnique
+      return $ mkSysLocal fs u ty
+
+newTyVar :: FastString -> Kind -> VM Var
+newTyVar fs k
+  = do
+      u <- liftDs newUnique
+      return $ mkTyVar (mkSysTvName u fs) k
+
+lookupTyCon :: TyCon -> VM (Maybe TyCon)
+lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
+
+extendTyVarPA :: Var -> CoreExpr -> VM ()
+extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
+
+initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
+initV hsc_env guts info p
+  = do
+      Just r <- initDs hsc_env (mg_module guts)
+                               (mg_rdr_env guts)
+                               (mg_types guts)
+                               go
+      return r
+  where
+    go = do
+           builtins <- initBuiltins
+           r <- runVM p builtins (initGlobalEnv info) emptyLocalEnv
+           case r of
+             Yes genv _ x -> return $ Just (new_info genv, x)
+             No           -> return Nothing
+
+    new_info genv = updVectInfo genv (mg_types guts) info
+

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 0358aca..df9fdb9 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -3,6 +3,8 @@ where
 
 #include "HsVersions.h"
 
+import VectMonad
+
 import DynFlags
 import HscTypes
 
@@ -35,215 +37,17 @@ vectorise hsc_env guts
       showPass dflags "Vectorisation"
       eps <- hscEPS hsc_env
       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
-      Just guts' <- initDs hsc_env (mg_module guts)
-                                   (mg_rdr_env guts)
-                                   (mg_types guts)
-                                   (vectoriseModule info guts)
+      Just (info', guts') <- initV hsc_env guts info (vectModule guts)
       endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
-      return guts'
+      return $ guts' { mg_vect_info = info' }
   where
     dflags = hsc_dflags hsc_env
 
--- ----------------------------------------------------------------------------
--- Vectorisation monad
-
-data Builtins = Builtins {
-                  parrayTyCon      :: TyCon
-                , paTyCon          :: TyCon
-                , closureTyCon     :: TyCon
-                , mkClosureVar     :: Var
-                , applyClosureVar  :: Var
-                , mkClosurePVar    :: Var
-                , applyClosurePVar :: Var
-                , closurePAVar     :: Var
-                , lengthPAVar      :: Var
-                , replicatePAVar   :: Var
-                }
-
-initBuiltins :: DsM Builtins
-initBuiltins
-  = do
-      parrayTyCon  <- dsLookupTyCon parrayTyConName
-      paTyCon      <- dsLookupTyCon paTyConName
-      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
-               , closureTyCon     = closureTyCon
-               , mkClosureVar     = mkClosureVar
-               , applyClosureVar  = applyClosureVar
-               , mkClosurePVar    = mkClosurePVar
-               , applyClosurePVar = applyClosurePVar
-               , closurePAVar     = closurePAVar
-               , lengthPAVar      = lengthPAVar
-               , replicatePAVar   = replicatePAVar
-               }
-
-data GlobalEnv = GlobalEnv {
-                  -- Mapping from global variables to their vectorised versions.
-                  -- 
-                  global_vars :: VarEnv CoreExpr
-
-                  -- Exported variables which have a vectorised version
-                  --
-                , global_exported_vars :: VarEnv (Var, Var)
-
-                  -- Mapping from TyCons to their vectorised versions.
-                  -- TyCons which do not have to be vectorised are mapped to
-                  -- themselves.
-                  --
-                , global_tycons :: NameEnv TyCon
-
-                  -- Mapping from TyCons to their PA dictionaries
-                  --
-                , global_tycon_pa :: NameEnv CoreExpr
-                }
-
-data LocalEnv = LocalEnv {
-                 -- Mapping from local variables to their vectorised and
-                 -- lifted versions
-                 --
-                 local_vars :: VarEnv (CoreExpr, CoreExpr)
-
-                 -- Mapping from tyvars to their PA dictionaries
-               , local_tyvar_pa :: VarEnv CoreExpr
-               }
-              
-
-initGlobalEnv :: VectInfo -> GlobalEnv
-initGlobalEnv info
-  = GlobalEnv {
-      global_vars          = mapVarEnv  (Var . snd) $ vectInfoCCVar   info
-    , global_exported_vars = emptyVarEnv
-    , global_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
-    , global_tycon_pa      = emptyNameEnv
-    }
-
-emptyLocalEnv = LocalEnv {
-                   local_vars     = emptyVarEnv
-                 , local_tyvar_pa = emptyVarEnv
-                 }
-
--- FIXME
-updVectInfo :: GlobalEnv -> ModGuts -> ModGuts
-updVectInfo env guts = guts { mg_vect_info = info' }
-  where
-    info' = info {
-              vectInfoCCVar   = global_exported_vars env
-            , vectInfoCCTyCon = tc_env
-            }
-
-    info  = mg_vect_info guts
-    tyenv = mg_types guts
-
-    tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
-                                            , let tc_name = tyConName tc
-                                            , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
-
-data VResult a = Yes GlobalEnv LocalEnv a | No
-
-newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
-
-instance Monad VM where
-  return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
-  VM p >>= f = VM $ \bi genv lenv -> do
-                                      r <- p bi genv lenv
-                                      case r of
-                                        Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
-                                        No                -> return No
-
-noV :: VM a
-noV = VM $ \_ _ _ -> return No
-
-tryV :: VM a -> VM (Maybe a)
-tryV (VM p) = VM $ \bi genv lenv ->
-  do
-    r <- p bi genv lenv
-    case r of
-      Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
-      No                -> return (Yes genv  lenv  Nothing)
-
-maybeV :: VM (Maybe a) -> VM a
-maybeV p = maybe noV return =<< p
-
-orElseV :: VM a -> VM a -> VM a
-orElseV p q = maybe q return =<< tryV p
-
-localV :: VM a -> VM a
-localV p = do
-             env <- readLEnv id
-             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) }
-
-builtin :: (Builtins -> a) -> VM a
-builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
-
-readGEnv :: (GlobalEnv -> a) -> VM a
-readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
-
-setGEnv :: GlobalEnv -> VM ()
-setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
-
-updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
-updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
-
-readLEnv :: (LocalEnv -> a) -> VM a
-readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
-
-setLEnv :: LocalEnv -> VM ()
-setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
-
-updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
-updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
-
-newLocalVar :: FastString -> Type -> VM Var
-newLocalVar fs ty
-  = do
-      u <- liftDs newUnique
-      return $ mkSysLocal fs u ty
-
-newTyVar :: FastString -> Kind -> VM Var
-newTyVar fs k
-  = do
-      u <- liftDs newUnique
-      return $ mkTyVar (mkSysTvName u fs) k
-
-lookupTyCon :: TyCon -> VM (Maybe TyCon)
-lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
-
-
-extendTyVarPA :: Var -> CoreExpr -> VM ()
-extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
-
--- ----------------------------------------------------------------------------
--- Bindings
-
-vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
-vectoriseModule info guts
-  = do
-      builtins <- initBuiltins
-      r <- runVM (vectModule guts) builtins (initGlobalEnv info) emptyLocalEnv
-      case r of
-        Yes genv _ guts' -> return $ updVectInfo genv guts'
-        No               -> return guts
-
 vectModule :: ModGuts -> VM ModGuts
 vectModule guts = return guts
 
-
+-- ----------------------------------------------------------------------------
+-- Bindings
 
 vectBndr :: Var -> VM (Var, Var)
 vectBndr v