feat: RBMap lemmas cont'd: find?_erase

Std/Data/List/Lemmas.lean

@@ -1348,6 +1348,13 @@ theorem any_eq_not_all_not (l : List α) (p : α → Bool) : l.any p = !l.all (!
 theorem all_eq_not_any_not (l : List α) (p : α → Bool) : l.all p = !l.any (!p .) := by
   simp only [not_any_eq_all_not, Bool.not_not]
 
+/-! ### all₂ -/
+
+theorem beq_eq_all₂ [BEq α] (l₁ l₂ : List α) : (l₁ == l₂) = l₁.all₂ BEq.beq l₂ := by
+  simp [BEq.beq]
+  induction l₁ generalizing l₂ <;> cases l₂ <;> simp [List.beq, all₂]
+  next a _ ih b _ => cases a == b <;> simp [ih]
+
 /-! ### reverse -/
 
 @[simp] theorem mem_reverseAux {x : α} : ∀ {as bs}, x ∈ reverseAux as bs ↔ x ∈ as ∨ x ∈ bs

Std/Data/RBMap/Alter.lean

@@ -71,8 +71,7 @@ theorem insertNew_eq_insert (h : zoom (cmp v) t = (nil, path)) :
 
 theorem zoom_del {t : RBNode α} :
     t.zoom cut path = (t', path') →
-    path.del (t.del cut) (match t with | node c .. => c | _ => red) =
-    path'.del t'.delRoot (match t' with | node c .. => c | _ => red) := by
+    path.del (t.del cut) t.isBlack = path'.del t'.delRoot t'.isBlack := by
   unfold RBNode.del; split <;> simp [zoom]
   · intro | rfl, rfl => rfl
   · next c a y b =>
@@ -279,19 +278,19 @@ protected theorem Ordered.fill : ∀ {path : Path α} {t},
       fun ⟨hp, ⟨ax, xb, ha, hb⟩, ⟨xp, ap, bp⟩⟩ => ⟨⟨hp, ax, xp, ap, ha⟩, hb, ⟨xb, bp⟩⟩,
       fun ⟨⟨hp, ax, xp, ap, ha⟩, hb, ⟨xb, bp⟩⟩ => ⟨hp, ⟨ax, xb, ha, hb⟩, ⟨xp, ap, bp⟩⟩⟩
 
-theorem _root_.Std.RBNode.Ordered.zoom' {t : RBNode α} {path : Path α}
+protected theorem _root_.Std.RBNode.Ordered.zoom' {t : RBNode α} {path : Path α}
     (ht : t.Ordered cmp) (hp : path.Ordered cmp) (tp : t.All (path.RootOrdered cmp))
     (pz : path.Zoomed cut) (eq : t.zoom cut path = (t', path')) :
     t'.Ordered cmp ∧ path'.Ordered cmp ∧ t'.All (path'.RootOrdered cmp) ∧ path'.Zoomed cut :=
   have ⟨hp', ht', tp'⟩ := Ordered.fill.1 <| zoom_fill eq ▸ Ordered.fill.2 ⟨hp, ht, tp⟩
   ⟨ht', hp', tp', zoom_zoomed₂ eq pz⟩
 
-theorem _root_.Std.RBNode.Ordered.zoom {t : RBNode α}
+protected theorem _root_.Std.RBNode.Ordered.zoom {t : RBNode α}
     (ht : t.Ordered cmp) (eq : t.zoom cut = (t', path')) :
     t'.Ordered cmp ∧ path'.Ordered cmp ∧ t'.All (path'.RootOrdered cmp) ∧ path'.Zoomed cut :=
   ht.zoom' (path := .root) ⟨⟩ (.trivial ⟨⟩) ⟨⟩ eq
 
-theorem Ordered.ins : ∀ {path : Path α} {t : RBNode α},
+protected theorem Ordered.ins : ∀ {path : Path α} {t : RBNode α},
     t.Ordered cmp → path.Ordered cmp → t.All (path.RootOrdered cmp) → (path.ins t).Ordered cmp
   | .root, t, ht, _, _ => Ordered.setBlack.2 ht
   | .left red parent x b, a, ha, ⟨hp, xb, xp, bp, hb⟩, H => by
@@ -305,18 +304,18 @@ theorem Ordered.ins : ∀ {path : Path α} {t : RBNode α},
     unfold ins; have ⟨xb, bp⟩ := All_and.1 H
     exact hp.ins (ha.balance2 ax xb hb) (balance2_All.2 ⟨xp, ap, bp⟩)
 
