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                 -- Hoisted bindings
 126                 , global_bindings :: [(Var, CoreExpr)]
 127                 }
 128
 129 data LocalEnv = LocalEnv {
 130                  -- Mapping from local variables to their vectorised and
 131                  -- lifted versions
 132                  --
 133                  local_vars :: VarEnv (CoreExpr, CoreExpr)
 134
 135                  -- Mapping from tyvars to their PA dictionaries
 136                , local_tyvar_pa :: VarEnv 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     , global_bindings      = []
 150     }
 151
 152 emptyLocalEnv = LocalEnv {
 153                    local_vars     = emptyVarEnv
 154                  , local_tyvar_pa = emptyVarEnv
 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 -> Var -> VM ()
 267 defGlobalVar v v' = updGEnv $ \env ->
 268   env { global_vars = extendVarEnv (global_vars env) v (Var v')
 269       , global_exported_vars = upd (global_exported_vars env)
 270       }
 271   where
 272     upd env | isExportedId v = extendVarEnv env v (v, v')
 273             | otherwise      = env
 274
 275 lookupVar :: Var -> VM (Scope CoreExpr (CoreExpr, CoreExpr))
 276 lookupVar v
 277   = do
 278       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 279       case r of
 280         Just e  -> return (Local e)
 281         Nothing -> liftM Global
 282                  $  maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 283
 284 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 285 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 286
 287 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 288 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 289
 290 extendTyVarPA :: Var -> CoreExpr -> VM ()
 291 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
 292
 293 deleteTyVarPA :: Var -> VM ()
 294 deleteTyVarPA tv = updLEnv $ \env -> env { local_tyvar_pa = delVarEnv (local_tyvar_pa env) tv }
 295
 296 -- Look up the dfun of a class instance.
 297 --
 298 -- The match must be unique - ie, match exactly one instance - but the
 299 -- type arguments used for matching may be more specific than those of
 300 -- the class instance declaration.  The found class instances must not have
 301 -- any type variables in the instance context that do not appear in the
 302 -- instances head (i.e., no flexi vars); for details for what this means,
 303 -- see the docs at InstEnv.lookupInstEnv.
 304 --
 305 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 306 lookupInst cls tys
 307   = do { instEnv <- getInstEnv
 308        ; case lookupInstEnv instEnv cls tys of
 309            ([(inst, inst_tys)], _)
 310              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 311              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 312                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 313              where
 314                inst_tys'  = [ty | Right ty <- inst_tys]
 315                noFlexiVar = all isRight inst_tys
 316            _other         -> noV
 317        }
 318   where
 319     isRight (Left  _) = False
 320     isRight (Right _) = True
 321
 322 -- Look up the representation tycon of a family instance.
 323 --
 324 -- The match must be unique - ie, match exactly one instance - but the
 325 -- type arguments used for matching may be more specific than those of
 326 -- the family instance declaration.
 327 --
 328 -- Return the instance tycon and its type instance.  For example, if we have
 329 --
 330 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 331 --
 332 -- then we have a coercion (ie, type instance of family instance coercion)
 333 --
 334 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 335 --
 336 -- which implies that :R42T was declared as 'data instance T [a]'.
 337 --
 338 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 339 lookupFamInst tycon tys
 340   = ASSERT( isOpenTyCon tycon )
 341     do { instEnv <- getFamInstEnv
 342        ; case lookupFamInstEnv instEnv tycon tys of
 343            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 344            _other                ->
 345              pprPanic "VectMonad.lookupFamInst: not found: "
 346                       (ppr $ mkTyConApp tycon tys)
 347        }
 348
 349 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 350 initV hsc_env guts info p
 351   = do
 352       eps <- hscEPS hsc_env
 353       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 354       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 355
 356       Just r <- initDs hsc_env (mg_module guts)
 357                                (mg_rdr_env guts)
 358                                (mg_types guts)
 359                                (go instEnvs famInstEnvs)
 360       return r
 361   where
 362
 363     go instEnvs famInstEnvs =
 364       do
 365         builtins <- initBuiltins
 366         r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs)
 367                    emptyLocalEnv
 368         case r of
 369           Yes genv _ x -> return $ Just (new_info genv, x)
 370           No           -> return Nothing
 371
 372     new_info genv = updVectInfo genv (mg_types guts) info
 373