diff --git a/src/Lean/Data/RBMap.lean b/src/Lean/Data/RBMap.lean index dd3186f49c..0aefdd4793 100644 --- a/src/Lean/Data/RBMap.lean +++ b/src/Lean/Data/RBMap.lean @@ -225,6 +225,28 @@ inductive WellFormed (cmp : α → α → Ordering) : RBNode α β → Prop wher | insertWff {n n' : RBNode α β} {k : α} {v : β k} : WellFormed cmp n → n' = insert cmp n k v → WellFormed cmp n' | eraseWff {n n' : RBNode α β} {k : α} : WellFormed cmp n → n' = erase cmp k n → WellFormed cmp n' +section Map + +@[specialize] def mapM {α : Type v} {β γ : α → Type v} {M : Type v → Type v} [Applicative M] + (f : (a : α) → β a → M (γ a)) + : RBNode α β → M (RBNode α γ) + | leaf => pure leaf + | node color lchild key val rchild => + pure (node color · key · ·) <*> lchild.mapM f <*> f _ val <*> rchild.mapM f + +@[specialize] def map {α : Type u} {β γ : α → Type v} + (f : {a : α} → β a → γ a) + : RBNode α β → RBNode α γ + | leaf => leaf + | node color lchild key val rchild => node color (lchild.map f) key (f val) (rchild.map f) + +end Map + +def toArray (n : RBNode α β) : Array (Sigma β) := + n.fold (init := ∅) fun acc k v => acc.push ⟨k,v⟩ + +instance : EmptyCollection (RBNode α β) := ⟨leaf⟩ + end RBNode open Std.RBNode @@ -276,12 +298,14 @@ instance : ForIn m (RBMap α β cmp) (α × β) where @[specialize] def toList : RBMap α β cmp → List (α × β) | ⟨t, _⟩ => t.revFold (fun ps k v => (k, v)::ps) [] +/-- Returns the kv pair `(a,b)` such that `a ≤ k` for all keys in the RBMap. -/ @[inline] protected def min : RBMap α β cmp → Option (α × β) | ⟨t, _⟩ => match t.min with | some ⟨k, v⟩ => some (k, v) | none => none +/-- Returns the kv pair `(a,b)` such that `a ≥ k` for all keys in the RBMap. -/ @[inline] protected def max : RBMap α β cmp → Option (α × β) | ⟨t, _⟩ => match t.max with @@ -315,18 +339,22 @@ instance [Repr α] [Repr β] : Repr (RBMap α β cmp) where @[inline] def lowerBound : RBMap α β cmp → α → Option (Sigma (fun (_ : α) => β)) | ⟨t, _⟩, x => t.lowerBound cmp x none +/-- Returns true if the given key `a` is in the RBMap.-/ @[inline] def contains (t : RBMap α β cmp) (a : α) : Bool := (t.find? a).isSome @[inline] def fromList (l : List (α × β)) (cmp : α → α → Ordering) : RBMap α β cmp := l.foldl (fun r p => r.insert p.1 p.2) (mkRBMap α β cmp) +/-- Returns true if the given predicate is true for all items in the RBMap. -/ @[inline] def all : RBMap α β cmp → (α → β → Bool) → Bool | ⟨t, _⟩, p => t.all p +/-- Returns true if the given predicate is true for any item in the RBMap. -/ @[inline] def any : RBMap α β cmp → (α → β → Bool) → Bool | ⟨t, _⟩, p => t.any p +/-- The number of items in the RBMap. -/ def size (m : RBMap α β cmp) : Nat := m.fold (fun sz _ _ => sz+1) 0 @@ -343,12 +371,24 @@ def maxDepth (t : RBMap α β cmp) : Nat := | some p => p | none => panic! "map is empty" +/-- Attempts to find the value with key `k : α` in `t` and panics if there is no such key.-/ @[inline] def find! [Inhabited β] (t : RBMap α β cmp) (k : α) : β := match t.find? k with | some b => b | none => panic! "key is not in the map" +/-- Merges the maps `t₁` and `t₂`, if a key `a : α` exists in both, +then use `mergeFn a b₁ b₂` to produce the new merged value. -/ +def mergeBy (mergeFn : α → β → β → β) (t₁ t₂ : RBMap α β cmp) : RBMap α β cmp := + t₂.fold (init := t₁) fun t₁ a b₂ => + t₁.insert a <| + match t₁.find? a with + | some b₁ => mergeFn a b₁ b₂ + | none => b₂ + end RBMap def rbmapOf {α : Type u} {β : Type v} (l : List (α × β)) (cmp : α → α → Ordering) : RBMap α β cmp := RBMap.fromList l cmp + +end Std