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