Skip to content

Commit 7632cef

Browse files
datokratdatokrat
authored
feat: hash map iterators (#10761)
This PR provides iterators on hash maps.
1 parent 7a47bfa commit 7632cef

File tree

13 files changed

+785
-2
lines changed

13 files changed

+785
-2
lines changed

‎src/Std/Data/DHashMap.lean‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,7 @@ module
77

88
prelude
99
public import Std.Data.DHashMap.Basic
10-
public import Std.Data.DHashMap.Lemmas
1110
public import Std.Data.DHashMap.AdditionalOperations
11+
public import Std.Data.DHashMap.Iterator
12+
public import Std.Data.DHashMap.Lemmas
13+
public import Std.Data.DHashMap.IteratorLemmas

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/-
2+
Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
3+
Released under Apache 2.0 license as described in the file LICENSE.
4+
Authors: Paul Reichert
5+
-/
6+
module
7+
8+
prelude
9+
import Init.Data.Nat.Lemmas
10+
11+
public import Init.Data.Iterators.Consumers
12+
import Init.Data.Iterators.Internal.Termination
13+
14+
public import Std.Data.DHashMap.Internal.AssocList.Basic
15+
16+
/-!
17+
# Iterators on associative lists
18+
-/
19+
20+
namespace Std.DHashMap.Internal.AssocList
21+
22+
open Std.Iterators
23+
24+
/-- Internal implementation detail of the hash map -/
25+
@[ext, unbox]
26+
public structure AssocListIterator (α : Type u) (β : α → Type v) where
27+
l : AssocList α β
28+
29+
public instance : Iterator (α := AssocListIterator α β) Id ((a : α) × β a) where
30+
IsPlausibleStep it
31+
| .yield it' out => it.internalState.l = .cons out.1 out.2 it'.internalState.l
32+
| .skip _ => False
33+
| .done => it.internalState.l = .nil
34+
step it := pure (match it with
35+
| ⟨⟨.nil⟩⟩ => .deflate ⟨.done, rfl⟩
36+
| ⟨⟨.cons k v l⟩⟩ => .deflate ⟨.yield (toIterM ⟨l⟩ Id _) ⟨k, v⟩, rfl⟩)
37+
38+
def AssocListIterator.finitenessRelation :
39+
FinitenessRelation (AssocListIterator α β) Id where
40+
rel := InvImage WellFoundedRelation.rel (AssocListIterator.l ∘ IterM.internalState)
41+
wf := InvImage.wf _ WellFoundedRelation.wf
42+
subrelation {it it'} h := by
43+
simp_wf
44+
obtain ⟨step, h, h'⟩ := h
45+
cases step <;> simp_all [IterStep.successor, IterM.IsPlausibleStep, Iterator.IsPlausibleStep]
46+
47+
public instance : Finite (AssocListIterator α β) Id :=
48+
Finite.of_finitenessRelation AssocListIterator.finitenessRelation
49+
50+
public instance {α : Type u} {β : α → Type v} {m : Type (max u v) → Type w''} [Monad m] :
51+
IteratorCollect (AssocListIterator α β) Id m :=
52+
.defaultImplementation
53+
54+
public instance {α : Type u} {β : α → Type v} {m : Type (max u v) → Type w''} [Monad m] :
55+
IteratorCollectPartial (AssocListIterator α β) Id m :=
56+
.defaultImplementation
57+
58+
public instance {α : Type u} {β : α → Type v} {m : Type (max u v) → Type w''} [Monad m] :
59+
IteratorLoop (AssocListIterator α β) Id m :=
60+
.defaultImplementation
61+
62+
public instance {α : Type u} {β : α → Type v} {m : Type (max u v) → Type w''} [Monad m] :
63+
IteratorLoopPartial (AssocListIterator α β) Id m :=
64+
.defaultImplementation
65+
66+
public instance : IteratorSize (AssocListIterator α β) Id :=
67+
.defaultImplementation
68+
69+
public instance : IteratorSizePartial (AssocListIterator α β) Id :=
70+
.defaultImplementation
71+
72+
/--
73+
Internal implementation detail of the hash map. Returns a finite iterator on an associative list.
74+
-/
75+
@[expose]
76+
public def iter {α : Type u} {β : α → Type v} (l : AssocList α β) :
77+
Iter (α := AssocListIterator α β) ((a : α) × β a) :=
78+
⟨⟨l⟩⟩
79+
80+
end Std.DHashMap.Internal.AssocList

‎src/Std/Data/DHashMap/Iterator.lean‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/-
2+
Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
3+
Released under Apache 2.0 license as described in the file LICENSE.
4+
Authors: Paul Reichert
5+
-/
6+
module
7+
8+
prelude
9+
public import Std.Data.Iterators.Producers.Array
10+
public import Init.Data.Iterators.Combinators.FlatMap
11+
public import Std.Data.DHashMap.Basic
12+
public import Std.Data.DHashMap.Internal.AssocList.Iterator
13+
14+
/-!
15+
# Iterators on `DHashMap` and `DHashMap.Raw`
16+
-/
17+
18+
namespace Std.DHashMap.Raw
19+
20+
/--
21+
Returns a finite iterator over the entries of a dependent hash map.
22+
The iterator yields the elements of the map in order and then terminates.
23+
24+
**Termination properties:**
25+
26+
* `Finite` instance: always
27+
* `Productive` instance: always
28+
-/
29+
@[inline]
30+
public def iter {α : Type u} {β : α → Type v} (m : Raw α β) :=
31+
(m.buckets.iter.flatMap fun b => b.iter : Iter ((a : α) × β a))
32+
33+
/--
34+
Returns a finite iterator over the keys of a dependent hash map.
35+
The iterator yields the keys in order and then terminates.
36+
37+
The key and value types must live in the same universe.
38+
39+
**Termination properties:**
40+
41+
* `Finite` instance: always
42+
* `Productive` instance: always
43+
-/
44+
@[inline]
45+
public def keysIter {α : Type u} {β : α → Type u} (m : Raw α β) :=
46+
(m.iter.map fun e => e.1 : Iter α)
47+
48+
/--
49+
Returns a finite iterator over the values of a hash map.
50+
The iterator yields the values in order and then terminates.
51+
52+
The key and value types must live in the same universe.
53+
54+
**Termination properties:**
55+
56+
* `Finite` instance: always
57+
* `Productive` instance: always
58+
-/
59+
@[inline]
60+
public def valuesIter {α : Type u} {β : Type u} (m : Raw α (fun _ => β)) :=
61+
(m.iter.map fun e => e.2 : Iter β)
62+
63+
end Std.DHashMap.Raw
64+
65+
namespace Std.DHashMap
66+
67+
@[inline, inherit_doc Raw.iter]
68+
public def iter {α : Type u} {β : α → Type v} [BEq α] [Hashable α] (m : DHashMap α β) :=
69+
(m.1.iter : Iter ((a : α) × β a))
70+
71+
@[inline, inherit_doc Raw.keysIter]
72+
public def keysIter {α : Type u} {β : α → Type u} [BEq α] [Hashable α] (m : DHashMap α β) :=
73+
(m.1.keysIter : Iter α)
74+
75+
@[inline, inherit_doc Raw.valuesIter]
76+
public def valuesIter {α : Type u} {β : Type u} [BEq α] [Hashable α]
77+
(m : DHashMap α (fun _ => β)) :=
78+
(m.iter.map fun e => e.2 : Iter β)
79+
80+
end Std.DHashMap

‎src/Std/Data/DHashMap/IteratorLemmas.lean‎

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
1+
/-
2+
Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
3+
Released under Apache 2.0 license as described in the file LICENSE.
4+
Authors: Paul Reichert
5+
-/
6+
module
7+
8+
prelude
9+
public import Std.Data.Iterators
10+
public import Std.Data.DHashMap.Iterator
11+
import all Std.Data.DHashMap.Basic
12+
import all Std.Data.DHashMap.Raw
13+
import all Std.Data.DHashMap.Iterator
14+
import Init.Data.Iterators.Lemmas.Combinators
15+
import Std.Data.DHashMap.Internal.WF
16+
import all Std.Data.DHashMap.Internal.Defs
17+
import Std.Data.DHashMap.RawLemmas
18+
import Std.Data.DHashMap.Lemmas
19+
20+
namespace Std.DHashMap.Internal.AssocList
21+
open Std.Iterators
22+
23+
@[simp]
24+
public theorem step_iter_nil {α : Type u} {β : α → Type v} :
25+
((.nil : AssocList α β).iter).step = ⟨.done, rfl⟩ := by
26+
simp [Iter.step, IterM.step, Iterator.step, Iter.toIterM, iter]
27+
28+
@[simp]
29+
public theorem step_iter_cons {α : Type u} {β : α → Type v} {k v} {l : AssocList α β} :
30+
((AssocList.cons k v l).iter).step = ⟨.yield l.iter ⟨k, v⟩, rfl⟩ := by
31+
simp [Iter.step, IterM.step, Iterator.step, Iter.toIterM, iter, toIterM, IterM.toIter]
32+
33+
@[simp]
34+
public theorem toList_iter {α : Type u} {β : α → Type v} {l : AssocList α β} :
35+
l.iter.toList = l.toList := by
36+
induction l
37+
· simp [Iter.toList_eq_match_step, step_iter_nil]
38+
· rw [Iter.toList_eq_match_step, step_iter_cons]
39+
simp [*]
40+
41+
end Std.DHashMap.Internal.AssocList
42+
43+
namespace Std.DHashMap.Raw
44+
open Std.Iterators
45+
46+
section EntriesIter
47+
48+
variable {α : Type u} {β : α → Type v} {m : Raw α β}
49+
50+
@[simp]
51+
public theorem toList_iter :
52+
m.iter.toList = m.toList := by
53+
simp [Raw.iter, Iter.toList_flatMap, Iter.toList_map, Internal.toListModel, List.flatMap,
54+
Internal.Raw.toList_eq_toListModel]
55+
56+
@[simp]
57+
public theorem toListRev_iter :
58+
m.iter.toListRev = m.toList.reverse := by
59+
simp [Iter.toListRev_eq, toList_iter]
60+
61+
@[simp]
62+
public theorem toArray_iter [BEq α] [Hashable α] (h : m.WF) :
63+
m.iter.toArray = m.toArray := by
64+
simp [← Iter.toArray_toList, ← Raw.toArray_toList h, toList_iter]
65+
66+
end EntriesIter
67+
68+
section KeysIter
69+
70+
variable {α : Type u} {β : α → Type u} {m : Raw α β}
71+
72+
@[simp]
73+
public theorem toList_keysIter [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (h : m.WF) :
74+
m.keysIter.toList = m.keys := by
75+
simp [keysIter, ← map_fst_toList_eq_keys h, toList_iter]
76+
77+
@[simp]
78+
public theorem toListRev_keysIter [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (h : m.WF) :
79+
m.keysIter.toListRev = m.keys.reverse := by
80+
simp [Iter.toListRev_eq, toList_keysIter h]
81+
82+
@[simp]
83+
public theorem toArray_keysIter [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (h : m.WF) :
84+
m.keysIter.toArray = m.keysArray := by
85+
simp [← Iter.toArray_toList, ← Raw.toArray_keys h, toList_keysIter h]
86+
87+
end KeysIter
88+
89+
section ValuesIter
90+
91+
variable {α β : Type u} {m : Raw α (fun _ => β)}
92+
93+
@[simp]
94+
public theorem toList_valuesIter_eq_toList_map_snd :
95+
m.valuesIter.toList = m.toList.map Sigma.snd := by
96+
simp [valuesIter, toList_iter]
97+
98+
@[simp]
99+
public theorem toListRev_valuesIter :
100+
m.valuesIter.toListRev = (m.toList.map Sigma.snd).reverse := by
101+
simp [Iter.toListRev_eq, toList_valuesIter_eq_toList_map_snd]
102+
103+
@[simp]
104+
public theorem toArray_valuesIter :
105+
m.valuesIter.toArray = (m.toList.map Sigma.snd).toArray := by
106+
simp [← Iter.toArray_toList, toList_valuesIter_eq_toList_map_snd]
107+
108+
end ValuesIter
109+
110+
end Std.DHashMap.Raw
111+
112+
namespace Std.DHashMap
113+
open Std.Iterators
114+
115+
section EntriesIter
116+
117+
variable {α : Type u} {β : α → Type v} [BEq α] [Hashable α] {m : DHashMap α β}
118+
119+
theorem toList_inner :
120+
m.inner.toList = m.toList :=
121+
rfl
122+
123+
@[simp]
124+
public theorem toList_iter :
125+
m.iter.toList = m.toList := by
126+
simp [iter, Raw.toList_iter, toList_inner]
127+
128+
@[simp]
129+
public theorem toListRev_iter :
130+
m.iter.toListRev = m.toList.reverse := by
131+
simp [Iter.toListRev_eq, toList_iter]
132+
133+
@[simp]
134+
public theorem toArray_iter :
135+
m.iter.toArray = m.toArray := by
136+
simp [← Iter.toArray_toList, ← toArray_toList, toList_iter]
137+
138+
end EntriesIter
139+
140+
section KeysIter
141+
142+
variable {α : Type u} {β : α → Type u} [BEq α] [Hashable α] {m : DHashMap α β}
143+
144+
theorem keys_inner :
145+
m.inner.keys = m.keys :=
146+
rfl
147+
148+
@[simp]
149+
public theorem toList_keysIter [EquivBEq α] [LawfulHashable α] :
150+
m.keysIter.toList = m.keys := by
151+
simp [keysIter, Raw.toList_keysIter m.wf, keys_inner]
152+
153+
@[simp]
154+
public theorem toListRev_keysIter [EquivBEq α] [LawfulHashable α] :
155+
m.keysIter.toListRev = m.keys.reverse := by
156+
simp [Iter.toListRev_eq, toList_keysIter]
157+
158+
@[simp]
159+
public theorem toArray_keysIter [EquivBEq α] [LawfulHashable α] :
160+
m.keysIter.toArray = m.keysArray := by
161+
simp [← Iter.toArray_toList, ← toArray_keys, toList_keysIter]
162+
163+
end KeysIter
164+
165+
section ValuesIter
166+
167+
variable {α : Type u} {β : Type u} [BEq α] [Hashable α] {m : DHashMap α (fun _ => β)}
168+
169+
@[simp]
170+
public theorem toList_valuesIter_eq_toList_map_snd :
171+
m.valuesIter.toList = m.toList.map Sigma.snd := by
172+
simp [valuesIter, toList_iter]
173+
174+
@[simp]
175+
public theorem toListRev_valuesIter :
176+
m.valuesIter.toListRev = (m.toList.map Sigma.snd).reverse := by
177+
simp [Iter.toListRev_eq, toList_valuesIter_eq_toList_map_snd]
178+
179+
@[simp]
180+
public theorem toArray_valuesIter :
181+
m.valuesIter.toArray = (m.toList.map Sigma.snd).toArray := by
182+
simp [← Iter.toArray_toList, toList_valuesIter_eq_toList_map_snd]
183+
184+
end ValuesIter
185+
186+
end Std.DHashMap

‎src/Std/Data/HashMap.lean‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,7 @@ module
77

88
prelude
99
public import Std.Data.HashMap.Basic
10-
public import Std.Data.HashMap.Lemmas
1110
public import Std.Data.HashMap.AdditionalOperations
11+
public import Std.Data.HashMap.Iterator
12+
public import Std.Data.HashMap.Lemmas
13+
public import Std.Data.HashMap.IteratorLemmas

0 commit comments

Comments
 (0)