compiler/vectorise/VectMonad.hs

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