-theorem Ordered.insertNew {path : Path α} (hp : path.Ordered cmp) (vp : path.RootOrdered cmp v) :
-    (path.insertNew v).Ordered cmp :=
+protected theorem Ordered.insertNew {path : Path α}
+    (hp : path.Ordered cmp) (vp : path.RootOrdered cmp v) : (path.insertNew v).Ordered cmp :=
   hp.ins ⟨⟨⟩, ⟨⟩, ⟨⟩, ⟨⟩⟩ ⟨vp, ⟨⟩, ⟨⟩⟩
 
-theorem Ordered.insert : ∀ {path : Path α} {t : RBNode α},
+protected theorem Ordered.insert : ∀ {path : Path α} {t : RBNode α},
     path.Ordered cmp → t.Ordered cmp → t.All (path.RootOrdered cmp) → path.RootOrdered cmp v →
     t.OnRoot (cmpEq cmp v) → (path.insert t v).Ordered cmp
   | _, nil, hp, _, _, vp, _ => hp.insertNew vp
   | _, node .., hp, ⟨ax, xb, ha, hb⟩, ⟨_, ap, bp⟩, vp, xv => Ordered.fill.2
     ⟨hp, ⟨ax.imp xv.lt_congr_right.2, xb.imp xv.lt_congr_left.2, ha, hb⟩, vp, ap, bp⟩
 
-theorem Ordered.del : ∀ {path : Path α} {t : RBNode α} {c},
+protected theorem Ordered.del : ∀ {path : Path α} {t : RBNode α} {c},
     t.Ordered cmp → path.Ordered cmp → t.All (path.RootOrdered cmp) → (path.del t c).Ordered cmp
   | .root, t, _, ht, _, _ => Ordered.setBlack.2 ht
   | .left _ parent x b, a, red, ha, ⟨hp, xb, xp, bp, hb⟩, H => by
@@ -330,7 +329,7 @@ theorem Ordered.del : ∀ {path : Path α} {t : RBNode α} {c},
     unfold del; have ⟨xb, bp⟩ := All_and.1 H
     exact hp.del (ha.balRight ax xb hb) (ap.balRight xp bp)
 
-theorem Ordered.erase : ∀ {path : Path α} {t : RBNode α},
+protected theorem Ordered.erase : ∀ {path : Path α} {t : RBNode α},
     path.Ordered cmp → t.Ordered cmp → t.All (path.RootOrdered cmp) → (path.erase t).Ordered cmp
   | _, nil, hp, ht, tp => Ordered.fill.2 ⟨hp, ht, tp⟩
   | _, node .., hp, ⟨ax, xb, ha, hb⟩, ⟨_, ap, bp⟩ => hp.del (ha.append ax xb hb) (ap.append bp)

Std/Data/RBMap/Basic.lean

@@ -119,11 +119,11 @@ An auxiliary data structure (an iterator) over an `RBNode` which lazily
 pulls elements from the left.
 -/
 protected inductive Stream (α : Type _)
-  | /-- The stream is empty. -/
-    nil
-  | /-- We are ready to deliver element `v` with right child `r`,
-    and where `tail` represents all the subtrees we have yet to destructure. -/
-    cons (v : α) (r : RBNode α) (tail : RBNode.Stream α)
+  /-- The stream is empty. -/
+  | nil
+  /-- We are ready to deliver element `v` with right child `r`,
+  and where `tail` represents all the subtrees we have yet to destructure. -/
+  | cons (v : α) (r : RBNode α) (tail : RBNode.Stream α)
 
 /-- `O(log n)`. Turn a node into a stream, by descending along the left spine. -/
 def toStream : RBNode α → (_ : RBNode.Stream α := .nil) → RBNode.Stream α
@@ -407,8 +407,17 @@ are not of this form as long as they are suitably monotonic.)
     | .gt => find? cut b
     | .eq => some y
 
+/-- `upperBound? cut` retrieves the smallest entry larger than or equal to `cut`, if it exists. -/
+@[specialize] def upperBound? (cut : α → Ordering) : RBNode α → (ub : Option α := .none) → Option α
+  | nil,          ub => ub
+  | node _ a y b, ub =>
+    match cut y with
+    | .lt => upperBound? cut a (some y)
+    | .gt => upperBound? cut b ub
+    | .eq => some y
+
 /-- `lowerBound? cut` retrieves the largest entry smaller than or equal to `cut`, if it exists. -/
-@[specialize] def lowerBound? (cut : α → Ordering) : RBNode α → Option α → Option α
+@[specialize] def lowerBound? (cut : α → Ordering) : RBNode α → (lb : Option α := .none) → Option α
   | nil,          lb => lb
   | node _ a y b, lb =>
     match cut y with
