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