compiler/vectorise/VectMonad.hs

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