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