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