compiler/vectorise/VectMonad.hs

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