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