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