compiler/vectorise/VectMonad.hs

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