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