feat: hash map iterators

src/Std/Data/DHashMap.lean

@@ -7,5 +7,7 @@ module
 
 prelude
 public import Std.Data.DHashMap.Basic
-public import Std.Data.DHashMap.Lemmas
 public import Std.Data.DHashMap.AdditionalOperations
+public import Std.Data.DHashMap.Iterator
+public import Std.Data.DHashMap.Lemmas
+public import Std.Data.DHashMap.IteratorLemmas

src/Std/Data/DHashMap/Internal/AssocList/Iterator.lean

@@ -0,0 +1,80 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+module
+
+prelude
+import Init.Data.Nat.Lemmas
+
+public import Init.Data.Iterators.Consumers
+import Init.Data.Iterators.Internal.Termination
+
+public import Std.Data.DHashMap.Internal.AssocList.Basic
+
+/-!
+# Iterators on associative lists
+-/
+
+namespace Std.DHashMap.Internal.AssocList
+
+open Std.Iterators
+
+/-- Internal implementation detail of the hash map -/
+@[ext, unbox]
+public structure AssocListIterator (α : Type u) (β : α → Type v) where
+  l : AssocList α β
+
+public instance : Iterator (α := AssocListIterator α β) Id ((a : α) × β a) where
+  IsPlausibleStep it
+    | .yield it' out => it.internalState.l = .cons out.1 out.2 it'.internalState.l
+    | .skip _ => False
+    | .done => it.internalState.l = .nil
+  step it := pure (match it with
+        | ⟨⟨.nil⟩⟩ => .deflate ⟨.done, rfl⟩
+        | ⟨⟨.cons k v l⟩⟩ => .deflate ⟨.yield (toIterM ⟨l⟩ Id _) ⟨k, v⟩, rfl⟩)
+
+def AssocListIterator.finitenessRelation :
+    FinitenessRelation (AssocListIterator α β) Id where
+  rel := InvImage WellFoundedRelation.rel (AssocListIterator.l ∘ IterM.internalState)
+  wf := InvImage.wf _ WellFoundedRelation.wf
+  subrelation {it it'} h := by
+    simp_wf
+    obtain ⟨step, h, h'⟩ := h
+    cases step <;> simp_all [IterStep.successor, IterM.IsPlausibleStep, Iterator.IsPlausibleStep]
+
+public instance : Finite (AssocListIterator α β) Id :=
+  Finite.of_finitenessRelation AssocListIterator.finitenessRelation
+
+public instance {α : Type u} {β : α → Type v} {m : Type (max u v) → Type w''} [Monad m] :
+    IteratorCollect (AssocListIterator α β) Id m :=
+  .defaultImplementation
+
+public instance {α : Type u} {β : α → Type v} {m : Type (max u v) → Type w''} [Monad m] :
+    IteratorCollectPartial (AssocListIterator α β) Id m :=
+  .defaultImplementation
+
+public instance {α : Type u} {β : α → Type v} {m : Type (max u v) → Type w''} [Monad m] :
+    IteratorLoop (AssocListIterator α β) Id m :=
+  .defaultImplementation
+
+public instance {α : Type u} {β : α → Type v} {m : Type (max u v) → Type w''} [Monad m] :
+    IteratorLoopPartial (AssocListIterator α β) Id m :=
+  .defaultImplementation
+
+public instance : IteratorSize (AssocListIterator α β) Id :=
+  .defaultImplementation
+
+public instance : IteratorSizePartial (AssocListIterator α β) Id :=
+  .defaultImplementation
+
+/--
+Internal implementation detail of the hash map. Returns a finite iterator on an associative list.
+-/
+@[expose]
+public def iter {α : Type u} {β : α → Type v} (l : AssocList α β) :
+    Iter (α := AssocListIterator α β) ((a : α) × β a) :=
+  ⟨⟨l⟩⟩
+
+end Std.DHashMap.Internal.AssocList

src/Std/Data/DHashMap/Iterator.lean

@@ -0,0 +1,80 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+module
+
+prelude
+public import Std.Data.Iterators.Producers.Array
+public import Init.Data.Iterators.Combinators.FlatMap
+public import Std.Data.DHashMap.Basic
+public import Std.Data.DHashMap.Internal.AssocList.Iterator
+
+/-!
+# Iterators on `DHashMap` and `DHashMap.Raw`
+-/
+
+namespace Std.DHashMap.Raw
+
+/--
+Returns a finite iterator over the entries of a dependent hash map.
+The iterator yields the elements of the map in order and then terminates.
+
+**Termination properties:**
+
+* `Finite` instance: always
+* `Productive` instance: always
+-/
+@[inline]
+public def iter {α : Type u} {β : α → Type v} (m : Raw α β) :=
+  (m.buckets.iter.flatMap fun b => b.iter : Iter ((a : α) × β a))
+
+/--
+Returns a finite iterator over the keys of a dependent hash map.
+The iterator yields the keys in order and then terminates.
+
+The key and value types must live in the same universe.
+
+**Termination properties:**
+
+* `Finite` instance: always
+* `Productive` instance: always
+-/
+@[inline]
+public def keysIter {α : Type u} {β : α → Type u} (m : Raw α β) :=
+  (m.iter.map fun e => e.1 : Iter α)
+
+/--
+Returns a finite iterator over the values of a hash map.
+The iterator yields the values in order and then terminates.
+
+The key and value types must live in the same universe.
+
+**Termination properties:**
+
+* `Finite` instance: always
+* `Productive` instance: always
+-/
+@[inline]
+public def valuesIter {α : Type u} {β : Type u} (m : Raw α (fun _ => β)) :=
+  (m.iter.map fun e => e.2 : Iter β)
+
+end Std.DHashMap.Raw
+
+namespace Std.DHashMap
+
+@[inline, inherit_doc Raw.iter]
+public def iter {α : Type u} {β : α → Type v} [BEq α] [Hashable α] (m : DHashMap α β) :=
+  (m.1.iter : Iter ((a : α) × β a))
+
+@[inline, inherit_doc Raw.keysIter]
+public def keysIter {α : Type u} {β : α → Type u} [BEq α] [Hashable α] (m : DHashMap α β) :=
+  (m.1.keysIter : Iter α)
+
+@[inline, inherit_doc Raw.valuesIter]
+public def valuesIter {α : Type u} {β : Type u} [BEq α] [Hashable α]
+    (m : DHashMap α (fun _ => β)) :=
+  (m.iter.map fun e => e.2 : Iter β)
+
+end Std.DHashMap

src/Std/Data/DHashMap/IteratorLemmas.lean

@@ -0,0 +1,186 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+module
+
+prelude
+public import Std.Data.Iterators
+public import Std.Data.DHashMap.Iterator
+import all Std.Data.DHashMap.Basic
+import all Std.Data.DHashMap.Raw
+import all Std.Data.DHashMap.Iterator
+import Init.Data.Iterators.Lemmas.Combinators
+import Std.Data.DHashMap.Internal.WF
+import all Std.Data.DHashMap.Internal.Defs
+import Std.Data.DHashMap.RawLemmas
+import Std.Data.DHashMap.Lemmas
+
+namespace Std.DHashMap.Internal.AssocList
+open Std.Iterators
+
+@[simp]
+public theorem step_iter_nil {α : Type u} {β : α → Type v} :
+    ((.nil : AssocList α β).iter).step = ⟨.done, rfl⟩ := by
+  simp [Iter.step, IterM.step, Iterator.step, Iter.toIterM, iter]
+
+@[simp]
+public theorem step_iter_cons {α : Type u} {β : α → Type v} {k v} {l : AssocList α β} :
+    ((AssocList.cons k v l).iter).step = ⟨.yield l.iter ⟨k, v⟩, rfl⟩ := by
+  simp [Iter.step, IterM.step, Iterator.step, Iter.toIterM, iter, toIterM, IterM.toIter]
+
+@[simp]
+public theorem toList_iter {α : Type u} {β : α → Type v} {l : AssocList α β} :
+    l.iter.toList = l.toList := by
+  induction l
+  · simp [Iter.toList_eq_match_step, step_iter_nil]
+  · rw [Iter.toList_eq_match_step, step_iter_cons]
+    simp [*]
+
+end Std.DHashMap.Internal.AssocList
+
+namespace Std.DHashMap.Raw
+open Std.Iterators
+
+section EntriesIter
+
+variable {α : Type u} {β : α → Type v} {m : Raw α β}
+
+@[simp]
+public theorem toList_iter :
+    m.iter.toList = m.toList := by
+  simp [Raw.iter, Iter.toList_flatMap, Iter.toList_map, Internal.toListModel, List.flatMap,
+    Internal.Raw.toList_eq_toListModel]
+
+@[simp]
+public theorem toListRev_iter :
+    m.iter.toListRev = m.toList.reverse := by
+  simp [Iter.toListRev_eq, toList_iter]
+
+@[simp]
+public theorem toArray_iter [BEq α] [Hashable α] (h : m.WF) :
+    m.iter.toArray = m.toArray := by
+  simp [← Iter.toArray_toList, ← Raw.toArray_toList h, toList_iter]
+
+end EntriesIter
+
+section KeysIter
+
+variable {α : Type u} {β : α → Type u} {m : Raw α β}
+
+@[simp]
+public theorem toList_keysIter [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (h : m.WF) :
+    m.keysIter.toList = m.keys := by
+  simp [keysIter, ← map_fst_toList_eq_keys h, toList_iter]
+
+@[simp]
+public theorem toListRev_keysIter [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (h : m.WF) :
+    m.keysIter.toListRev = m.keys.reverse := by
+  simp [Iter.toListRev_eq, toList_keysIter h]
+
+@[simp]
+public theorem toArray_keysIter [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (h : m.WF) :
+    m.keysIter.toArray = m.keysArray := by
+  simp [← Iter.toArray_toList, ← Raw.toArray_keys h, toList_keysIter h]
+
+end KeysIter
+
+section ValuesIter
+
+variable {α β : Type u} {m : Raw α (fun _ => β)}
+
+@[simp]
+public theorem toList_valuesIter_eq_toList_map_snd :
+    m.valuesIter.toList = m.toList.map Sigma.snd := by
+  simp [valuesIter, toList_iter]
+
+@[simp]
+public theorem toListRev_valuesIter :
+    m.valuesIter.toListRev = (m.toList.map Sigma.snd).reverse := by
+  simp [Iter.toListRev_eq, toList_valuesIter_eq_toList_map_snd]
+
+@[simp]
+public theorem toArray_valuesIter :
+    m.valuesIter.toArray = (m.toList.map Sigma.snd).toArray := by
+  simp [← Iter.toArray_toList, toList_valuesIter_eq_toList_map_snd]
+
+end ValuesIter
+
+end Std.DHashMap.Raw
+
+namespace Std.DHashMap
+open Std.Iterators
+
+section EntriesIter
+
+variable {α : Type u} {β : α → Type v} [BEq α] [Hashable α] {m : DHashMap α β}
+
+theorem toList_inner :
+    m.inner.toList = m.toList :=
+  rfl
+
+@[simp]
+public theorem toList_iter :
+    m.iter.toList = m.toList := by
+  simp [iter, Raw.toList_iter, toList_inner]
+
+@[simp]
+public theorem toListRev_iter :
+    m.iter.toListRev = m.toList.reverse := by
+  simp [Iter.toListRev_eq, toList_iter]
+
+@[simp]
+public theorem toArray_iter :
+    m.iter.toArray = m.toArray := by
+  simp [← Iter.toArray_toList, ← toArray_toList, toList_iter]
+
+end EntriesIter
+
+section KeysIter
+
+variable {α : Type u} {β : α → Type u} [BEq α] [Hashable α] {m : DHashMap α β}
+
+theorem keys_inner :
+    m.inner.keys = m.keys :=
+  rfl
+
+@[simp]
+public theorem toList_keysIter [EquivBEq α] [LawfulHashable α] :
+    m.keysIter.toList = m.keys := by
+  simp [keysIter, Raw.toList_keysIter m.wf, keys_inner]
+
+@[simp]
+public theorem toListRev_keysIter [EquivBEq α] [LawfulHashable α] :
+    m.keysIter.toListRev = m.keys.reverse := by
+  simp [Iter.toListRev_eq, toList_keysIter]
+
+@[simp]
+public theorem toArray_keysIter [EquivBEq α] [LawfulHashable α] :
+    m.keysIter.toArray = m.keysArray := by
+  simp [← Iter.toArray_toList, ← toArray_keys, toList_keysIter]
+
+end KeysIter
+
+section ValuesIter
+
+variable {α : Type u} {β : Type u} [BEq α] [Hashable α] {m : DHashMap α (fun _ => β)}
+
+@[simp]
+public theorem toList_valuesIter_eq_toList_map_snd :
+    m.valuesIter.toList = m.toList.map Sigma.snd := by
+  simp [valuesIter, toList_iter]
+
+@[simp]
+public theorem toListRev_valuesIter :
+    m.valuesIter.toListRev = (m.toList.map Sigma.snd).reverse := by
+  simp [Iter.toListRev_eq, toList_valuesIter_eq_toList_map_snd]
+
+@[simp]
+public theorem toArray_valuesIter :
+    m.valuesIter.toArray = (m.toList.map Sigma.snd).toArray := by
+  simp [← Iter.toArray_toList, toList_valuesIter_eq_toList_map_snd]
+
+end ValuesIter
+
+end Std.DHashMap

src/Std/Data/HashMap.lean

@@ -7,5 +7,7 @@ module
 
 prelude
 public import Std.Data.HashMap.Basic
-public import Std.Data.HashMap.Lemmas
 public import Std.Data.HashMap.AdditionalOperations
+public import Std.Data.HashMap.Iterator
+public import Std.Data.HashMap.Lemmas
+public import Std.Data.HashMap.IteratorLemmas