compiler/vectorise/VectMonad.hs

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