kamstrup · kamstrup · Nov 30, 2024 · Nov 29, 2024
diff --git a/go.mod b/go.mod
@@ -1,3 +1,3 @@
 module github.com/kamstrup/intmap
 
-go 1.18
+go 1.23
diff --git a/map64.go b/map64.go
@@ -2,6 +2,7 @@
 package intmap
 
 import (
+	"iter"
 	"math"
 )
 
@@ -229,6 +230,58 @@ func (m *Map[K, V]) ForEach(f func(K, V) bool) {
 	forEach64(m.data, f)
 }
 
+// All returns an iterator over key-value pairs from m.
+// The iterator returns immediately if invoked on a nil map.
+//
+// The iteration order of a Map is not defined, so please avoid relying on it.
+func (m *Map[K, V]) All() iter.Seq2[K, V] {
+	return m.ForEach
+}
+
+// Keys returns an iterator over keys in m.
+// The iterator returns immediately if invoked on a nil map.
+//
+// The iteration order of a Map is not defined, so please avoid relying on it.
+func (m *Map[K, V]) Keys() iter.Seq[K] {
+	return func(yield func(k K) bool) {
+		if m == nil {
+			return
+		}
+
+		if m.hasZeroKey && !yield(K(0)) {
+			return
+		}
+
+		for _, p := range m.data {
+			if p.K != K(0) && !yield(p.K) {
+				return
+			}
+		}
+	}
+}
+
+// Values returns an iterator over values in m.
+// The iterator returns immediately if invoked on a nil map.
+//
+// The iteration order of a Map is not defined, so please avoid relying on it.
+func (m *Map[K, V]) Values() iter.Seq[V] {
+	return func(yield func(v V) bool) {
+		if m == nil {
+			return
+		}
+
+		if m.hasZeroKey && !yield(m.zeroVal) {
+			return
+		}
+
+		for _, p := range m.data {
+			if p.K != K(0) && !yield(p.V) {
+				return
+			}
+		}
+	}
+}
+
 // Clear removes all items from the map, but keeps the internal buffers for reuse.
8000func (m *Map[K, V]) Clear() {
 	var zero V

diff --git a/map64_test.go b/map64_test.go
@@ -1,6 +1,8 @@
 package intmap
 
-import "testing"
+import (
+	"testing"
+)
 
 func TestMap64(t *testing.T) {
 	type pairs [][2]int64
@@ -184,3 +186,40 @@ func TestMap64ForEachStop(t *testing.T) {
 		t.Fatalf("unexpected number of elements processed: %d, want %d", have, want)
 	}
 }
+
+func TestMap64Iterators(t *testing.T) {
+	m := New[int, int](10)
+	for i := 0; i < 100; i++ {
+		m.Put(i, 99-i) // 0:99, 1:98, 2:97, ...
+	}
+	const sumTo99 = 99 * (99 + 1) / 2
+
+	sum := 0
+	for k, v := range m.All() {
+		sum += k + v
+	}
+
+	if sum != sumTo99*2 {
+		t.Fatalf("unexpected sum when iterating over keys and values: %d, want %d", sum, sumTo99*2)
+	}
+
+	sum = 0
+	for k := range m.Keys() {
+		if k == 9 {
+			continue
+		}
+		sum += k
+	}
+	expected := sumTo99 - 9
+	if sum != expected {
+		t.Fatalf("unexpected sum when iterating over keys: %d, want %d", sum, expected)
+	}
+
+	sum = 0
+	for v := range m.Values() {
+		sum += v
+	}
+	if sum != sumTo99 {
+		t.Fatalf("unexpected sum when iterating over values: %d, want %d", sum, sumTo99*2)
+	}
+}
diff --git a/set.go b/set.go
@@ -1,5 +1,7 @@
 package intmap
 
+import "iter"
+
 // Set is a specialization of Map modelling a set of integers.
 // Like Map, methods that read from the set are valid on the nil Set.
 // This include Has, Len, and ForEach.
@@ -47,3 +49,11 @@ func (s *Set[K]) ForEach(visit func(k K) bool) {
 		return visit(k)
 	})
 }
+
+// All returns an iterator over keys from the set.
+// The iterator returns immediately if the set is nil.
+//
+// The iteration order of a Set is not defined, so please avoid relying on it.
+func (s *Set[K]) All() iter.Seq[K] {
+	return s.ForEach
+}
diff --git a/set_test.go b/set_test.go
@@ -130,3 +130,20 @@ func TestNilSet(t *testing.T) {
 	}
 
 }
+
+func TestSetIter(t *testing.T) {
+	s := NewSet[int](10)
+	for i := 0; i < 100; i++ {
+		s.Add(i)
+	}
+
+	sum := 0
+	for k := range s.All() {
+		sum += k
+	}
+44BB
+	const sumTo99 = 99 * (99 + 1) / 2
+	if sum != sumTo99 {
+		t.Fatalf("unexpected sum: %d, want %d", sum, sumTo99)
+	}
+}