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(..),
  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, 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                 , liftingContext   :: Var
  77                 }
  78
  79 initBuiltins :: DsM Builtins
  80 initBuiltins
  81   = do
  82       parrayTyCon  <- dsLookupTyCon parrayTyConName
  83       paTyCon      <- dsLookupTyCon paTyConName
  84       let paDataCon = case tyConDataCons paTyCon of [dc] -> dc
  85       closureTyCon <- dsLookupTyCon closureTyConName
  86
  87       mkClosureVar     <- dsLookupGlobalId mkClosureName
  88       applyClosureVar  <- dsLookupGlobalId applyClosureName
  89       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  90       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  91       lengthPAVar      <- dsLookupGlobalId lengthPAName
  92       replicatePAVar   <- dsLookupGlobalId replicatePAName
  93       emptyPAVar       <- dsLookupGlobalId emptyPAName
  94
  95       liftingContext <- liftM (\u -> mkSysLocal FSLIT("lc") u intPrimTy)
  96                               newUnique
  97
  98       return $ Builtins {
  99                  parrayTyCon      = parrayTyCon
 100                , paTyCon          = paTyCon
 101                , paDataCon        = paDataCon
 102                , closureTyCon     = closureTyCon
 103                , mkClosureVar     = mkClosureVar
 104                , applyClosureVar  = applyClosureVar
 105                , mkClosurePVar    = mkClosurePVar
 106                , applyClosurePVar = applyClosurePVar
 107                , lengthPAVar      = lengthPAVar
 108                , replicatePAVar   = replicatePAVar
 109                , emptyPAVar       = emptyPAVar
 110                , liftingContext   = liftingContext
 111                }
 112
 113 data GlobalEnv = GlobalEnv {
 114                   -- Mapping from global variables to their vectorised versions.
 115                   --
 116                   global_vars :: VarEnv Var
 117
 118                   -- Exported variables which have a vectorised version
 119                   --
 120                 , global_exported_vars :: VarEnv (Var, Var)
 121
 122                   -- Mapping from TyCons to their vectorised versions.
 123                   -- TyCons which do not have to be vectorised are mapped to
 124                   -- themselves.
 125                   --
 126                 , global_tycons :: NameEnv TyCon
 127
 128                   -- Mapping from DataCons to their vectorised versions
 129                   --
 130                 , global_datacons :: NameEnv DataCon
 131
 132                   -- Mapping from TyCons to their PA dfuns
 133                   --
 134                 , global_pa_funs :: NameEnv Var
 135
 136                 -- External package inst-env & home-package inst-env for class
 137                 -- instances
 138                 --
 139                 , global_inst_env :: (InstEnv, InstEnv)
 140
 141                 -- External package inst-env & home-package inst-env for family
 142                 -- instances
 143                 --
 144                 , global_fam_inst_env :: FamInstEnvs
 145
 146                 -- Hoisted bindings
 147                 , global_bindings :: [(Var, CoreExpr)]
 148
 149                   -- Global Rdr environment (from ModGuts)
 150                   --
 151                 , global_rdr_env :: GlobalRdrEnv
 152                 }
 153
 154 data LocalEnv = LocalEnv {
 155                  -- Mapping from local variables to their vectorised and
 156                  -- lifted versions
 157                  --
 158                  local_vars :: VarEnv (Var, Var)
 159
 160                  -- In-scope type variables
 161                  --
 162                , local_tyvars :: [TyVar]
 163
 164                  -- Mapping from tyvars to their PA dictionaries
 165                , local_tyvar_pa :: VarEnv CoreExpr
 166
 167                  -- Local binding name
 168                , local_bind_name :: FastString
 169                }
 170
 171
 172 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> Builtins -> GlobalRdrEnv
 173               -> GlobalEnv
 174 initGlobalEnv info instEnvs famInstEnvs bi rdr_env
 175   = GlobalEnv {
 176       global_vars          = mapVarEnv snd $ vectInfoVar info
 177     , global_exported_vars = emptyVarEnv
 178     , global_tycons        = extendNameEnv (mapNameEnv snd (vectInfoTyCon info))
 179                                            (tyConName funTyCon) (closureTyCon bi)
 180
 181     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 182     , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
 183     , global_inst_env      = instEnvs
 184     , global_fam_inst_env  = famInstEnvs
 185     , global_bindings      = []
 186     , global_rdr_env       = rdr_env
 187     }
 188
 189 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 190 setFamInstEnv l_fam_inst genv
 191   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 192   where
 193     (g_fam_inst, _) = global_fam_inst_env genv
 194
 195 emptyLocalEnv = LocalEnv {
 196                    local_vars     = emptyVarEnv
 197                  , local_tyvars   = []
 198                  , local_tyvar_pa = emptyVarEnv
 199                  , local_bind_name  = FSLIT("fn")
 200                  }
 201
 202 -- FIXME
 203 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 204 updVectInfo env tyenv info
 205   = info {
 206       vectInfoVar     = global_exported_vars env
 207     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
 208     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
 209     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
 210     }
 211   where
 212     mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
 213                                    | from <- from_tyenv tyenv
 214                                    , let name = getName from
 215                                    , Just to <- [lookupNameEnv (from_env env) name]]
 216
 217 data VResult a = Yes GlobalEnv LocalEnv a | No
 218
 219 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 220
 221 instance Monad VM where
 222   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 223   VM p >>= f = VM $ \bi genv lenv -> do
 224                                       r <- p bi genv lenv
 225                                       case r of
 226                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 227                                         No                -> return No
 228
 229 noV :: VM a
 230 noV = VM $ \_ _ _ -> return No
 231
 232 traceNoV :: String -> SDoc -> VM a
 233 traceNoV s d = pprTrace s d noV
 234
 235 tryV :: VM a -> VM (Maybe a)
 236 tryV (VM p) = VM $ \bi genv lenv ->
 237   do
 238     r <- p bi genv lenv
 239     case r of
 240       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 241       No                -> return (Yes genv  lenv  Nothing)
 242
 243 maybeV :: VM (Maybe a) -> VM a
 244 maybeV p = maybe noV return =<< p
 245
 246 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 247 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 248
 249 orElseV :: VM a -> VM a -> VM a
 250 orElseV p q = maybe q return =<< tryV p
 251
 252 fixV :: (a -> VM a) -> VM a
 253 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 254   where
 255     unYes (Yes _ _ x) = x
 256
 257 localV :: VM a -> VM a
 258 localV p = do
 259              env <- readLEnv id
 260              x <- p
 261              setLEnv env
 262              return x
 263
 264 closedV :: VM a -> VM a
 265 closedV p = do
 266               env <- readLEnv id
 267               setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
 268               x <- p
 269               setLEnv env
 270               return x
 271
 272 liftDs :: DsM a -> VM a
 273 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 274
 275 builtin :: (Builtins -> a) -> VM a
 276 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 277
 278 readGEnv :: (GlobalEnv -> a) -> VM a
 279 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 280
 281 setGEnv :: GlobalEnv -> VM ()
 282 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 283
 284 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 285 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 286
 287 readLEnv :: (LocalEnv -> a) -> VM a
 288 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 289
 290 setLEnv :: LocalEnv -> VM ()
 291 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 292
 293 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 294 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 295
 296 getInstEnv :: VM (InstEnv, InstEnv)
 297 getInstEnv = readGEnv global_inst_env
 298
 299 getFamInstEnv :: VM FamInstEnvs
 300 getFamInstEnv = readGEnv global_fam_inst_env
 301
 302 getBindName :: VM FastString
 303 getBindName = readLEnv local_bind_name
 304
 305 inBind :: Id -> VM a -> VM a
 306 inBind id p
 307   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
 308        p
 309
 310 lookupRdrName :: RdrName -> VM Name
 311 lookupRdrName rdr_name
 312   = do
 313       rdr_env <- readGEnv global_rdr_env
 314       case lookupGRE_RdrName rdr_name rdr_env of
 315         [gre] -> return (gre_name gre)
 316         []    -> pprPanic "VectMonad.lookupRdrName: not found" (ppr rdr_name)
 317         _     -> pprPanic "VectMonad.lookupRdrName: ambiguous" (ppr rdr_name)
 318
 319 lookupRdrVar :: RdrName -> VM Var
 320 lookupRdrVar rdr_name
 321   = do
 322       name <- lookupRdrName rdr_name
 323       liftDs (dsLookupGlobalId name)
 324
 325 cloneName :: (OccName -> OccName) -> Name -> VM Name
 326 cloneName mk_occ name = liftM make (liftDs newUnique)
 327   where
 328     occ_name = mk_occ (nameOccName name)
 329
 330     make u | isExternalName name = mkExternalName u (nameModule name)
 331                                                     occ_name
 332                                                     (nameSrcSpan name)
 333            | otherwise           = mkSystemName u occ_name
 334
 335 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
 336 cloneId mk_occ id ty
 337   = do
 338       name <- cloneName mk_occ (getName id)
 339       let id' | isExportedId id = Id.mkExportedLocalId name ty
 340               | otherwise       = Id.mkLocalId         name ty
 341       return id'
 342
 343 newExportedVar :: OccName -> Type -> VM Var
 344 newExportedVar occ_name ty
 345   = do
 346       mod <- liftDs getModuleDs
 347       u   <- liftDs newUnique
 348
 349       let name = mkExternalName u mod occ_name noSrcSpan
 350
 351       return $ Id.mkExportedLocalId name ty
 352
 353 newLocalVar :: FastString -> Type -> VM Var
 354 newLocalVar fs ty
 355   = do
 356       u <- liftDs newUnique
 357       return $ mkSysLocal fs u ty
 358
 359 newDummyVar :: Type -> VM Var
 360 newDummyVar = newLocalVar FSLIT("ds")
 361
 362 newTyVar :: FastString -> Kind -> VM Var
 363 newTyVar fs k
 364   = do
 365       u <- liftDs newUnique
 366       return $ mkTyVar (mkSysTvName u fs) k
 367
 368 defGlobalVar :: Var -> Var -> VM ()
 369 defGlobalVar v v' = updGEnv $ \env ->
 370   env { global_vars = extendVarEnv (global_vars env) v v'
 371       , global_exported_vars = upd (global_exported_vars env)
 372       }
 373   where
 374     upd env | isExportedId v = extendVarEnv env v (v, v')
 375             | otherwise      = env
 376
 377 lookupVar :: Var -> VM (Scope Var (Var, Var))
 378 lookupVar v
 379   = do
 380       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 381       case r of
 382         Just e  -> return (Local e)
 383         Nothing -> liftM Global
 384                  $  traceMaybeV "lookupVar" (ppr v)
 385                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 386
 387 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 388 lookupTyCon tc
 389   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 390
 391   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 392
 393 defTyCon :: TyCon -> TyCon -> VM ()
 394 defTyCon tc tc' = updGEnv $ \env ->
 395   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 396
 397 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 398 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 399
 400 defDataCon :: DataCon -> DataCon -> VM ()
 401 defDataCon dc dc' = updGEnv $ \env ->
 402   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 403
 404 lookupTyConPA :: TyCon -> VM (Maybe Var)
 405 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 406
 407 defTyConPA :: TyCon -> Var -> VM ()
 408 defTyConPA tc pa = updGEnv $ \env ->
 409   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 410
 411 defTyConRdrPAs :: [(Name, RdrName)] -> VM ()
 412 defTyConRdrPAs ps
 413   = do
 414       pas <- mapM lookupRdrVar rdr_names
 415       updGEnv $ \env ->
 416         env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 417                                                  (zip tcs pas) }
 418   where
 419     (tcs, rdr_names) = unzip 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                                              (mg_rdr_env guts))
 516                    emptyLocalEnv
 517         case r of
 518           Yes genv _ x -> return $ Just (new_info genv, x)
 519           No           -> return Nothing
 520
 521     new_info genv = updVectInfo genv (mg_types guts) info
 522