compiler/vectorise/VectMonad.hs

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