Finish breaking up VectBuiltIn and VectMonad, and add comments

[ghc-hetmet.git] / compiler / vectorise / Vectorise / Monad.hs

diff --git a/compiler/vectorise/Vectorise/Monad.hs b/compiler/vectorise/Vectorise/Monad.hs
new file mode 100644 (file)
index 0000000..42c1435
--- /dev/null
+++ b/compiler/vectorise/Vectorise/Monad.hs
@@ -0,0 +1,146 @@
+
+module Vectorise.Monad (
+       module Vectorise.Monad.Base,
+       module Vectorise.Monad.Naming,
+       module Vectorise.Monad.Local,
+       module Vectorise.Monad.Global,
+       module Vectorise.Monad.InstEnv,
+       initV,
+
+       -- * Builtins
+       liftBuiltinDs,
+       builtin,
+       builtins,
+       
+       -- * Variables
+       lookupVar,
+       maybeCantVectoriseVarM,
+       dumpVar,
+
+       -- * Primitives
+       lookupPrimPArray,
+       lookupPrimMethod
+)
+where
+import Vectorise.Monad.Base
+import Vectorise.Monad.Naming
+import Vectorise.Monad.Local
+import Vectorise.Monad.Global
+import Vectorise.Monad.InstEnv
+import Vectorise.Builtins
+import Vectorise.Env
+
+import HscTypes hiding  ( MonadThings(..) )
+import Module
+import TyCon
+import Var
+import VarEnv
+import Id
+import DsMonad
+import Outputable
+import Control.Monad
+
+
+-- | Run a vectorisation computation.
+initV  :: PackageId
+       -> HscEnv
+       -> ModGuts
+       -> VectInfo
+       -> VM a
+       -> IO (Maybe (VectInfo, a))
+
+initV pkg hsc_env guts info p
+  = do
+         -- XXX: ignores error messages and warnings, check that this is
+         -- indeed ok (the use of "Just r" suggests so)
+      (_,Just r) <- initDs hsc_env (mg_module guts)
+                               (mg_rdr_env guts)
+                               (mg_types guts)
+                               go
+      return r
+  where
+    go 
+     = do
+        builtins       <- initBuiltins pkg
+        builtin_vars   <- initBuiltinVars builtins
+        builtin_tycons <- initBuiltinTyCons builtins
+        let builtin_datacons = initBuiltinDataCons builtins
+        builtin_boxed  <- initBuiltinBoxedTyCons builtins
+        builtin_scalars        <- initBuiltinScalars builtins
+
+        eps <- liftIO $ hscEPS hsc_env
+        let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
+            instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
+
+        builtin_prs    <- initBuiltinPRs builtins instEnvs
+        builtin_pas    <- initBuiltinPAs builtins instEnvs
+
+        let genv = extendImportedVarsEnv builtin_vars
+                 . extendScalars builtin_scalars
+                 . extendTyConsEnv builtin_tycons
+                 . extendDataConsEnv builtin_datacons
+                 . extendPAFunsEnv builtin_pas
+                 . setPRFunsEnv    builtin_prs
+                 . setBoxedTyConsEnv builtin_boxed
+                 $ initGlobalEnv info instEnvs famInstEnvs
+
+        r <- runVM p builtins genv 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
+
+
+-- Builtins -------------------------------------------------------------------
+-- | Lift a desugaring computation using the `Builtins` into the vectorisation monad.
+liftBuiltinDs :: (Builtins -> DsM a) -> VM a
+liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
+
+
+-- | Project something from the set of builtins.
+builtin :: (Builtins -> a) -> VM a
+builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
+
+
+-- | Lift a function using the `Builtins` into the vectorisation monad.
+builtins :: (a -> Builtins -> b) -> VM (a -> b)
+builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
+
+
+-- Var ------------------------------------------------------------------------
+-- | Lookup the vectorised and\/or lifted versions of this variable.
+--     If it's in the global environment we get the vectorised version.
+--      If it's in the local environment we get both the vectorised and lifted version.
+lookupVar :: Var -> VM (Scope Var (Var, Var))
+lookupVar v
+ = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
+      case r of
+        Just e  -> return (Local e)
+        Nothing -> liftM Global
+                . maybeCantVectoriseVarM v
+                . readGEnv $ \env -> lookupVarEnv (global_vars env) v
+
+maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
+maybeCantVectoriseVarM v p
+ = do r <- p
+      case r of
+        Just x  -> return x
+        Nothing -> dumpVar v
+
+dumpVar :: Var -> a
+dumpVar var
+       | Just _                <- isClassOpId_maybe var
+       = cantVectorise "ClassOpId not vectorised:" (ppr var)
+
+       | otherwise
+       = cantVectorise "Variable not vectorised:" (ppr var)
+
+
+-- Primitives -----------------------------------------------------------------
+lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
+lookupPrimPArray = liftBuiltinDs . primPArray
+
+lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
+lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
+