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