From 910b4e8e6a72f147cf957a1b51a066802caf5a9b Mon Sep 17 00:00:00 2001
From: Joe Tsai <joetsai@digital-static.net>
Date: Wed, 9 Oct 2024 10:28:12 -0700
Subject: [PATCH] syncs: add iterators to Map (#13739)

Add Keys, Values, and All to iterate over
all keys, values, and entries, respectively.

Updates #11038

Signed-off-by: Joe Tsai <joetsai@digital-static.net>
---
 syncs/syncs.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/syncs/syncs.go b/syncs/syncs.go
index 0d40204d2..bfb7c1e04 100644
--- a/syncs/syncs.go
+++ b/syncs/syncs.go
@@ -6,6 +6,7 @@
 
 import (
 	"context"
+	"iter"
 	"sync"
 	"sync/atomic"
 
@@ -256,6 +257,8 @@ func (m *Map[K, V]) Delete(key K) {
 // Iteration stops if f returns false. Map changes are blocked during iteration.
 // A read lock is held for the entire duration of the iteration.
 // Use the [WithLock] method instead to mutate the map during iteration.
+//
+// Deprecated: Use [All], [Keys], or [Values] instead.
 func (m *Map[K, V]) Range(f func(key K, value V) bool) {
 	m.mu.RLock()
 	defer m.mu.RUnlock()
@@ -266,6 +269,51 @@ func (m *Map[K, V]) Range(f func(key K, value V) bool) {
 	}
 }
 
+// Keys iterates over all keys in the map in an undefined order.
+// A read lock is held for the entire duration of the iteration.
+// Use the [WithLock] method instead to mutate the map during iteration.
+func (m *Map[K, V]) Keys() iter.Seq[K] {
+	return func(yield func(K) bool) {
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		for k := range m.m {
+			if !yield(k) {
+				return
+			}
+		}
+	}
+}
+
+// Values iterates over all values in the map in an undefined order.
+// A read lock is held for the entire duration of the iteration.
+// Use the [WithLock] method instead to mutate the map during iteration.
+func (m *Map[K, V]) Values() iter.Seq[V] {
+	return func(yield func(V) bool) {
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		for _, v := range m.m {
+			if !yield(v) {
+				return
+			}
+		}
+	}
+}
+
+// All iterates over all entries in the map in an undefined order.
+// A read lock is held for the entire duration of the iteration.
+// Use the [WithLock] method instead to mutate the map during iteration.
+func (m *Map[K, V]) All() iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		for k, v := range m.m {
+			if !yield(k, v) {
+				return
+			}
+		}
+	}
+}
+
 // WithLock calls f with the underlying map.
 // Use of m2 must not escape the duration of this call.
 // The write-lock is held for the entire duration of this call.