compiler/vectorise/VectMonad.hs

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