compiler/vectorise/VectMonad.hs

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