compiler/vectorise/VectMonad.hs

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