@@ -677,9 +686,12 @@ if it returns `.eq` we will remove the element.
 (The function `cmp k` for some key `k` is a valid cut function, but we can also use cuts that
 are not of this form as long as they are suitably monotonic.)
 -/
-@[inline] def erase (t : RBSet α cmp) (cut : α → Ordering) : RBSet α cmp :=
+@[inline] def eraseP (t : RBSet α cmp) (cut : α → Ordering) : RBSet α cmp :=
   ⟨t.1.erase cut, t.2.erase⟩
 
+/-- `O(log n)`. Remove element `x` from the tree, if present. -/
+@[inline] def erase (t : RBSet α cmp) (x : α) : RBSet α cmp := t.eraseP (cmp x)
+
 /-- `O(log n)`. Find an element in the tree using a cut function. -/
 @[inline] def findP? (t : RBSet α cmp) (cut : α → Ordering) : Option α := t.1.find? cut
 
@@ -690,17 +702,28 @@ are not of this form as long as they are suitably monotonic.)
 @[inline] def findPD (t : RBSet α cmp) (cut : α → Ordering) (v₀ : α) : α := (t.findP? cut).getD v₀
 
 /--
-`O(log n)`. `lowerBoundP cut` retrieves the largest entry comparing `lt` or `eq` under `cut`,
+`O(log n)`. `upperBoundP cut` retrieves the smallest entry comparing `gt` or `eq` under `cut`,
+if it exists. If multiple keys in the map return `eq` under `cut`, any of them may be returned.
+-/
+@[inline] def upperBoundP? (t : RBSet α cmp) (cut : α → Ordering) : Option α := t.1.upperBound? cut
+
+/--
+`O(log n)`. `upperBound? k` retrieves the smallest entry larger than or equal to `k`,
 if it exists.
 -/
-@[inline] def lowerBoundP? (t : RBSet α cmp) (cut : α → Ordering) : Option α :=
-  t.1.lowerBound? cut none
+@[inline] def upperBound? (t : RBSet α cmp) (k : α) : Option α := t.upperBoundP? (cmp k)
+
+/--
+`O(log n)`. `lowerBoundP cut` retrieves the largest entry comparing `lt` or `eq` under `cut`,
+if it exists. If multiple keys in the map return `eq` under `cut`, any of them may be returned.
+-/
+@[inline] def lowerBoundP? (t : RBSet α cmp) (cut : α → Ordering) : Option α := t.1.lowerBound? cut
 
 /--
 `O(log n)`. `lowerBound? k` retrieves the largest entry smaller than or equal to `k`,
 if it exists.
 -/
-@[inline] def lowerBound? (t : RBSet α cmp) (k : α) : Option α := t.1.lowerBound? (cmp k) none
+@[inline] def lowerBound? (t : RBSet α cmp) (k : α) : Option α := t.lowerBoundP? (cmp k)
 
 /-- `O(log n)`. Returns true if the given cut returns `eq` for something in the RBSet. -/
 @[inline] def containsP (t : RBSet α cmp) (cut : α → Ordering) : Bool := (t.findP? cut).isSome
@@ -1029,7 +1052,7 @@ instance [Repr α] [Repr β] : Repr (RBMap α β cmp) where
 @[inline] def insert (t : RBMap α β cmp) (k : α) (v : β) : RBMap α β cmp := RBSet.insert t (k, v)
 
 /-- `O(log n)`. Remove an element `k` from the map. -/
-@[inline] def erase (t : RBMap α β cmp) (k : α) : RBMap α β cmp := RBSet.erase t (cmp k ·.1)
+@[inline] def erase (t : RBMap α β cmp) (k : α) : RBMap α β cmp := RBSet.eraseP t (cmp k ·.1)
 
 /-- `O(n log n)`. Build a tree from an unsorted list by inserting them one at a time. -/
 @[inline] def ofList (l : List (α × β)) (cmp : α → α → Ordering) : RBMap α β cmp :=
@@ -1048,6 +1071,13 @@ instance [Repr α] [Repr β] : Repr (RBMap α β cmp) where
 /-- `O(log n)`. Find the value corresponding to key `k`, or return `v₀` if it is not in the map. -/
 @[inline] def findD (t : RBMap α β cmp) (k : α) (v₀ : β) : β := (t.find? k).getD v₀
 
+/--
+`O(log n)`. `upperBound? k` retrieves the key-value pair of the smallest key
+larger than or equal to `k`, if it exists.
+-/
+@[inline] def upperBound? (t : RBMap α β cmp) (k : α) : Option (α × β) :=
+   RBSet.upperBoundP? t (cmp k ·.1)
+
 /--
 `O(log n)`. `lowerBound? k` retrieves the key-value pair of the largest key
 smaller than or equal to `k`, if it exists.