feat: extensional hash maps (leanprover#8004)

This PR adds extensional hash maps and hash sets under the names
`Std.ExtDHashMap`, `Std.ExtHashMap` and `Std.ExtHashSet`. Extensional
hash maps work like regular hash maps, except that they have
extensionality lemmas which make them easier to use in proofs. This
however makes it also impossible to regularly iterate over its entries.

Author: Rob23oba

src/Init/Core.lean

@@ -938,6 +938,34 @@ term.
 theorem eqRec_heq {α : Sort u} {φ : α → Sort v} {a a' : α} : (h : a = a') → (p : φ a) → HEq (Eq.recOn (motive := fun x _ => φ x) h p) p
   | rfl, p => HEq.refl p
 
+/--
+Heterogenous equality with an `Eq.rec` application on the left is equivalent to a heterogenous
+equality on the original term.
+-/
+theorem eqRec_heq_iff {α : Sort u} {a : α} {motive : (b : α) → a = b → Sort v}
+    {b : α} {refl : motive a (Eq.refl a)} {h : a = b} {c : motive b h} :
+    HEq (@Eq.rec α a motive refl b h) c ↔ HEq refl c :=
+  h.rec (fun _ => ⟨id, id⟩) c
+
+/--
+Heterogenous equality with an `Eq.rec` application on the right is equivalent to a heterogenous
+equality on the original term.
+-/
+theorem heq_eqRec_iff {α : Sort u} {a : α} {motive : (b : α) → a = b → Sort v}
+    {b : α} {refl : motive a (Eq.refl a)} {h : a = b} {c : motive b h} :
+    HEq c (@Eq.rec α a motive refl b h) ↔ HEq c refl :=
+  h.rec (fun _ => ⟨id, id⟩) c
+
+/--
+Moves an cast using `Eq.rec` from the function to the argument.
+Note: because the motive isn't reliably detected by unification,
+it needs to be provided as an explicit parameter.
+-/
+theorem apply_eqRec {α : Sort u} {a : α} (motive : (b : α) → a = b → Sort v)
+    {b : α} {h : a = b} {c : motive a (Eq.refl a) → β} {d : motive b h} :
+    @Eq.rec α a (fun b h => motive b h → β) c b h d = c (h.symm ▸ d) := by
+  cases h; rfl
+
 /--
 If casting a term with `Eq.rec` to another type makes it equal to some other term, then the two
 terms are heterogeneously equal.
@@ -983,7 +1011,7 @@ theorem HEq.comm {a : α} {b : β} : HEq a b ↔ HEq b a := Iff.intro HEq.symm H
 theorem heq_comm {a : α} {b : β} : HEq a b ↔ HEq b a := HEq.comm
 
 @[symm] theorem Iff.symm (h : a ↔ b) : b ↔ a := Iff.intro h.mpr h.mp
-theorem Iff.comm: (a ↔ b) ↔ (b ↔ a) := Iff.intro Iff.symm Iff.symm
+theorem Iff.comm : (a ↔ b) ↔ (b ↔ a) := Iff.intro Iff.symm Iff.symm
 theorem iff_comm : (a ↔ b) ↔ (b ↔ a) := Iff.comm
 
 @[symm] theorem And.symm : a ∧ b → b ∧ a := fun ⟨ha, hb⟩ => ⟨hb, ha⟩
@@ -1148,12 +1176,12 @@ theorem dif_eq_if (c : Prop) {h : Decidable c} {α : Sort u} (t : α) (e : α) :
   | isTrue _    => rfl
   | isFalse _   => rfl
 
-instance {c t e : Prop} [dC : Decidable c] [dT : Decidable t] [dE : Decidable e] : Decidable (if c then t else e)  :=
+instance {c t e : Prop} [dC : Decidable c] [dT : Decidable t] [dE : Decidable e] : Decidable (if c then t else e) :=
   match dC with
   | isTrue _   => dT
   | isFalse _  => dE
 
-instance {c : Prop} {t : c → Prop} {e : ¬c → Prop} [dC : Decidable c] [dT : ∀ h, Decidable (t h)] [dE : ∀ h, Decidable (e h)] : Decidable (if h : c then t h else e h)  :=
+instance {c : Prop} {t : c → Prop} {e : ¬c → Prop} [dC : Decidable c] [dT : ∀ h, Decidable (t h)] [dE : ∀ h, Decidable (e h)] : Decidable (if h : c then t h else e h) :=
   match dC with
   | isTrue hc  => dT hc
   | isFalse hc => dE hc
@@ -1869,9 +1897,7 @@ protected abbrev hrecOn
     (f : (a : α) → motive (Quot.mk r a))
     (c : (a b : α) → (p : r a b) → HEq (f a) (f b))
     : motive q :=
-  Quot.recOn q f fun a b p => eq_of_heq <|
-    have p₁ : HEq (Eq.ndrec (f a) (sound p)) (f a) := eqRec_heq (sound p) (f a)
-    HEq.trans p₁ (c a b p)
+  Quot.recOn q f fun a b p => eq_of_heq (eqRec_heq_iff.mpr (c a b p))
 
 end
 end Quot
@@ -2234,6 +2260,27 @@ theorem funext {α : Sort u} {β : α → Sort v} {f g : (x : α) → β x}
   show extfunApp (Quot.mk eqv f) = extfunApp (Quot.mk eqv g)
   exact congrArg extfunApp (Quot.sound h)
 
+/--
+Like `Quot.liftOn q f h` but allows `f a` to "know" that `q = Quot.mk r a`.
+-/
+protected abbrev Quot.pliftOn {α : Sort u} {r : α → α → Prop}
+    (q : Quot r)
+    (f : (a : α) → q = Quot.mk r a → β)
+    (h : ∀ (a b : α) (h h'), r a b → f a h = f b h') : β :=
+  q.rec (motive := fun q' => q = q' → β) f
+    (fun a b p => funext fun h' =>
+      (apply_eqRec (motive := fun b _ => q = b)).trans
+        (@h a b (h'.trans (sound p).symm) h' p)) rfl
+
+/--
+Like `Quotient.liftOn q f h` but allows `f a` to "know" that `q = Quotient.mk s a`.
+-/
+protected abbrev Quotient.pliftOn {α : Sort u} {s : Setoid α}
+    (q : Quotient s)
+    (f : (a : α) → q = Quotient.mk s a → β)
+    (h : ∀ (a b : α) (h h'), a ≈ b → f a h = f b h') : β :=
+  Quot.pliftOn q f h
+
 instance Pi.instSubsingleton {α : Sort u} {β : α → Sort v} [∀ a, Subsingleton (β a)] :
     Subsingleton (∀ a, β a) where
   allEq f g := funext fun a => Subsingleton.elim (f a) (g a)

src/Std/Data.lean

@@ -10,6 +10,9 @@ import Std.Data.HashSet
 import Std.Data.DTreeMap
 import Std.Data.TreeMap
 import Std.Data.TreeSet
+import Std.Data.ExtDHashMap
+import Std.Data.ExtHashMap
+import Std.Data.ExtHashSet
 
 -- The imports above only import the modules needed to work with the version which bundles
 -- the well-formedness invariant, so we need to additionally import the files that deal with the

src/Std/Data/DHashMap/Basic.lean

@@ -18,7 +18,7 @@ Lemmas about the operations on `Std.Data.DHashMap` are available in the
 module `Std.Data.DHashMap.Lemmas`.
 
 See the module `Std.Data.DHashMap.Raw` for a variant of this type which is safe to use in
-nested inductive types.
+nested inductive types and the module `Std.Data.ExtDHashMap` for a variant with extensionality.
 
 For implementation notes, see the docstring of the module `Std.Data.DHashMap.Internal.Defs`.
 -/
@@ -54,9 +54,12 @@ should be an equivalence relation and `a == b` should imply `hash a = hash b` (s
 instance is lawful, i.e., if `a == b` implies `a = b`.
 
 These hash maps contain a bundled well-formedness invariant, which means that they cannot
-be used in nested inductive types. For these use cases, `Std.Data.DHashMap.Raw` and
-`Std.Data.DHashMap.Raw.WF` unbundle the invariant from the hash map. When in doubt, prefer
+be used in nested inductive types. For these use cases, `Std.DHashMap.Raw` and
+`Std.DHashMap.Raw.WF` unbundle the invariant from the hash map. When in doubt, prefer
 `DHashMap` over `DHashMap.Raw`.
+
+For a variant that is more convenient for use in proofs because of extensionalities, see
+`Std.ExtDHashMap` which is defined in the module `Std.Data.ExtDHashMap`.
 -/
 def DHashMap (α : Type u) (β : α → Type v) [BEq α] [Hashable α] := { m : DHashMap.Raw α β // m.WF }
 

src/Std/Data/DHashMap/Internal/RawLemmas.lean

@@ -1178,8 +1178,6 @@ theorem foldRevM_eq_foldrM_toList [Monad m'] [LawfulMonad m']
     {f : δ → α → β → m' δ} {init : δ} :
     Raw.Internal.foldRevM f init m.1 =
       (Raw.Const.toList m.1).foldrM (fun a b => f b a.1 a.2) init := by
-  have :=Raw.foldRevM_eq_foldrM_toListModel (m := m') (b := m.1) (init := init) (f := f)
-
   simp_to_model [foldRevM, Const.toList] using List.foldrM_eq_foldrM_toProd'
 
 theorem foldRev_eq_foldr_toList {f : δ → α → β → δ} {init : δ} :
@@ -1199,6 +1197,16 @@ end Const
 
 end monadic
 
+section insertMany
+
+variable {ρ : Type w} [ForIn Id ρ ((a : α) × β a)]
+
+@[elab_as_elim]
+theorem insertMany_ind {motive : Raw₀ α β → Prop} (m : Raw₀ α β) (l : ρ)
+    (init : motive m) (insert : ∀ m a b, motive m → motive (m.insert a b)) :
+    motive (m.insertMany l).1 :=
+  (m.insertMany l).2 motive (insert _ _ _) init
+
 @[simp]
 theorem insertMany_nil :
     m.insertMany [] = m := by
@@ -1227,6 +1235,14 @@ theorem contains_of_contains_insertMany_list [EquivBEq α] [LawfulHashable α] (
     (m.insertMany l).1.contains k → (l.map Sigma.fst).contains k = false → m.contains k := by
   simp_to_model [insertMany, contains] using List.containsKey_of_containsKey_insertList
 
+theorem contains_insertMany_of_contains [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} {k : α} (h' : m.contains k) : (m.insertMany l).1.contains k := by
+  refine (?_ : _ ∧ (m.insertMany l).1.1.WF).1
+  refine insertMany_ind m l ⟨h', h⟩ ?_
+  intro m a b ⟨h', h⟩
+  simp only [h, contains_insert, h', Bool.or_true, true_and]
+  exact h.insert₀
+
 theorem get?_insertMany_list_of_contains_eq_false [LawfulBEq α] (h : m.1.WF)
     {l : List ((a : α) × β a)} {k : α}
     (h' : (l.map Sigma.fst).contains k = false) :
@@ -1352,6 +1368,15 @@ theorem size_le_size_insertMany_list [EquivBEq α] [LawfulHashable α] (h : m.1.
     m.1.size ≤ (m.insertMany l).1.1.size := by
   simp_to_model [insertMany, size] using List.length_le_length_insertList
 
+theorem size_le_size_insertMany [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} : m.1.size ≤ (m.insertMany l).1.1.size := by
+  refine (?_ : _ ∧ (m.insertMany l).1.1.WF).1
+  refine insertMany_ind m l ⟨Nat.le_refl _, h⟩ ?_
+  intro m' a b ⟨h', h⟩
+  constructor
+  · exact Nat.le_trans h' (size_le_size_insert m' h)
+  · exact h.insert₀
+
 theorem size_insertMany_list_le [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
     {l : List ((a : α) × β a)} :
     (m.insertMany l).1.1.size ≤ m.1.size + l.length := by
@@ -1363,9 +1388,26 @@ theorem isEmpty_insertMany_list [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
     (m.insertMany l).1.1.isEmpty = (m.1.isEmpty && l.isEmpty) := by
   simp_to_model [insertMany, isEmpty] using List.isEmpty_insertList
 
+theorem isEmpty_of_isEmpty_insertMany [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} : (m.insertMany l).1.1.isEmpty → m.1.isEmpty := by
+  refine (?_ : _ ∧ (m.insertMany l).1.1.WF).1
+  refine insertMany_ind m l ⟨id, h⟩ ?_
+  intro m' a b ⟨h', h⟩
+  constructor
+  · intro h''
+    simp only [isEmpty_insert, h, Bool.false_eq_true] at h''
+  · exact h.insert₀
+
 namespace Const
 
 variable {β : Type v} (m : Raw₀ α (fun _ => β))
+variable {ρ : Type w} [ForIn Id ρ (α × β)]
+
+@[elab_as_elim]
+theorem insertMany_ind {motive : Raw₀ α (fun _ => β) → Prop} (m : Raw₀ α fun _ => β) (l : ρ)
+    (init : motive m) (insert : ∀ m a b, motive m → motive (m.insert a b)) :
+    motive (insertMany m l).1 :=
+  (insertMany m l).2 motive (insert _ _ _) init
 
 @[simp]
 theorem insertMany_nil :
@@ -1395,6 +1437,14 @@ theorem contains_of_contains_insertMany_list [EquivBEq α] [LawfulHashable α] (
     (insertMany m l).1.contains k → (l.map Prod.fst).contains k = false → m.contains k := by
   simp_to_model [Const.insertMany, contains] using List.containsKey_of_containsKey_insertListConst
 
+theorem contains_insertMany_of_contains [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} {k : α} (h' : m.contains k) : (insertMany m l).1.contains k := by
+  refine (?_ : _ ∧ (insertMany m l).1.1.WF).1
+  refine insertMany_ind m l ⟨h', h⟩ ?_
+  intro m a b ⟨h', h⟩
+  simp only [h, contains_insert, h', Bool.or_true, true_and]
+  exact h.insert₀
+
 theorem getKey?_insertMany_list_of_contains_eq_false [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
     {l : List (α × β)} {k : α}
     (h' : (l.map Prod.fst).contains k = false) :
@@ -1466,6 +1516,15 @@ theorem size_le_size_insertMany_list [EquivBEq α] [LawfulHashable α] (h : m.1.
     m.1.size ≤ (insertMany m l).1.1.size := by
   simp_to_model [Const.insertMany, size] using List.length_le_length_insertListConst
 
+theorem size_le_size_insertMany [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} : m.1.size ≤ (insertMany m l).1.1.size := by
+  refine (?_ : _ ∧ (insertMany m l).1.1.WF).1
+  refine insertMany_ind m l ⟨Nat.le_refl _, h⟩ ?_
+  intro m' a b ⟨h', h⟩
+  constructor
+  · exact Nat.le_trans h' (size_le_size_insert m' h)
+  · exact h.insert₀
+
 theorem size_insertMany_list_le [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
     {l : List (α × β)} :
     (insertMany m l).1.1.size ≤ m.1.size + l.length := by
@@ -1477,6 +1536,16 @@ theorem isEmpty_insertMany_list [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
     (insertMany m l).1.1.isEmpty = (m.1.isEmpty && l.isEmpty) := by
   simp_to_model [Const.insertMany, isEmpty] using List.isEmpty_insertListConst
 
+theorem isEmpty_of_isEmpty_insertMany [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} : (insertMany m l).1.1.isEmpty → m.1.isEmpty := by
+  refine (?_ : _ ∧ (insertMany m l).1.1.WF).1
+  refine insertMany_ind m l ⟨id, h⟩ ?_
+  intro m' a b ⟨h', h⟩
+  constructor
+  · intro h''
+    simp only [isEmpty_insert, h, Bool.false_eq_true] at h''
+  · exact h.insert₀
+
 theorem get?_insertMany_list_of_contains_eq_false [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
     {l : List (α × β)} {k : α}
     (h' : (l.map Prod.fst).contains k = false) :
@@ -1532,6 +1601,15 @@ theorem getD_insertMany_list_of_mem [EquivBEq α] [LawfulHashable α] (h : m.1.W
 
 variable (m : Raw₀ α (fun _ => Unit))
 
+variable {ρ : Type w} [ForIn Id ρ α]
+
+@[elab_as_elim]
+theorem insertManyIfNewUnit_ind {motive : Raw₀ α (fun _ => Unit) → Prop}
+    (m : Raw₀ α fun _ => Unit) (l : ρ)
+    (init : motive m) (insert : ∀ m a, motive m → motive (m.insertIfNew a ())) :
+    motive (insertManyIfNewUnit m l).1 :=
+  (insertManyIfNewUnit m l).2 motive (insert _ _) init
+
 @[simp]
 theorem insertManyIfNewUnit_nil :
     insertManyIfNewUnit m [] = m := by
@@ -1561,6 +1639,14 @@ theorem contains_of_contains_insertManyIfNewUnit_list [EquivBEq α] [LawfulHasha
   simp_to_model [Const.insertManyIfNewUnit, contains]
     using List.containsKey_of_containsKey_insertListIfNewUnit
 
+theorem contains_insertManyIfNewUnit_of_contains [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} {k : α} (h' : m.contains k) : (insertManyIfNewUnit m l).1.contains k := by
+  refine (?_ : _ ∧ (insertManyIfNewUnit m l).1.1.WF).1
+  refine insertManyIfNewUnit_ind m l ⟨h', h⟩ ?_
+  intro m a ⟨h', h⟩
+  simp only [h, contains_insertIfNew, h', Bool.or_true, true_and]
+  exact h.insertIfNew₀
+
 theorem getKey?_insertManyIfNewUnit_list_of_contains_eq_false_of_contains_eq_false
     [EquivBEq α] [LawfulHashable α]
     (h : m.1.WF) {l : List α} {k : α} :
@@ -1652,6 +1738,15 @@ theorem size_le_size_insertManyIfNewUnit_list [EquivBEq α] [LawfulHashable α]
     m.1.size ≤ (insertManyIfNewUnit m l).1.1.size := by
   simp_to_model [Const.insertManyIfNewUnit, size] using List.length_le_length_insertListIfNewUnit
 
+theorem size_le_size_insertManyIfNewUnit [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} : m.1.size ≤ (insertManyIfNewUnit m l).1.1.size := by
+  refine (?_ : _ ∧ (insertManyIfNewUnit m l).1.1.WF).1
+  refine insertManyIfNewUnit_ind m l ⟨Nat.le_refl _, h⟩ ?_
+  intro m' a ⟨h', h⟩
+  constructor
+  · exact Nat.le_trans h' (size_le_size_insertIfNew m' h)
+  · exact h.insertIfNew₀
+
 theorem size_insertManyIfNewUnit_list_le [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
     {l : List α} :
     (insertManyIfNewUnit m l).1.1.size ≤ m.1.size + l.length := by
@@ -1663,6 +1758,16 @@ theorem isEmpty_insertManyIfNewUnit_list [EquivBEq α] [LawfulHashable α] (h :
     (insertManyIfNewUnit m l).1.1.isEmpty = (m.1.isEmpty && l.isEmpty) := by
   simp_to_model [Const.insertManyIfNewUnit, isEmpty] using List.isEmpty_insertListIfNewUnit
 
+theorem isEmpty_of_isEmpty_insertManyIfNewUnit [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
+    {l : ρ} : (insertManyIfNewUnit m l).1.1.isEmpty → m.1.isEmpty := by
+  refine (?_ : _ ∧ (insertManyIfNewUnit m l).1.1.WF).1
+  refine insertManyIfNewUnit_ind m l ⟨id, h⟩ ?_
+  intro m' a ⟨h', h⟩
+  constructor
+  · intro h''
+    simp only [isEmpty_insertIfNew, h, Bool.false_eq_true] at h''
+  · exact h.insertIfNew₀
+
 theorem get?_insertManyIfNewUnit_list [EquivBEq α] [LawfulHashable α] (h : m.1.WF)
     {l : List α} {k : α} :
     get? (insertManyIfNewUnit m l).1 k =
@@ -2329,6 +2434,8 @@ abbrev getD_insertManyIfNewUnit_empty_list := @getD_insertManyIfNewUnit_emptyWit
 
 end Const
 
+end insertMany
+
 section Alter
 
 theorem isEmpty_alter_eq_isEmpty_erase [LawfulBEq α] (h : m.1.WF) {k : α}
@@ -3069,10 +3176,20 @@ theorem modify_equiv_congr (h₁ : m₁.1.WF) (h₂ : m₂.1.WF) (h : m₁.1 ~m
     {k : α} (f : β → β) : (modify m₁ k f).1 ~m (modify m₂ k f).1 := by
   simp_to_model [Equiv, Const.modify] using List.Const.modifyKey_of_perm _ h.1
 
+theorem equiv_of_forall_getKey_eq_of_forall_get?_eq (h₁ : m₁.1.WF) (h₂ : m₂.1.WF) :
+    (∀ k hk hk', m₁.getKey k hk = m₂.getKey k hk') → (∀ k, get? m₁ k = get? m₂ k) → m₁.1 ~m m₂.1 := by
+  simp_to_model [getKey, Const.get?, contains, Equiv] using List.getKey_getValue?_ext
+
+@[deprecated equiv_of_forall_getKey_eq_of_forall_get?_eq (since := "2025-04-25")]
 theorem equiv_of_forall_getKey?_eq_of_forall_get?_eq (h₁ : m₁.1.WF) (h₂ : m₂.1.WF) :
     (∀ k, m₁.getKey? k = m₂.getKey? k) → (∀ k, get? m₁ k = get? m₂ k) → m₁.1 ~m m₂.1 := by
   simp_to_model [getKey?, Const.get?, Equiv] using List.getKey?_getValue?_ext
 
+theorem equiv_of_forall_get?_eq {α : Type u} [BEq α] [Hashable α] [LawfulBEq α]
+    {m₁ m₂ : Raw₀ α fun _ => β} (h₁ : m₁.1.WF) (h₂ : m₂.1.WF) :
+    (∀ k, get? m₁ k = get? m₂ k) → m₁.1 ~m m₂.1 := by
+  simpa only [Const.get?_eq_get?, h₁, h₂] using Raw₀.equiv_of_forall_get?_eq m₁ m₂ h₁ h₂
+
 theorem equiv_of_forall_getKey?_unit_eq {m₁ m₂ : Raw₀ α fun _ => Unit}
     (h₁ : m₁.1.WF) (h₂ : m₂.1.WF) : (∀ k, m₁.getKey? k = m₂.getKey? k) → m₁.1 ~m m₂.1 := by
   simp_to_model [getKey?, Equiv] using List.getKey?_ext