compiler/vectorise/VectMonad.hs

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