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