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 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                  -- In-scope type variables
 136                  --
 137                , local_tyvars :: [TyVar]
 138
 139                  -- Mapping from tyvars to their PA dictionaries
 140                , local_tyvar_pa :: VarEnv CoreExpr
 141                }
 142
 143
 144 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 145 initGlobalEnv info instEnvs famInstEnvs
 146   = GlobalEnv {
 147       global_vars          = mapVarEnv  (Var . snd) $ vectInfoVar   info
 148     , global_exported_vars = emptyVarEnv
 149     , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
 150     , global_tycon_pa      = emptyNameEnv
 151     , global_inst_env      = instEnvs
 152     , global_fam_inst_env  = famInstEnvs
 153     , global_bindings      = []
 154     }
 155
 156 emptyLocalEnv = LocalEnv {
 157                    local_vars     = emptyVarEnv
 158                  , local_tyvars   = []
 159                  , local_tyvar_pa = emptyVarEnv
 160                  }
 161
 162 -- FIXME
 163 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 164 updVectInfo env tyenv info
 165   = info {
 166       vectInfoVar   = global_exported_vars env
 167     , vectInfoTyCon = tc_env
 168     }
 169   where
 170     tc_env = mkNameEnv [(tc_name, (tc,tc'))
 171                | tc <- typeEnvTyCons tyenv
 172                , let tc_name = tyConName tc
 173                , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
 174
 175 data VResult a = Yes GlobalEnv LocalEnv a | No
 176
 177 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 178
 179 instance Monad VM where
 180   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 181   VM p >>= f = VM $ \bi genv lenv -> do
 182                                       r <- p bi genv lenv
 183                                       case r of
 184                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 185                                         No                -> return No
 186
 187 noV :: VM a
 188 noV = VM $ \_ _ _ -> return No
 189
 190 tryV :: VM a -> VM (Maybe a)
 191 tryV (VM p) = VM $ \bi genv lenv ->
 192   do
 193     r <- p bi genv lenv
 194     case r of
 195       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 196       No                -> return (Yes genv  lenv  Nothing)
 197
 198 maybeV :: VM (Maybe a) -> VM a
 199 maybeV p = maybe noV return =<< p
 200
 201 orElseV :: VM a -> VM a -> VM a
 202 orElseV p q = maybe q return =<< tryV p
 203
 204 fixV :: (a -> VM a) -> VM a
 205 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 206   where
 207     unYes (Yes _ _ x) = x
 208
 209 localV :: VM a -> VM a
 210 localV p = do
 211              env <- readLEnv id
 212              x <- p
 213              setLEnv env
 214              return x
 215
 216 closedV :: VM a -> VM a
 217 closedV p = do
 218               env <- readLEnv id
 219               setLEnv emptyLocalEnv
 220               x <- p
 221               setLEnv env
 222               return x
 223
 224 liftDs :: DsM a -> VM a
 225 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 226
 227 builtin :: (Builtins -> a) -> VM a
 228 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 229
 230 readGEnv :: (GlobalEnv -> a) -> VM a
 231 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 232
 233 setGEnv :: GlobalEnv -> VM ()
 234 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 235
 236 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 237 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 238
 239 readLEnv :: (LocalEnv -> a) -> VM a
 240 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 241
 242 setLEnv :: LocalEnv -> VM ()
 243 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 244
 245 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 246 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 247
 248 getInstEnv :: VM (InstEnv, InstEnv)
 249 getInstEnv = readGEnv global_inst_env
 250
 251 getFamInstEnv :: VM FamInstEnvs
 252 getFamInstEnv = readGEnv global_fam_inst_env
 253
 254 cloneName :: (OccName -> OccName) -> Name -> VM Name
 255 cloneName mk_occ name = liftM make (liftDs newUnique)
 256   where
 257     occ_name = mk_occ (nameOccName name)
 258
 259     make u | isExternalName name = mkExternalName u (nameModule name)
 260                                                     occ_name
 261                                                     (nameSrcSpan name)
 262            | otherwise           = mkSystemName u occ_name
 263
 264 newLocalVar :: FastString -> Type -> VM Var
 265 newLocalVar fs ty
 266   = do
 267       u <- liftDs newUnique
 268       return $ mkSysLocal fs u ty
 269
 270 newTyVar :: FastString -> Kind -> VM Var
 271 newTyVar fs k
 272   = do
 273       u <- liftDs newUnique
 274       return $ mkTyVar (mkSysTvName u fs) k
 275
 276 defGlobalVar :: Var -> Var -> VM ()
 277 defGlobalVar v v' = updGEnv $ \env ->
 278   env { global_vars = extendVarEnv (global_vars env) v (Var v')
 279       , global_exported_vars = upd (global_exported_vars env)
 280       }
 281   where
 282     upd env | isExportedId v = extendVarEnv env v (v, v')
 283             | otherwise      = env
 284
 285 lookupVar :: Var -> VM (Scope CoreExpr (CoreExpr, CoreExpr))
 286 lookupVar v
 287   = do
 288       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 289       case r of
 290         Just e  -> return (Local e)
 291         Nothing -> liftM Global
 292                  $  maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 293
 294 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 295 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 296
 297 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 298 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 299
 300 defLocalTyVar :: TyVar -> VM ()
 301 defLocalTyVar tv = updLEnv $ \env ->
 302   env { local_tyvars   = tv : local_tyvars env
 303       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 304       }
 305
 306 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 307 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 308   env { local_tyvars   = tv : local_tyvars env
 309       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 310       }
 311
 312 localTyVars :: VM [TyVar]
 313 localTyVars = readLEnv (reverse . local_tyvars)
 314
 315 -- Look up the dfun of a class instance.
 316 --
 317 -- The match must be unique - ie, match exactly one instance - but the
 318 -- type arguments used for matching may be more specific than those of
 319 -- the class instance declaration.  The found class instances must not have
 320 -- any type variables in the instance context that do not appear in the
 321 -- instances head (i.e., no flexi vars); for details for what this means,
 322 -- see the docs at InstEnv.lookupInstEnv.
 323 --
 324 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 325 lookupInst cls tys
 326   = do { instEnv <- getInstEnv
 327        ; case lookupInstEnv instEnv cls tys of
 328            ([(inst, inst_tys)], _)
 329              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 330              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 331                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 332              where
 333                inst_tys'  = [ty | Right ty <- inst_tys]
 334                noFlexiVar = all isRight inst_tys
 335            _other         -> noV
 336        }
 337   where
 338     isRight (Left  _) = False
 339     isRight (Right _) = True
 340
 341 -- Look up the representation tycon of a family instance.
 342 --
 343 -- The match must be unique - ie, match exactly one instance - but the
 344 -- type arguments used for matching may be more specific than those of
 345 -- the family instance declaration.
 346 --
 347 -- Return the instance tycon and its type instance.  For example, if we have
 348 --
 349 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 350 --
 351 -- then we have a coercion (ie, type instance of family instance coercion)
 352 --
 353 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 354 --
 355 -- which implies that :R42T was declared as 'data instance T [a]'.
 356 --
 357 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 358 lookupFamInst tycon tys
 359   = ASSERT( isOpenTyCon tycon )
 360     do { instEnv <- getFamInstEnv
 361        ; case lookupFamInstEnv instEnv tycon tys of
 362            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 363            _other                ->
 364              pprPanic "VectMonad.lookupFamInst: not found: "
 365                       (ppr $ mkTyConApp tycon tys)
 366        }
 367
 368 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 369 initV hsc_env guts info p
 370   = do
 371       eps <- hscEPS hsc_env
 372       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 373       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 374
 375       Just r <- initDs hsc_env (mg_module guts)
 376                                (mg_rdr_env guts)
 377                                (mg_types guts)
 378                                (go instEnvs famInstEnvs)
 379       return r
 380   where
 381
 382     go instEnvs famInstEnvs =
 383       do
 384         builtins <- initBuiltins
 385         r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs)
 386                    emptyLocalEnv
 387         case r of
 388           Yes genv _ x -> return $ Just (new_info genv, x)
 389           No           -> return Nothing
 390
 391     new_info genv = updVectInfo genv (mg_types guts) info
 392