From 350b36411c94a10fab03ce4e6beba502deb45bb1 Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Thu, 14 Nov 2024 10:22:01 +1100
Subject: [PATCH] chore: upstream some NameMap functions (#6056)

---
 src/Lean/Data/NameMap.lean | 16 ++++++++++++++++
 src/Lean/Data/RBMap.lean   | 18 ++++++++++++++++++
 src/Lean/Data/RBTree.lean  |  7 +++++++
 3 files changed, 41 insertions(+)

diff --git a/src/Lean/Data/NameMap.lean b/src/Lean/Data/NameMap.lean
index 6593e748bc..04cd225102 100644
--- a/src/Lean/Data/NameMap.lean
+++ b/src/Lean/Data/NameMap.lean
@@ -33,6 +33,16 @@ def find? (m : NameMap α) (n : Name) : Option α := RBMap.find? m n
 instance : ForIn m (NameMap α) (Name × α) :=
   inferInstanceAs (ForIn _ (RBMap ..) ..)
 
+/-- `filter f m` returns the `NameMap` consisting of all
+"`key`/`val`"-pairs in `m` where `f key val` returns `true`. -/
+def filter (f : Name → α → Bool) (m : NameMap α) : NameMap α := RBMap.filter f m
+
+/-- `filterMap f m` filters an `NameMap` and simultaneously modifies the filtered values.
+
+It takes a function `f : Name → α → Option β` and applies `f name` to the value with key `name`.
+The resulting entries with non-`none` value are collected to form the output `NameMap`. -/
+def filterMap (f : Name → α → Option β) (m : NameMap α) : NameMap β := RBMap.filterMap f m
+
 end NameMap
 
 def NameSet := RBTree Name Name.quickCmp
@@ -53,6 +63,9 @@ def append (s t : NameSet) : NameSet :=
 instance : Append NameSet where
   append := NameSet.append
 
+/-- `filter f s` returns the `NameSet` consisting of all `x` in `s` where `f x` returns `true`. -/
+def filter (f : Name → Bool) (s : NameSet) : NameSet := RBTree.filter f s
+
 end NameSet
 
 def NameSSet := SSet Name
@@ -73,6 +86,9 @@ instance : EmptyCollection NameHashSet := ⟨empty⟩
 instance : Inhabited NameHashSet := ⟨{}⟩
 def insert (s : NameHashSet) (n : Name) := Std.HashSet.insert s n
 def contains (s : NameHashSet) (n : Name) : Bool := Std.HashSet.contains s n
+
+/-- `filter f s` returns the `NameHashSet` consisting of all `x` in `s` where `f x` returns `true`. -/
+def filter (f : Name → Bool) (s : NameHashSet) : NameHashSet := Std.HashSet.filter f s
 end NameHashSet
 
 def MacroScopesView.isPrefixOf (v₁ v₂ : MacroScopesView) : Bool :=
diff --git a/src/Lean/Data/RBMap.lean b/src/Lean/Data/RBMap.lean
index 4172b70000..2ac041f43c 100644
--- a/src/Lean/Data/RBMap.lean
+++ b/src/Lean/Data/RBMap.lean
@@ -404,6 +404,24 @@ def intersectBy {γ : Type v₁} {δ : Type v₂} (mergeFn : α → β → γ 
       | some b₂ => acc.insert a <| mergeFn a b₁ b₂
       | none => acc
 
+/--
+`filter f m` returns the `RBMap` consisting of all
+"`key`/`val`"-pairs in `m` where `f key val` returns `true`.
+-/
+def filter (f : α → β → Bool) (m : RBMap α β cmp) : RBMap α β cmp :=
+  m.fold (fun r k v => if f k v then r.insert k v else r) {}
+
+/--
+`filterMap f m` filters an `RBMap` and simultaneously modifies the filtered values.
+
+It takes a function `f : α → β → Option γ` and applies `f k v` to the value with key `k`.
+The resulting entries with non-`none` value are collected to form the output `RBMap`.
+-/
+def filterMap (f : α → β → Option γ) (m : RBMap α β cmp) : RBMap α γ cmp :=
+  m.fold (fun r k v => match f k v with
+    | none => r
+    | some b => r.insert k b) {}
+
 end RBMap
 
 def rbmapOf {α : Type u} {β : Type v} (l : List (α × β)) (cmp : α → α → Ordering) : RBMap α β cmp :=
diff --git a/src/Lean/Data/RBTree.lean b/src/Lean/Data/RBTree.lean
index f749d80f8f..4265bb1ace 100644
--- a/src/Lean/Data/RBTree.lean
+++ b/src/Lean/Data/RBTree.lean
@@ -114,6 +114,13 @@ def union (t₁ t₂ : RBTree α cmp) : RBTree α cmp :=
 def diff (t₁ t₂ : RBTree α cmp) : RBTree α cmp :=
   t₂.fold .erase t₁
 
+/--
+`filter f m` returns the `RBTree` consisting of all
+`x` in `m` where `f x` returns `true`.
+-/
+def filter (f : α → Bool) (m : RBTree α cmp) : RBTree α cmp :=
+  RBMap.filter (fun a _ => f a) m
+
 end RBTree
 
 def rbtreeOf {α : Type u} (l : List α) (cmp : α → α → Ordering) : RBTree α cmp :=