compiler/vectorise/Vectorise/Monad.hs

   1
   2 module Vectorise.Monad (
   3         module Vectorise.Monad.Base,
   4         module Vectorise.Monad.Naming,
   5         module Vectorise.Monad.Local,
   6         module Vectorise.Monad.Global,
   7         module Vectorise.Monad.InstEnv,
   8         initV,
   9
  10         -- * Builtins
  11         liftBuiltinDs,
  12         builtin,
  13         builtins,
  14
  15         -- * Variables
  16         lookupVar,
  17         maybeCantVectoriseVarM,
  18         dumpVar,
  19
  20         -- * Primitives
  21         lookupPrimPArray,
  22         lookupPrimMethod
  23 )
  24 where
  25 import Vectorise.Monad.Base
  26 import Vectorise.Monad.Naming
  27 import Vectorise.Monad.Local
  28 import Vectorise.Monad.Global
  29 import Vectorise.Monad.InstEnv
  30 import Vectorise.Builtins
  31 import Vectorise.Env
  32
  33 import HscTypes hiding  ( MonadThings(..) )
  34 import Module
  35 import TyCon
  36 import Var
  37 import VarEnv
  38 import Id
  39 import DsMonad
  40 import Outputable
  41 import Control.Monad
  42
  43
  44 -- | Run a vectorisation computation.
  45 initV   :: PackageId
  46         -> HscEnv
  47         -> ModGuts
  48         -> VectInfo
  49         -> VM a
  50         -> IO (Maybe (VectInfo, a))
  51
  52 initV pkg hsc_env guts info p
  53   = do
  54          -- XXX: ignores error messages and warnings, check that this is
  55          -- indeed ok (the use of "Just r" suggests so)
  56       (_,Just r) <- initDs hsc_env (mg_module guts)
  57                                (mg_rdr_env guts)
  58                                (mg_types guts)
  59                                go
  60       return r
  61   where
  62     go
  63      = do
  64         builtins        <- initBuiltins pkg
  65         builtin_vars    <- initBuiltinVars builtins
  66         builtin_tycons  <- initBuiltinTyCons builtins
  67         let builtin_datacons = initBuiltinDataCons builtins
  68         builtin_boxed   <- initBuiltinBoxedTyCons builtins
  69         builtin_scalars <- initBuiltinScalars builtins
  70
  71         eps <- liftIO $ hscEPS hsc_env
  72         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
  73             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
  74
  75         builtin_prs     <- initBuiltinPRs builtins instEnvs
  76         builtin_pas     <- initBuiltinPAs builtins instEnvs
  77
  78         let genv = extendImportedVarsEnv builtin_vars
  79                  . extendScalars builtin_scalars
  80                  . extendTyConsEnv builtin_tycons
  81                  . extendDataConsEnv builtin_datacons
  82                  . extendPAFunsEnv builtin_pas
  83                  . setPRFunsEnv    builtin_prs
  84                  . setBoxedTyConsEnv builtin_boxed
  85                  $ initGlobalEnv info instEnvs famInstEnvs
  86
  87         r <- runVM p builtins genv emptyLocalEnv
  88         case r of
  89           Yes genv _ x -> return $ Just (new_info genv, x)
  90           No           -> return Nothing
  91
  92     new_info genv = updVectInfo genv (mg_types guts) info
  93
  94
  95 -- Builtins -------------------------------------------------------------------
  96 -- | Lift a desugaring computation using the `Builtins` into the vectorisation monad.
  97 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
  98 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
  99
 100
 101 -- | Project something from the set of builtins.
 102 builtin :: (Builtins -> a) -> VM a
 103 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 104
 105
 106 -- | Lift a function using the `Builtins` into the vectorisation monad.
 107 builtins :: (a -> Builtins -> b) -> VM (a -> b)
 108 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 109
 110
 111 -- Var ------------------------------------------------------------------------
 112 -- | Lookup the vectorised and\/or lifted versions of this variable.
 113 --      If it's in the global environment we get the vectorised version.
 114 --      If it's in the local environment we get both the vectorised and lifted version.
 115 lookupVar :: Var -> VM (Scope Var (Var, Var))
 116 lookupVar v
 117  = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 118       case r of
 119         Just e  -> return (Local e)
 120         Nothing -> liftM Global
 121                 . maybeCantVectoriseVarM v
 122                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
 123
 124 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
 125 maybeCantVectoriseVarM v p
 126  = do r <- p
 127       case r of
 128         Just x  -> return x
 129         Nothing -> dumpVar v
 130
 131 dumpVar :: Var -> a
 132 dumpVar var
 133         | Just _                <- isClassOpId_maybe var
 134         = cantVectorise "ClassOpId not vectorised:" (ppr var)
 135
 136         | otherwise
 137         = cantVectorise "Variable not vectorised:" (ppr var)
 138
 139
 140 -- Primitives -----------------------------------------------------------------
 141 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 142 lookupPrimPArray = liftBuiltinDs . primPArray
 143
 144 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 145 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
 146