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