compiler/vectorise/VectMonad.hs

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