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 $ \r -> runVM (f (unYes r)) bi genv lenv )
 232   where
 233     -- NOTE: It is essential that we are lazy in r above so do not replace
 234     --       calls to this function by an explicit case.
 235     unYes (Yes _ _ x) = x
 236     unYes No          = panic "VectMonad.fixV: no result"
 237
 238 localV :: VM a -> VM a
 239 localV p = do
 240              env <- readLEnv id
 241              x <- p
 242              setLEnv env
 243              return x
 244
 245 closedV :: VM a -> VM a
 246 closedV p = do
 247               env <- readLEnv id
 248               setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
 249               x <- p
 250               setLEnv env
 251               return x
 252
 253 liftDs :: DsM a -> VM a
 254 liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
 255
 256 builtin :: (Builtins -> a) -> VM a
 257 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 258
 259 builtins :: (a -> Builtins -> b) -> VM (a -> b)
 260 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 261
 262 readGEnv :: (GlobalEnv -> a) -> VM a
 263 readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
 264
 265 setGEnv :: GlobalEnv -> VM ()
 266 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 267
 268 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 269 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 270
 271 readLEnv :: (LocalEnv -> a) -> VM a
 272 readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
 273
 274 setLEnv :: LocalEnv -> VM ()
 275 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 276
 277 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 278 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 279
 280 {-
 281 getInstEnv :: VM (InstEnv, InstEnv)
 282 getInstEnv = readGEnv global_inst_env
 283 -}
 284
 285 getFamInstEnv :: VM FamInstEnvs
 286 getFamInstEnv = readGEnv global_fam_inst_env
 287
 288 getBindName :: VM FastString
 289 getBindName = readLEnv local_bind_name
 290
 291 inBind :: Id -> VM a -> VM a
 292 inBind id p
 293   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
 294        p
 295
 296 cloneName :: (OccName -> OccName) -> Name -> VM Name
 297 cloneName mk_occ name = liftM make (liftDs newUnique)
 298   where
 299     occ_name = mk_occ (nameOccName name)
 300
 301     make u | isExternalName name = mkExternalName u (nameModule name)
 302                                                     occ_name
 303                                                     (nameSrcSpan name)
 304            | otherwise           = mkSystemName u occ_name
 305
 306 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
 307 cloneId mk_occ id ty
 308   = do
 309       name <- cloneName mk_occ (getName id)
 310       let id' | isExportedId id = Id.mkExportedLocalId name ty
 311               | otherwise       = Id.mkLocalId         name ty
 312       return id'
 313
 314 cloneVar :: Var -> VM Var
 315 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
 316
 317 newExportedVar :: OccName -> Type -> VM Var
 318 newExportedVar occ_name ty
 319   = do
 320       mod <- liftDs getModuleDs
 321       u   <- liftDs newUnique
 322
 323       let name = mkExternalName u mod occ_name noSrcSpan
 324
 325       return $ Id.mkExportedLocalId name ty
 326
 327 newLocalVar :: FastString -> Type -> VM Var
 328 newLocalVar fs ty
 329   = do
 330       u <- liftDs newUnique
 331       return $ mkSysLocal fs u ty
 332
 333 newDummyVar :: Type -> VM Var
 334 newDummyVar = newLocalVar (fsLit "ds")
 335
 336 newTyVar :: FastString -> Kind -> VM Var
 337 newTyVar fs k
 338   = do
 339       u <- liftDs newUnique
 340       return $ mkTyVar (mkSysTvName u fs) k
 341
 342 defGlobalVar :: Var -> Var -> VM ()
 343 defGlobalVar v v' = updGEnv $ \env ->
 344   env { global_vars = extendVarEnv (global_vars env) v v'
 345       , global_exported_vars = upd (global_exported_vars env)
 346       }
 347   where
 348     upd env | isExportedId v = extendVarEnv env v (v, v')
 349             | otherwise      = env
 350
 351 lookupVar :: Var -> VM (Scope Var (Var, Var))
 352 lookupVar v
 353   = do
 354       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 355       case r of
 356         Just e  -> return (Local e)
 357         Nothing -> liftM Global
 358                  $  traceMaybeV "lookupVar" (ppr v)
 359                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 360
 361 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 362 lookupTyCon tc
 363   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 364
 365   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 366
 367 defTyCon :: TyCon -> TyCon -> VM ()
 368 defTyCon tc tc' = updGEnv $ \env ->
 369   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 370
 371 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 372 lookupDataCon dc
 373   | isTupleTyCon (dataConTyCon dc) = return (Just dc)
 374   | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 375
 376 defDataCon :: DataCon -> DataCon -> VM ()
 377 defDataCon dc dc' = updGEnv $ \env ->
 378   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 379
 380 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 381 lookupPrimPArray = liftDs . primPArray
 382
 383 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 384 lookupPrimMethod tycon = liftDs . primMethod tycon
 385
 386 lookupTyConPA :: TyCon -> VM (Maybe Var)
 387 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 388
 389 defTyConPA :: TyCon -> Var -> VM ()
 390 defTyConPA tc pa = updGEnv $ \env ->
 391   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 392
 393 defTyConPAs :: [(TyCon, Var)] -> VM ()
 394 defTyConPAs ps = updGEnv $ \env ->
 395   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 396                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 397
 398 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 399 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 400
 401 lookupTyConPR :: TyCon -> VM (Maybe Var)
 402 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 403
 404 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
 405 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
 406                                                        (tyConName tc)
 407
 408 defLocalTyVar :: TyVar -> VM ()
 409 defLocalTyVar tv = updLEnv $ \env ->
 410   env { local_tyvars   = tv : local_tyvars env
 411       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 412       }
 413
 414 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 415 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 416   env { local_tyvars   = tv : local_tyvars env
 417       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 418       }
 419
 420 localTyVars :: VM [TyVar]
 421 localTyVars = readLEnv (reverse . local_tyvars)
 422
 423 -- Look up the dfun of a class instance.
 424 --
 425 -- The match must be unique - ie, match exactly one instance - but the
 426 -- type arguments used for matching may be more specific than those of
 427 -- the class instance declaration.  The found class instances must not have
 428 -- any type variables in the instance context that do not appear in the
 429 -- instances head (i.e., no flexi vars); for details for what this means,
 430 -- see the docs at InstEnv.lookupInstEnv.
 431 --
 432 {-
 433 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 434 lookupInst cls tys
 435   = do { instEnv <- getInstEnv
 436        ; case lookupInstEnv instEnv cls tys of
 437            ([(inst, inst_tys)], _)
 438              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 439              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 440                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 441              where
 442                inst_tys'  = [ty | Right ty <- inst_tys]
 443                noFlexiVar = all isRight inst_tys
 444            _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
 445        }
 446   where
 447     isRight (Left  _) = False
 448     isRight (Right _) = True
 449 -}
 450
 451 -- Look up the representation tycon of a family instance.
 452 --
 453 -- The match must be unique - ie, match exactly one instance - but the
 454 -- type arguments used for matching may be more specific than those of
 455 -- the family instance declaration.
 456 --
 457 -- Return the instance tycon and its type instance.  For example, if we have
 458 --
 459 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 460 --
 461 -- then we have a coercion (ie, type instance of family instance coercion)
 462 --
 463 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 464 --
 465 -- which implies that :R42T was declared as 'data instance T [a]'.
 466 --
 467 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 468 lookupFamInst tycon tys
 469   = ASSERT( isOpenTyCon tycon )
 470     do { instEnv <- getFamInstEnv
 471        ; case lookupFamInstEnv instEnv tycon tys of
 472            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 473            _other                ->
 474              pprPanic "VectMonad.lookupFamInst: not found: "
 475                       (ppr $ mkTyConApp tycon tys)
 476        }
 477
 478 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 479 initV hsc_env guts info p
 480   = do
 481       Just r <- initDs hsc_env (mg_module guts)
 482                                (mg_rdr_env guts)
 483                                (mg_types guts)
 484                                go
 485       return r
 486   where
 487
 488     go =
 489       do
 490         builtins       <- initBuiltins
 491         builtin_vars   <- initBuiltinVars builtins
 492         builtin_tycons <- initBuiltinTyCons builtins
 493         let builtin_datacons = initBuiltinDataCons builtins
 494         builtin_pas    <- initBuiltinPAs builtins
 495         builtin_prs    <- initBuiltinPRs builtins
 496         builtin_boxed  <- initBuiltinBoxedTyCons builtins
 497
 498         eps <- liftIO $ hscEPS hsc_env
 499         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 500             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 501
 502         let genv = extendImportedVarsEnv builtin_vars
 503                  . extendTyConsEnv builtin_tycons
 504                  . extendDataConsEnv builtin_datacons
 505                  . extendPAFunsEnv builtin_pas
 506                  . setPRFunsEnv    builtin_prs
 507                  . setBoxedTyConsEnv builtin_boxed
 508                  $ initGlobalEnv info instEnvs famInstEnvs
 509
 510         r <- runVM p builtins genv emptyLocalEnv
 511         case r of
 512           Yes genv _ x -> return $ Just (new_info genv, x)
 513           No           -> return Nothing
 514
 515     new_info genv = updVectInfo genv (mg_types guts) info
 516