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