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