compiler/vectorise/VectMonad.hs

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