compiler/vectorise/VectMonad.hs

   1 module VectMonad (
   2   VM,
   3
   4   noV, tryV, maybeV, orElseV, localV, initV,
   5   newLocalVar, newTyVar,
   6
   7   Builtins(..), paDictTyCon,
   8   builtin,
   9
  10   GlobalEnv(..),
  11   readGEnv, setGEnv, updGEnv,
  12
  13   LocalEnv(..),
  14   readLEnv, setLEnv, updLEnv,
  15
  16   lookupTyCon, extendTyVarPA, deleteTyVarPA,
  17
  18   lookupInst, lookupFamInst
  19 ) where
  20
  21 #include "HsVersions.h"
  22
  23 import HscTypes
  24 import CoreSyn
  25 import Class
  26 import TyCon
  27 import Type
  28 import Var
  29 import VarEnv
  30 import Id
  31 import Name
  32 import NameEnv
  33
  34 import DsMonad
  35 import PrelNames
  36
  37 import InstEnv
  38 import FamInstEnv
  39
  40 import Panic
  41 import Outputable
  42 import FastString
  43
  44 -- ----------------------------------------------------------------------------
  45 -- Vectorisation monad
  46
  47 data Builtins = Builtins {
  48                   parrayTyCon      :: TyCon
  49                 , paClass          :: Class
  50                 , closureTyCon     :: TyCon
  51                 , mkClosureVar     :: Var
  52                 , applyClosureVar  :: Var
  53                 , mkClosurePVar    :: Var
  54                 , applyClosurePVar :: Var
  55                 , lengthPAVar      :: Var
  56                 , replicatePAVar   :: Var
  57                 }
  58
  59 paDictTyCon :: Builtins -> TyCon
  60 paDictTyCon = classTyCon . paClass
  61
  62 initBuiltins :: DsM Builtins
  63 initBuiltins
  64   = do
  65       parrayTyCon  <- dsLookupTyCon parrayTyConName
  66       paClass      <- dsLookupClass paClassName
  67       closureTyCon <- dsLookupTyCon closureTyConName
  68
  69       mkClosureVar     <- dsLookupGlobalId mkClosureName
  70       applyClosureVar  <- dsLookupGlobalId applyClosureName
  71       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  72       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  73       lengthPAVar      <- dsLookupGlobalId lengthPAName
  74       replicatePAVar   <- dsLookupGlobalId replicatePAName
  75
  76       return $ Builtins {
  77                  parrayTyCon      = parrayTyCon
  78                , paClass          = paClass
  79                , closureTyCon     = closureTyCon
  80                , mkClosureVar     = mkClosureVar
  81                , applyClosureVar  = applyClosureVar
  82                , mkClosurePVar    = mkClosurePVar
  83                , applyClosurePVar = applyClosurePVar
  84                , lengthPAVar      = lengthPAVar
  85                , replicatePAVar   = replicatePAVar
  86                }
  87
  88 data GlobalEnv = GlobalEnv {
  89                   -- Mapping from global variables to their vectorised versions.
  90                   --
  91                   global_vars :: VarEnv CoreExpr
  92
  93                   -- Exported variables which have a vectorised version
  94                   --
  95                 , global_exported_vars :: VarEnv (Var, Var)
  96
  97                   -- Mapping from TyCons to their vectorised versions.
  98                   -- TyCons which do not have to be vectorised are mapped to
  99                   -- themselves.
 100                   --
 101                 , global_tycons :: NameEnv TyCon
 102
 103                   -- Mapping from TyCons to their PA dictionaries
 104                   --
 105                 , global_tycon_pa :: NameEnv CoreExpr
 106
 107                 -- External package inst-env & home-package inst-env for class
 108                 -- instances
 109                 --
 110                 , global_inst_env :: (InstEnv, InstEnv)
 111
 112                 -- External package inst-env & home-package inst-env for family
 113                 -- instances
 114                 --
 115                 , global_fam_inst_env :: FamInstEnvs
 116                 }
 117
 118 data LocalEnv = LocalEnv {
 119                  -- Mapping from local variables to their vectorised and
 120                  -- lifted versions
 121                  --
 122                  local_vars :: VarEnv (CoreExpr, CoreExpr)
 123
 124                  -- Mapping from tyvars to their PA dictionaries
 125                , local_tyvar_pa :: VarEnv CoreExpr
 126                }
 127
 128
 129 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 130 initGlobalEnv info instEnvs famInstEnvs
 131   = GlobalEnv {
 132       global_vars          = mapVarEnv  (Var . snd) $ vectInfoCCVar   info
 133     , global_exported_vars = emptyVarEnv
 134     , global_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
 135     , global_tycon_pa      = emptyNameEnv
 136     , global_inst_env      = instEnvs
 137     , global_fam_inst_env  = famInstEnvs
 138     }
 139
 140 emptyLocalEnv = LocalEnv {
 141                    local_vars     = emptyVarEnv
 142                  , local_tyvar_pa = emptyVarEnv
 143                  }
 144
 145 -- FIXME
 146 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 147 updVectInfo env tyenv info
 148   = info {
 149       vectInfoCCVar   = global_exported_vars env
 150     , vectInfoCCTyCon = tc_env
 151     }
 152   where
 153     tc_env = mkNameEnv [(tc_name, (tc,tc'))
 154                | tc <- typeEnvTyCons tyenv
 155                , let tc_name = tyConName tc
 156                , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
 157
 158 data VResult a = Yes GlobalEnv LocalEnv a | No
 159
 160 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 161
 162 instance Monad VM where
 163   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 164   VM p >>= f = VM $ \bi genv lenv -> do
 165                                       r <- p bi genv lenv
 166                                       case r of
 167                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 168                                         No                -> return No
 169
 170 noV :: VM a
 171 noV = VM $ \_ _ _ -> return No
 172
 173 tryV :: VM a -> VM (Maybe a)
 174 tryV (VM p) = VM $ \bi genv lenv ->
 175   do
 176     r <- p bi genv lenv
 177     case r of
 178       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 179       No                -> return (Yes genv  lenv  Nothing)
 180
 181 maybeV :: VM (Maybe a) -> VM a
 182 maybeV p = maybe noV return =<< p
 183
 184 orElseV :: VM a -> VM a -> VM a
 185 orElseV p q = maybe q return =<< tryV p
 186
 187 localV :: VM a -> VM a
 188 localV p = do
 189              env <- readLEnv id
 190              x <- p
 191              setLEnv env
 192              return x
 193
 194 liftDs :: DsM a -> VM a
 195 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 196
 197 builtin :: (Builtins -> a) -> VM a
 198 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 199
 200 readGEnv :: (GlobalEnv -> a) -> VM a
 201 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 202
 203 setGEnv :: GlobalEnv -> VM ()
 204 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 205
 206 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 207 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 208
 209 readLEnv :: (LocalEnv -> a) -> VM a
 210 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 211
 212 setLEnv :: LocalEnv -> VM ()
 213 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 214
 215 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 216 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 217
 218 getInstEnv :: VM (InstEnv, InstEnv)
 219 getInstEnv = readGEnv global_inst_env
 220
 221 getFamInstEnv :: VM FamInstEnvs
 222 getFamInstEnv = readGEnv global_fam_inst_env
 223
 224 newLocalVar :: FastString -> Type -> VM Var
 225 newLocalVar fs ty
 226   = do
 227       u <- liftDs newUnique
 228       return $ mkSysLocal fs u ty
 229
 230 newTyVar :: FastString -> Kind -> VM Var
 231 newTyVar fs k
 232   = do
 233       u <- liftDs newUnique
 234       return $ mkTyVar (mkSysTvName u fs) k
 235
 236 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 237 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 238
 239 extendTyVarPA :: Var -> CoreExpr -> VM ()
 240 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
 241
 242 deleteTyVarPA :: Var -> VM ()
 243 deleteTyVarPA tv = updLEnv $ \env -> env { local_tyvar_pa = delVarEnv (local_tyvar_pa env) tv }
 244
 245 -- Look up the dfun of a class instance.
 246 --
 247 -- The match must be unique - ie, match exactly one instance - but the
 248 -- type arguments used for matching may be more specific than those of
 249 -- the class instance declaration.  The found class instances must not have
 250 -- any type variables in the instance context that do not appear in the
 251 -- instances head (i.e., no flexi vars); for details for what this means,
 252 -- see the docs at InstEnv.lookupInstEnv.
 253 --
 254 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 255 lookupInst cls tys
 256   = do { instEnv <- getInstEnv
 257        ; case lookupInstEnv instEnv cls tys of
 258            ([(inst, inst_tys)], _)
 259              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 260              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 261                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 262              where
 263                inst_tys'  = [ty | Right ty <- inst_tys]
 264                noFlexiVar = all isRight inst_tys
 265            _other                  ->
 266              pprPanic "VectMonad.lookupInst: not found: "
 267                       (ppr $ mkTyConApp (classTyCon cls) tys)
 268        }
 269   where
 270     isRight (Left  _) = False
 271     isRight (Right _) = True
 272
 273 -- Look up the representation tycon of a family instance.
 274 --
 275 -- The match must be unique - ie, match exactly one instance - but the
 276 -- type arguments used for matching may be more specific than those of
 277 -- the family instance declaration.
 278 --
 279 -- Return the instance tycon and its type instance.  For example, if we have
 280 --
 281 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 282 --
 283 -- then we have a coercion (ie, type instance of family instance coercion)
 284 --
 285 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 286 --
 287 -- which implies that :R42T was declared as 'data instance T [a]'.
 288 --
 289 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 290 lookupFamInst tycon tys
 291   = ASSERT( isOpenTyCon tycon )
 292     do { instEnv <- getFamInstEnv
 293        ; case lookupFamInstEnv instEnv tycon tys of
 294            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 295            _other                ->
 296              pprPanic "VectMonad.lookupFamInst: not found: "
 297                       (ppr $ mkTyConApp tycon tys)
 298        }
 299
 300 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 301 initV hsc_env guts info p
 302   = do
 303       eps <- hscEPS hsc_env
 304       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 305       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 306
 307       Just r <- initDs hsc_env (mg_module guts)
 308                                (mg_rdr_env guts)
 309                                (mg_types guts)
 310                                (go instEnvs famInstEnvs)
 311       return r
 312   where
 313
 314     go instEnvs famInstEnvs =
 315       do
 316         builtins <- initBuiltins
 317         r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs)
 318                    emptyLocalEnv
 319         case r of
 320           Yes genv _ x -> return $ Just (new_info genv, x)
 321           No           -> return Nothing
 322
 323     new_info genv = updVectInfo genv (mg_types guts) info
 324