feat: OrderedAssocList

Batteries.lean

@@ -27,6 +27,7 @@ import Batteries.Data.LazyList
 import Batteries.Data.List
 import Batteries.Data.MLList
 import Batteries.Data.Nat
+import Batteries.Data.OrderedAssocList
 import Batteries.Data.PairingHeap
 import Batteries.Data.RBMap
 import Batteries.Data.Range

Batteries/Classes/Order.lean

@@ -30,11 +30,20 @@ end Ordering
 
 namespace Batteries
 
+/-- Construct a `BEq` instance from a comparator function. -/
+def beqOfCmp (cmp : α → α → Ordering) : BEq α where
+  beq x y := cmp x y = .eq
+
 /-- `TotalBLE le` asserts that `le` has a total order, that is, `le a b ∨ le b a`. -/
 class TotalBLE (le : α → α → Bool) : Prop where
   /-- `le` is total: either `le a b` or `le b a`. -/
   total : le a b ∨ le b a
 
+/-- `AntisymmCmp cmp` asserts that `cmp x y = .eq` only if `x = y`. -/
+class AntisymmCmp (cmp : α → α → Ordering) : Prop where
+  /-- If two terms compare as `.eq`, then they are equal. -/
+  eq_of_cmp_eq : cmp x y = .eq → x = y
+
 /-- `OrientedCmp cmp` asserts that `cmp` is determined by the relation `cmp x y = .lt`. -/
 class OrientedCmp (cmp : α → α → Ordering) : Prop where
   /-- The comparator operation is symmetric, in the sense that if `cmp x y` equals `.lt` then
@@ -59,6 +68,21 @@ theorem cmp_refl [OrientedCmp cmp] : cmp x x = .eq :=
 
 end OrientedCmp
 
+/--
+The `BEq` instance from a comparator `cmp : α → α → Ordering` with `[AntisymmCmp cmp]` and
+`[OrientedCmp cmp]` is a `LawfulBEq`.
+-/
+theorem lawfulBEqOfCmp (cmp : α → α → Ordering) [AntisymmCmp cmp] [OrientedCmp cmp] :
+    letI := beqOfCmp cmp
+    LawfulBEq α :=
+  let _ := beqOfCmp cmp
+  { eq_of_beq := fun h => by
+      apply AntisymmCmp.eq_of_cmp_eq (cmp := cmp)
+      simpa [beqOfCmp] using h
+    rfl := by
+      intro a
+      simpa [beqOfCmp] using OrientedCmp.cmp_refl }
+
 /-- `TransCmp cmp` asserts that `cmp` induces a transitive relation. -/
 class TransCmp (cmp : α → α → Ordering) extends OrientedCmp cmp : Prop where
   /-- The comparator operation is transitive. -/

Batteries/Data/AssocList.lean

@@ -113,6 +113,43 @@ def toListTR (as : AssocList α β) : List (α × β) :=
 @[simp] theorem length_mapVal : (mapVal f l).length = l.length := by
   induction l <;> simp_all
 
+
+/-- `O(n)`. Map a function `f` over the values of the list, dropping `none`. -/
+def filterMapVal (f : α → β → Option δ) : AssocList α β → AssocList α δ
+  | nil        => nil
+  | cons k v t => match f k v with
+    | none => filterMapVal f t
+    | some d => cons k d (filterMapVal f t)
+
+@[simp] theorem filterMapVal_nil : filterMapVal f nil = nil := rfl
+
+theorem filterMapVal_cons (f : α → β → Option γ) (k) (v) (t) :
+    filterMapVal f (.cons k v t) =
+      match f k v with
+      | none => filterMapVal f t
+      | some d => .cons k d (filterMapVal f t) := rfl
+
+@[simp] theorem toList_filterMapVal (f : α → β → Option δ) (l : AssocList α β) :
+    (filterMapVal f l).toList =
+      l.toList.filterMap (fun (a, b) => (f a b).map fun v => (a, v)) := by
+  induction l with
+  | nil => simp
+  | cons k v t ih =>
+    revert ih
+    simp only [filterMapVal, toList, List.filterMap_cons]
+    match f k v with
+    | none
+    | some d => simp
+
+theorem length_filterMapVal : (filterMapVal f l).length ≤ l.length := by
+  induction l with
+  | nil => simp
+  | cons k v t ih =>
+    simp_all only [filterMapVal, length_cons]
+    match f k v with
+    | none => exact Nat.le_trans ih (Nat.le_succ _)
+    | some _ => exact Nat.succ_le_succ ih
+
 /-- `O(n)`. Returns the first entry in the list whose entry satisfies `p`. -/
 @[specialize] def findEntryP? (p : α → β → Bool) : AssocList α β → Option (α × β)
   | nil         => none
@@ -140,7 +177,7 @@ theorem find?_eq_findEntry? [BEq α] (a : α) (l : AssocList α β) :
     find? a l = (l.findEntry? a).map (·.2) := by
   induction l <;> simp [find?, List.find?_cons]; split <;> simp [*]
 
-@[simp] theorem find?_eq [BEq α] (a : α) (l : AssocList α β) :
+theorem find?_eq [BEq α] (a : α) (l : AssocList α β) :
     find? a l = (l.toList.find? (·.1 == a)).map (·.2) := by simp [find?_eq_findEntry?]
 
 /-- `O(n)`. Returns true if any entry in the list satisfies `p`. -/

Batteries/Data/HashMap/Lemmas.lean

@@ -24,7 +24,7 @@ namespace Imp
 end Imp
 
 @[simp] theorem empty_find? [BEq α] [Hashable α] {a : α} :
-    (∅ : HashMap α β).find? a = none := by simp [find?, Imp.find?]
+    (∅ : HashMap α β).find? a = none := by simp [find?, Imp.find?, AssocList.find?]
 
 -- `Std.HashMap` has this lemma (as `getElem?_insert`) and many more, so working on this
 -- `proof_wanted` is likely not a good use of your time.