From 9f9fe0b372ebd30d9d4c3da869d4d88eb7f360d3 Mon Sep 17 00:00:00 2001
From: sheaf <sam.derbyshire@gmail.com>
Date: Fri, 7 Mar 2025 19:11:23 +0100
Subject: [PATCH] Add mapMaybeTM method to TrieMap class
This commit adds a new method to the TrieMap class, mapMaybeTM, and
adds implementations to all the instances.
mapMaybeTM is useful when filtering containers that contain other
containers.
---
compiler/GHC/Cmm/Dataflow/Label.hs | 7 ++++++-
compiler/GHC/Core/Map/Expr.hs | 23 +++++++++++++++++++++++
compiler/GHC/Core/Map/Type.hs | 29 +++++++++++++++++++++++++++++
compiler/GHC/Data/TrieMap.hs | 24 +++++++++++++++++++++++-
compiler/GHC/Stg/CSE.hs | 3 +++
compiler/GHC/Tc/Solver/Types.hs | 11 +++++++++++
compiler/GHC/Types/Var/Env.hs | 6 +++++-
7 files changed, 100 insertions(+), 3 deletions(-)
diff --git a/compiler/GHC/Cmm/Dataflow/Label.hs b/compiler/GHC/Cmm/Dataflow/Label.hs
index 837ffcd0e6a..a213ada0743 100644
--- a/compiler/GHC/Cmm/Dataflow/Label.hs
+++ b/compiler/GHC/Cmm/Dataflow/Label.hs
@@ -63,6 +63,7 @@ module GHC.Cmm.Dataflow.Label
, mapToList
, mapFromList
, mapFromListWith
+ , mapMapMaybe
) where
import GHC.Prelude
@@ -280,6 +281,9 @@ mapFromList assocs = LM (M.fromList [(lblToUnique k, v) | (k, v) <- assocs])
mapFromListWith :: (v -> v -> v) -> [(Label, v)] -> LabelMap v
mapFromListWith f assocs = LM (M.fromListWith f [(lblToUnique k, v) | (k, v) <- assocs])
+mapMapMaybe :: (a -> Maybe b) -> LabelMap a -> LabelMap b
+mapMapMaybe f (LM m) = LM (M.mapMaybe f m)
+
-----------------------------------------------------------------------------
-- Instances
@@ -298,7 +302,8 @@ instance TrieMap LabelMap where
lookupTM k m = mapLookup k m
alterTM k f m = mapAlter f k m
foldTM k m z = mapFoldr k z m
- filterTM f m = mapFilter f m
+ filterTM f = mapFilter f
+ mapMaybeTM f = mapMapMaybe f
-----------------------------------------------------------------------------
-- FactBase
diff --git a/compiler/GHC/Core/Map/Expr.hs b/compiler/GHC/Core/Map/Expr.hs
index 6d22e3a102e..b6dafb35118 100644
--- a/compiler/GHC/Core/Map/Expr.hs
+++ b/compiler/GHC/Core/Map/Expr.hs
@@ -122,6 +122,7 @@ instance TrieMap CoreMap where
alterTM k f (CoreMap m) = CoreMap (alterTM (deBruijnize k) f m)
foldTM k (CoreMap m) = foldTM k m
filterTM f (CoreMap m) = CoreMap (filterTM f m)
+ mapMaybeTM f (CoreMap m) = CoreMap (mapMaybeTM f m)
-- | @CoreMapG a@ is a map from @DeBruijn CoreExpr@ to @a@. The extended
-- key makes it suitable for recursive traversal, since it can track binders,
@@ -271,6 +272,7 @@ instance TrieMap CoreMapX where
alterTM = xtE
foldTM = fdE
filterTM = ftE
+ mapMaybeTM = mpE
--------------------------
ftE :: (a->Bool) -> CoreMapX a -> CoreMapX a
@@ -287,6 +289,20 @@ ftE f (CM { cm_var = cvar, cm_lit = clit
, cm_letr = fmap (fmap (filterTM f)) cletr, cm_case = fmap (filterTM f) ccase
, cm_ecase = fmap (filterTM f) cecase, cm_tick = fmap (filterTM f) ctick }
+mpE :: (a -> Maybe b) -> CoreMapX a -> CoreMapX b
+mpE f (CM { cm_var = cvar, cm_lit = clit
+ , cm_co = cco, cm_type = ctype
+ , cm_cast = ccast , cm_app = capp
+ , cm_lam = clam, cm_letn = cletn
+ , cm_letr = cletr, cm_case = ccase
+ , cm_ecase = cecase, cm_tick = ctick })
+ = CM { cm_var = mapMaybeTM f cvar, cm_lit = mapMaybeTM f clit
+ , cm_co = mapMaybeTM f cco, cm_type = mapMaybeTM f ctype
+ , cm_cast = fmap (mapMaybeTM f) ccast, cm_app = fmap (mapMaybeTM f) capp
+ , cm_lam = fmap (mapMaybeTM f) clam, cm_letn = fmap (fmap (mapMaybeTM f)) cletn
+ , cm_letr = fmap (fmap (mapMaybeTM f)) cletr, cm_case = fmap (mapMaybeTM f) ccase
+ , cm_ecase = fmap (mapMaybeTM f) cecase, cm_tick = fmap (mapMaybeTM f) ctick }
+
--------------------------
lookupCoreMap :: CoreMap a -> CoreExpr -> Maybe a
lookupCoreMap cm e = lookupTM e cm
@@ -409,6 +425,7 @@ instance TrieMap AltMap where
alterTM = xtA emptyCME
foldTM = fdA
filterTM = ftA
+ mapMaybeTM = mpA
instance Eq (DeBruijn CoreAlt) where
D env1 a1 == D env2 a2 = go a1 a2 where
@@ -446,3 +463,9 @@ fdA :: (a -> b -> b) -> AltMap a -> b -> b
fdA k m = foldTM k (am_deflt m)
. foldTM (foldTM k) (am_data m)
. foldTM (foldTM k) (am_lit m)
+
+mpA :: (a -> Maybe b) -> AltMap a -> AltMap b
+mpA f (AM { am_deflt = adeflt, am_data = adata, am_lit = alit })
+ = AM { am_deflt = mapMaybeTM f adeflt
+ , am_data = fmap (mapMaybeTM f) adata
+ , am_lit = fmap (mapMaybeTM f) alit }
diff --git a/compiler/GHC/Core/Map/Type.hs b/compiler/GHC/Core/Map/Type.hs
index 545c3897418..4e3442f65c5 100644
--- a/compiler/GHC/Core/Map/Type.hs
+++ b/compiler/GHC/Core/Map/Type.hs
@@ -96,6 +96,7 @@ instance TrieMap CoercionMap where
alterTM k f (CoercionMap m) = CoercionMap (alterTM (deBruijnize k) f m)
foldTM k (CoercionMap m) = foldTM k m
filterTM f (CoercionMap m) = CoercionMap (filterTM f m)
+ mapMaybeTM f (CoercionMap m) = CoercionMap (mapMaybeTM f m)
type CoercionMapG = GenMap CoercionMapX
newtype CoercionMapX a = CoercionMapX (TypeMapX a)
@@ -112,6 +113,7 @@ instance TrieMap CoercionMapX where
alterTM = xtC
foldTM f (CoercionMapX core_tm) = foldTM f core_tm
filterTM f (CoercionMapX core_tm) = CoercionMapX (filterTM f core_tm)
+ mapMaybeTM f (CoercionMapX core_tm) = CoercionMapX (mapMaybeTM f core_tm)
instance Eq (DeBruijn Coercion) where
D env1 co1 == D env2 co2
@@ -189,6 +191,7 @@ instance TrieMap TypeMapX where
alterTM = xtT
foldTM = fdT
filterTM = filterT
+ mapMaybeTM = mpT
instance Eq (DeBruijn Type) where
(==) = eqDeBruijnType
@@ -380,6 +383,7 @@ instance TrieMap TyLitMap where
alterTM = xtTyLit
foldTM = foldTyLit
filterTM = filterTyLit
+ mapMaybeTM = mpTyLit
emptyTyLitMap :: TyLitMap a
emptyTyLitMap = TLM { tlm_number = Map.empty, tlm_string = emptyUFM, tlm_char = Map.empty }
@@ -407,6 +411,10 @@ filterTyLit :: (a -> Bool) -> TyLitMap a -> TyLitMap a
filterTyLit f (TLM { tlm_number = tn, tlm_string = ts, tlm_char = tc })
= TLM { tlm_number = Map.filter f tn, tlm_string = filterUFM f ts, tlm_char = Map.filter f tc }
+mpTyLit :: (a -> Maybe b) -> TyLitMap a -> TyLitMap b
+mpTyLit f (TLM { tlm_number = tn, tlm_string = ts, tlm_char = tc })
+ = TLM { tlm_number = Map.mapMaybe f tn, tlm_string = mapMaybeUFM f ts, tlm_char = Map.mapMaybe f tc }
+
-------------------------------------------------
-- | @TypeMap a@ is a map from 'Type' to @a@. If you are a client, this
-- is the type you want. The keys in this map may have different kinds.
@@ -435,6 +443,7 @@ instance TrieMap TypeMap where
alterTM k f m = xtTT (deBruijnize k) f m
foldTM k (TypeMap m) = foldTM (foldTM k) m
filterTM f (TypeMap m) = TypeMap (fmap (filterTM f) m)
+ mapMaybeTM f (TypeMap m) = TypeMap (fmap (mapMaybeTM f) m)
foldTypeMap :: (a -> b -> b) -> b -> TypeMap a -> b
foldTypeMap k z m = foldTM k m z
@@ -479,6 +488,7 @@ instance TrieMap LooseTypeMap where
alterTM k f (LooseTypeMap m) = LooseTypeMap (alterTM (deBruijnize k) f m)
foldTM f (LooseTypeMap m) = foldTM f m
filterTM f (LooseTypeMap m) = LooseTypeMap (filterTM f m)
+ mapMaybeTM f (LooseTypeMap m) = LooseTypeMap (mapMaybeTM f m)
{-
************************************************************************
@@ -558,10 +568,13 @@ instance TrieMap BndrMap where
alterTM = xtBndr emptyCME
foldTM = fdBndrMap
filterTM = ftBndrMap
+ mapMaybeTM = mpBndrMap
fdBndrMap :: (a -> b -> b) -> BndrMap a -> b -> b
fdBndrMap f (BndrMap tm) = foldTM (foldTM f) tm
+mpBndrMap :: (a -> Maybe b) -> BndrMap a -> BndrMap b
+mpBndrMap f (BndrMap tm) = BndrMap (fmap (mapMaybeTM f) tm)
-- We need to use 'BndrMap' for 'Coercion', 'CoreExpr' AND 'Type', since all
-- of these data types have binding forms.
@@ -594,6 +607,7 @@ instance TrieMap VarMap where
alterTM = xtVar emptyCME
foldTM = fdVar
filterTM = ftVar
+ mapMaybeTM = mpVar
lkVar :: CmEnv -> Var -> VarMap a -> Maybe a
lkVar env v
@@ -619,9 +633,24 @@ ftVar :: (a -> Bool) -> VarMap a -> VarMap a
ftVar f (VM { vm_bvar = bv, vm_fvar = fv })
= VM { vm_bvar = filterTM f bv, vm_fvar = filterTM f fv }
+mpVar :: (a -> Maybe b) -> VarMap a -> VarMap b
+mpVar f (VM { vm_bvar = bv, vm_fvar = fv })
+ = VM { vm_bvar = mapMaybeTM f bv, vm_fvar = mapMaybeTM f fv }
+
-------------------------------------------------
lkDNamed :: NamedThing n => n -> DNameEnv a -> Maybe a
lkDNamed n env = lookupDNameEnv env (getName n)
xtDNamed :: NamedThing n => n -> XT a -> DNameEnv a -> DNameEnv a
xtDNamed tc f m = alterDNameEnv f m (getName tc)
+
+mpT :: (a -> Maybe b) -> TypeMapX a -> TypeMapX b
+mpT f (TM { tm_var = tvar, tm_app = tapp, tm_tycon = ttycon
+ , tm_forall = tforall, tm_tylit = tlit
+ , tm_coerce = tcoerce })
+ = TM { tm_var = mapMaybeTM f tvar
+ , tm_app = fmap (mapMaybeTM f) tapp
+ , tm_tycon = mapMaybeTM f ttycon
+ , tm_forall = fmap (mapMaybeTM f) tforall
+ , tm_tylit = mapMaybeTM f tlit
+ , tm_coerce = tcoerce >>= f }
diff --git a/compiler/GHC/Data/TrieMap.hs b/compiler/GHC/Data/TrieMap.hs
index 1f0f5446f4b..e36387b5b00 100644
--- a/compiler/GHC/Data/TrieMap.hs
+++ b/compiler/GHC/Data/TrieMap.hs
@@ -69,7 +69,7 @@ class Functor m => TrieMap m where
lookupTM :: forall b. Key m -> m b -> Maybe b
alterTM :: forall b. Key m -> XT b -> m b -> m b
filterTM :: (a -> Bool) -> m a -> m a
-
+ mapMaybeTM :: (a -> Maybe b) -> m a -> m b
foldTM :: (a -> b -> b) -> m a -> b -> b
-- The unusual argument order here makes
-- it easy to compose calls to foldTM;
@@ -146,6 +146,7 @@ instance TrieMap IntMap.IntMap where
alterTM = xtInt
foldTM k m z = IntMap.foldr k z m
filterTM f m = IntMap.filter f m
+ mapMaybeTM f m = IntMap.mapMaybe f m
xtInt :: Int -> XT a -> IntMap.IntMap a -> IntMap.IntMap a
xtInt k f m = IntMap.alter f k m
@@ -157,6 +158,7 @@ instance Ord k => TrieMap (Map.Map k) where
alterTM k f m = Map.alter f k m
foldTM k m z = Map.foldr k z m
filterTM f m = Map.filter f m
+ mapMaybeTM f m = Map.mapMaybe f m
{-
@@ -233,6 +235,7 @@ instance forall key. Uniquable key => TrieMap (UniqDFM key) where
alterTM k f m = alterUDFM f m k
foldTM k m z = foldUDFM k z m
filterTM f m = filterUDFM f m
+ mapMaybeTM f m = mapMaybeUDFM f m
{-
************************************************************************
@@ -259,6 +262,7 @@ instance TrieMap m => TrieMap (MaybeMap m) where
alterTM = xtMaybe alterTM
foldTM = fdMaybe
filterTM = ftMaybe
+ mapMaybeTM = mpMaybe
instance TrieMap m => Foldable (MaybeMap m) where
foldMap = foldMapTM
@@ -281,6 +285,10 @@ ftMaybe :: TrieMap m => (a -> Bool) -> MaybeMap m a -> MaybeMap m a
ftMaybe f (MM { mm_nothing = mn, mm_just = mj })
= MM { mm_nothing = filterMaybe f mn, mm_just = filterTM f mj }
+mpMaybe :: TrieMap m => (a -> Maybe b) -> MaybeMap m a -> MaybeMap m b
+mpMaybe f (MM { mm_nothing = mn, mm_just = mj })
+ = MM { mm_nothing = mn >>= f, mm_just = mapMaybeTM f mj }
+
foldMaybe :: (a -> b -> b) -> Maybe a -> b -> b
foldMaybe _ Nothing b = b
foldMaybe k (Just a) b = k a b
@@ -314,6 +322,7 @@ instance TrieMap m => TrieMap (ListMap m) where
alterTM = xtList alterTM
foldTM = fdList
filterTM = ftList
+ mapMaybeTM = mpList
instance TrieMap m => Foldable (ListMap m) where
foldMap = foldMapTM
@@ -340,6 +349,10 @@ ftList :: TrieMap m => (a -> Bool) -> ListMap m a -> ListMap m a
ftList f (LM { lm_nil = mnil, lm_cons = mcons })
= LM { lm_nil = filterMaybe f mnil, lm_cons = fmap (filterTM f) mcons }
+mpList :: TrieMap m => (a -> Maybe b) -> ListMap m a -> ListMap m b
+mpList f (LM { lm_nil = mnil, lm_cons = mcons })
+ = LM { lm_nil = mnil >>= f, lm_cons = fmap (mapMaybeTM f) mcons }
+
{-
************************************************************************
* *
@@ -395,6 +408,7 @@ instance (Eq (Key m), TrieMap m) => TrieMap (GenMap m) where
alterTM = xtG
foldTM = fdG
filterTM = ftG
+ mapMaybeTM = mpG
instance (Eq (Key m), TrieMap m) => Foldable (GenMap m) where
foldMap = foldMapTM
@@ -457,3 +471,11 @@ ftG f input@(SingletonMap _ v)
ftG f (MultiMap m) = MultiMap (filterTM f m)
-- we don't have enough information to reconstruct the key to make
-- a SingletonMap
+
+{-# INLINEABLE mpG #-}
+mpG :: TrieMap m => (a -> Maybe b) -> GenMap m a -> GenMap m b
+mpG _ EmptyMap = EmptyMap
+mpG f (SingletonMap k v) = case f v of
+ Just v' -> SingletonMap k v'
+ Nothing -> EmptyMap
+mpG f (MultiMap m) = MultiMap (mapMaybeTM f m)
diff --git a/compiler/GHC/Stg/CSE.hs b/compiler/GHC/Stg/CSE.hs
index 0afe917465a..4e01666246c 100644
--- a/compiler/GHC/Stg/CSE.hs
+++ b/compiler/GHC/Stg/CSE.hs
@@ -142,6 +142,8 @@ instance TrieMap StgArgMap where
foldTM k m = foldTM k (sam_var m) . foldTM k (sam_lit m)
filterTM f (SAM {sam_var = varm, sam_lit = litm}) =
SAM { sam_var = filterTM f varm, sam_lit = filterTM f litm }
+ mapMaybeTM f (SAM {sam_var = varm, sam_lit = litm}) =
+ SAM { sam_var = mapMaybeTM f varm, sam_lit = mapMaybeTM f litm }
newtype ConAppMap a = CAM { un_cam :: DNameEnv (ListMap StgArgMap a) }
@@ -158,6 +160,7 @@ instance TrieMap ConAppMap where
m { un_cam = un_cam m |> xtDNamed dataCon |>> alterTM args f }
foldTM k = un_cam >.> foldTM (foldTM k)
filterTM f = un_cam >.> fmap (filterTM f) >.> CAM
+ mapMaybeTM f = un_cam >.> fmap (mapMaybeTM f) >.> CAM
-----------------
-- The CSE Env --
diff --git a/compiler/GHC/Tc/Solver/Types.hs b/compiler/GHC/Tc/Solver/Types.hs
index 5c3f9d9b102..3fa7ab27f7d 100644
--- a/compiler/GHC/Tc/Solver/Types.hs
+++ b/compiler/GHC/Tc/Solver/Types.hs
@@ -14,6 +14,7 @@ module GHC.Tc.Solver.Types (
TcAppMap, emptyTcAppMap, isEmptyTcAppMap,
insertTcApp, alterTcApp, filterTcAppMap,
+ mapMaybeTcAppMap,
tcAppMapToBag, foldTcAppMap, delTcApp,
EqualCtList, filterEqualCtList, addToEqualCtList
@@ -114,6 +115,16 @@ filterTcAppMap f m = mapMaybeDTyConEnv one_tycon m
where
filtered_tm = filterTM f tm
+mapMaybeTcAppMap :: forall a b. (a -> Maybe b) -> TcAppMap a -> TcAppMap b
+mapMaybeTcAppMap f m = mapMaybeDTyConEnv one_tycon m
+ where
+ one_tycon :: ListMap LooseTypeMap a -> Maybe (ListMap LooseTypeMap b)
+ one_tycon tm
+ | isEmptyTM mapped_tm = Nothing
+ | otherwise = Just mapped_tm
+ where
+ mapped_tm = mapMaybeTM f tm
+
tcAppMapToBag :: TcAppMap a -> Bag a
tcAppMapToBag m = foldTcAppMap consBag m emptyBag
diff --git a/compiler/GHC/Types/Var/Env.hs b/compiler/GHC/Types/Var/Env.hs
index d1543e8eca6..1c4398a62a4 100644
--- a/compiler/GHC/Types/Var/Env.hs
+++ b/compiler/GHC/Types/Var/Env.hs
@@ -74,7 +74,8 @@ module GHC.Types.Var.Env (
-- * TidyEnv and its operation
TidyEnv,
- emptyTidyEnv, mkEmptyTidyEnv, delTidyEnvList
+ emptyTidyEnv, mkEmptyTidyEnv, delTidyEnvList,
+ mapMaybeDVarEnv
) where
import GHC.Prelude
@@ -656,6 +657,9 @@ mapDVarEnv = mapUDFM
filterDVarEnv :: (a -> Bool) -> DVarEnv a -> DVarEnv a
filterDVarEnv = filterUDFM
+mapMaybeDVarEnv :: (a -> Maybe b) -> DVarEnv a -> DVarEnv b
+mapMaybeDVarEnv f = mapMaybeUDFM f
+
alterDVarEnv :: (Maybe a -> Maybe a) -> DVarEnv a -> Var -> DVarEnv a
alterDVarEnv = alterUDFM
--
GitLab