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