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