compiler/vectorise/VectMonad.hs

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