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                 , closureTyCon     :: TyCon
  70                 , mkClosureVar     :: Var
  71                 , applyClosureVar  :: Var
  72                 , mkClosurePVar    :: Var
  73                 , applyClosurePVar :: Var
  74                 , lengthPAVar      :: Var
  75                 , replicatePAVar   :: Var
  76                 , emptyPAVar       :: Var
  77                 -- , packPAVar        :: Var
  78                 -- , combinePAVar     :: Var
  79                 , intEqPAVar       :: Var
  80                 , liftingContext   :: Var
  81                 }
  82
  83 initBuiltins :: DsM Builtins
  84 initBuiltins
  85   = do
  86       parrayTyCon  <- dsLookupTyCon parrayTyConName
  87       paTyCon      <- dsLookupTyCon paTyConName
  88       let paDataCon = case tyConDataCons paTyCon of [dc] -> dc
  89       closureTyCon <- dsLookupTyCon closureTyConName
  90
  91       mkClosureVar     <- dsLookupGlobalId mkClosureName
  92       applyClosureVar  <- dsLookupGlobalId applyClosureName
  93       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  94       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  95       lengthPAVar      <- dsLookupGlobalId lengthPAName
  96       replicatePAVar   <- dsLookupGlobalId replicatePAName
  97       emptyPAVar       <- dsLookupGlobalId emptyPAName
  98       -- packPAVar        <- dsLookupGlobalId packPAName
  99       -- combinePAVar     <- dsLookupGlobalId combinePAName
 100       intEqPAVar       <- dsLookupGlobalId intEqPAName
 101
 102       liftingContext <- liftM (\u -> mkSysLocal FSLIT("lc") u intPrimTy)
 103                               newUnique
 104
 105       return $ Builtins {
 106                  parrayTyCon      = parrayTyCon
 107                , paTyCon          = paTyCon
 108                , paDataCon        = paDataCon
 109                , closureTyCon     = closureTyCon
 110                , mkClosureVar     = mkClosureVar
 111                , applyClosureVar  = applyClosureVar
 112                , mkClosurePVar    = mkClosurePVar
 113                , applyClosurePVar = applyClosurePVar
 114                , lengthPAVar      = lengthPAVar
 115                , replicatePAVar   = replicatePAVar
 116                , emptyPAVar       = emptyPAVar
 117                -- , packPAVar        = packPAVar
 118                -- , combinePAVar     = combinePAVar
 119                , intEqPAVar       = intEqPAVar
 120                , liftingContext   = liftingContext
 121                }
 122
 123 data GlobalEnv = GlobalEnv {
 124                   -- Mapping from global variables to their vectorised versions.
 125                   --
 126                   global_vars :: VarEnv Var
 127
 128                   -- Exported variables which have a vectorised version
 129                   --
 130                 , global_exported_vars :: VarEnv (Var, Var)
 131
 132                   -- Mapping from TyCons to their vectorised versions.
 133                   -- TyCons which do not have to be vectorised are mapped to
 134                   -- themselves.
 135                   --
 136                 , global_tycons :: NameEnv TyCon
 137
 138                   -- Mapping from DataCons to their vectorised versions
 139                   --
 140                 , global_datacons :: NameEnv DataCon
 141
 142                   -- Mapping from TyCons to their PA dfuns
 143                   --
 144                 , global_pa_funs :: NameEnv Var
 145
 146                 -- External package inst-env & home-package inst-env for class
 147                 -- instances
 148                 --
 149                 , global_inst_env :: (InstEnv, InstEnv)
 150
 151                 -- External package inst-env & home-package inst-env for family
 152                 -- instances
 153                 --
 154                 , global_fam_inst_env :: FamInstEnvs
 155
 156                 -- Hoisted bindings
 157                 , global_bindings :: [(Var, CoreExpr)]
 158                 }
 159
 160 data LocalEnv = LocalEnv {
 161                  -- Mapping from local variables to their vectorised and
 162                  -- lifted versions
 163                  --
 164                  local_vars :: VarEnv (Var, Var)
 165
 166                  -- In-scope type variables
 167                  --
 168                , local_tyvars :: [TyVar]
 169
 170                  -- Mapping from tyvars to their PA dictionaries
 171                , local_tyvar_pa :: VarEnv CoreExpr
 172
 173                  -- Local binding name
 174                , local_bind_name :: FastString
 175                }
 176
 177
 178 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> Builtins
 179               -> GlobalEnv
 180 initGlobalEnv info instEnvs famInstEnvs bi
 181   = GlobalEnv {
 182       global_vars          = mapVarEnv snd $ vectInfoVar info
 183     , global_exported_vars = emptyVarEnv
 184     , global_tycons        = extendNameEnv (mapNameEnv snd (vectInfoTyCon info))
 185                                            (tyConName funTyCon) (closureTyCon bi)
 186
 187     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 188     , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
 189     , global_inst_env      = instEnvs
 190     , global_fam_inst_env  = famInstEnvs
 191     , global_bindings      = []
 192     }
 193
 194 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 195 setFamInstEnv l_fam_inst genv
 196   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 197   where
 198     (g_fam_inst, _) = global_fam_inst_env genv
 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 lookupExternalVar :: Module -> FastString -> VM Var
 316 lookupExternalVar mod fs
 317   = liftDs
 318   $ dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
 319
 320 cloneName :: (OccName -> OccName) -> Name -> VM Name
 321 cloneName mk_occ name = liftM make (liftDs newUnique)
 322   where
 323     occ_name = mk_occ (nameOccName name)
 324
 325     make u | isExternalName name = mkExternalName u (nameModule name)
 326                                                     occ_name
 327                                                     (nameSrcSpan name)
 328            | otherwise           = mkSystemName u occ_name
 329
 330 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
 331 cloneId mk_occ id ty
 332   = do
 333       name <- cloneName mk_occ (getName id)
 334       let id' | isExportedId id = Id.mkExportedLocalId name ty
 335               | otherwise       = Id.mkLocalId         name ty
 336       return id'
 337
 338 newExportedVar :: OccName -> Type -> VM Var
 339 newExportedVar occ_name ty
 340   = do
 341       mod <- liftDs getModuleDs
 342       u   <- liftDs newUnique
 343
 344       let name = mkExternalName u mod occ_name noSrcSpan
 345
 346       return $ Id.mkExportedLocalId name ty
 347
 348 newLocalVar :: FastString -> Type -> VM Var
 349 newLocalVar fs ty
 350   = do
 351       u <- liftDs newUnique
 352       return $ mkSysLocal fs u ty
 353
 354 newDummyVar :: Type -> VM Var
 355 newDummyVar = newLocalVar FSLIT("ds")
 356
 357 newTyVar :: FastString -> Kind -> VM Var
 358 newTyVar fs k
 359   = do
 360       u <- liftDs newUnique
 361       return $ mkTyVar (mkSysTvName u fs) k
 362
 363 defGlobalVar :: Var -> Var -> VM ()
 364 defGlobalVar v v' = updGEnv $ \env ->
 365   env { global_vars = extendVarEnv (global_vars env) v v'
 366       , global_exported_vars = upd (global_exported_vars env)
 367       }
 368   where
 369     upd env | isExportedId v = extendVarEnv env v (v, v')
 370             | otherwise      = env
 371
 372 lookupVar :: Var -> VM (Scope Var (Var, Var))
 373 lookupVar v
 374   = do
 375       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 376       case r of
 377         Just e  -> return (Local e)
 378         Nothing -> liftM Global
 379                  $  traceMaybeV "lookupVar" (ppr v)
 380                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 381
 382 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 383 lookupTyCon tc
 384   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 385
 386   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 387
 388 defTyCon :: TyCon -> TyCon -> VM ()
 389 defTyCon tc tc' = updGEnv $ \env ->
 390   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 391
 392 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 393 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 394
 395 defDataCon :: DataCon -> DataCon -> VM ()
 396 defDataCon dc dc' = updGEnv $ \env ->
 397   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 398
 399 lookupTyConPA :: TyCon -> VM (Maybe Var)
 400 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 401
 402 defTyConPA :: TyCon -> Var -> VM ()
 403 defTyConPA tc pa = updGEnv $ \env ->
 404   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 405
 406 defTyConPAs :: [(TyCon, Var)] -> VM ()
 407 defTyConPAs ps = updGEnv $ \env ->
 408   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 409                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 410
 411 defTyConBuiltinPAs :: [(Name, Module, FastString)] -> VM ()
 412 defTyConBuiltinPAs ps
 413   = do
 414       pas <- zipWithM lookupExternalVar mods fss
 415       updGEnv $ \env ->
 416         env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 417                                                  (zip tcs pas) }
 418   where
 419     (tcs, mods, fss) = unzip3 ps
 420
 421 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 422 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 423
 424 defLocalTyVar :: TyVar -> VM ()
 425 defLocalTyVar tv = updLEnv $ \env ->
 426   env { local_tyvars   = tv : local_tyvars env
 427       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 428       }
 429
 430 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 431 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 432   env { local_tyvars   = tv : local_tyvars env
 433       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 434       }
 435
 436 localTyVars :: VM [TyVar]
 437 localTyVars = readLEnv (reverse . local_tyvars)
 438
 439 -- Look up the dfun of a class instance.
 440 --
 441 -- The match must be unique - ie, match exactly one instance - but the
 442 -- type arguments used for matching may be more specific than those of
 443 -- the class instance declaration.  The found class instances must not have
 444 -- any type variables in the instance context that do not appear in the
 445 -- instances head (i.e., no flexi vars); for details for what this means,
 446 -- see the docs at InstEnv.lookupInstEnv.
 447 --
 448 {-
 449 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 450 lookupInst cls tys
 451   = do { instEnv <- getInstEnv
 452        ; case lookupInstEnv instEnv cls tys of
 453            ([(inst, inst_tys)], _)
 454              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 455              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 456                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 457              where
 458                inst_tys'  = [ty | Right ty <- inst_tys]
 459                noFlexiVar = all isRight inst_tys
 460            _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
 461        }
 462   where
 463     isRight (Left  _) = False
 464     isRight (Right _) = True
 465 -}
 466
 467 -- Look up the representation tycon of a family 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 family instance declaration.
 472 --
 473 -- Return the instance tycon and its type instance.  For example, if we have
 474 --
 475 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 476 --
 477 -- then we have a coercion (ie, type instance of family instance coercion)
 478 --
 479 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 480 --
 481 -- which implies that :R42T was declared as 'data instance T [a]'.
 482 --
 483 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 484 lookupFamInst tycon tys
 485   = ASSERT( isOpenTyCon tycon )
 486     do { instEnv <- getFamInstEnv
 487        ; case lookupFamInstEnv instEnv tycon tys of
 488            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 489            _other                ->
 490              pprPanic "VectMonad.lookupFamInst: not found: "
 491                       (ppr $ mkTyConApp tycon tys)
 492        }
 493
 494 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 495 initV hsc_env guts info p
 496   = do
 497       eps <- hscEPS hsc_env
 498       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 499       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 500
 501       Just r <- initDs hsc_env (mg_module guts)
 502                                (mg_rdr_env guts)
 503                                (mg_types guts)
 504                                (go instEnvs famInstEnvs)
 505       return r
 506   where
 507
 508     go instEnvs famInstEnvs =
 509       do
 510         builtins <- initBuiltins
 511         r <- runVM p builtins (initGlobalEnv info
 512                                              instEnvs
 513                                              famInstEnvs
 514                                              builtins)
 515                    emptyLocalEnv
 516         case r of
 517           Yes genv _ x -> return $ Just (new_info genv, x)
 518           No           -> return Nothing
 519
 520     new_info genv = updVectInfo genv (mg_types guts) info
 521