From 94c79659fac7985e1ae6fce87cb2d708c9d64d1f Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Wed, 9 Oct 2024 08:02:45 -0700
Subject: [PATCH] types/views: add iterators to the three Map view types

Their callers using Range are all kinda clunky feeling. Iterators
should make them more readable.

Updates #12912

Change-Id: I93461eba8e735276fda4a8558a4ae4bfd6c04922
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 types/views/views.go      | 34 +++++++++++++++++++++++++++++++
 types/views/views_test.go | 43 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/types/views/views.go b/types/views/views.go
index 5fe88fa6c..19aa69d4a 100644
--- a/types/views/views.go
+++ b/types/views/views.go
@@ -440,6 +440,17 @@ func (m MapSlice[K, V]) AsMap() map[K][]V {
 	return out
 }
 
+// All returns an iterator iterating over the keys and values of m.
+func (m MapSlice[K, V]) All() iter.Seq2[K, Slice[V]] {
+	return func(yield func(K, Slice[V]) bool) {
+		for k, v := range m.ж {
+			if !yield(k, SliceOf(v)) {
+				return
+			}
+		}
+	}
+}
+
 // Map provides a read-only view of a map. It is the caller's responsibility to
 // make sure V is immutable.
 type Map[K comparable, V any] struct {
@@ -526,6 +537,18 @@ func (m Map[K, V]) Range(f MapRangeFn[K, V]) {
 	}
 }
 
+// All returns an iterator iterating over the keys
+// and values of m.
+func (m Map[K, V]) All() iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
+		for k, v := range m.ж {
+			if !yield(k, v) {
+				return
+			}
+		}
+	}
+}
+
 // MapFnOf returns a MapFn for m.
 func MapFnOf[K comparable, T any, V any](m map[K]T, f func(T) V) MapFn[K, T, V] {
 	return MapFn[K, T, V]{
@@ -587,6 +610,17 @@ func (m MapFn[K, T, V]) Range(f MapRangeFn[K, V]) {
 	}
 }
 
+// All returns an iterator iterating over the keys and value views of m.
+func (m MapFn[K, T, V]) All() iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
+		for k, v := range m.ж {
+			if !yield(k, m.wrapv(v)) {
+				return
+			}
+		}
+	}
+}
+
 // ContainsPointers reports whether T contains any pointers,
 // either explicitly or implicitly.
 // It has special handling for some types that contain pointers
diff --git a/types/views/views_test.go b/types/views/views_test.go
index ec7dcec4c..8a1ff3fdd 100644
--- a/types/views/views_test.go
+++ b/types/views/views_test.go
@@ -446,6 +446,7 @@ func (v testStructView) AsStruct() *testStruct {
 	}
 	return v.p.Clone()
 }
+func (v testStructView) ValueForTest() string { return v.p.value }
 
 func TestSliceViewRange(t *testing.T) {
 	vs := SliceOfViews([]*testStruct{{value: "foo"}, {value: "bar"}})
@@ -458,3 +459,45 @@ func TestSliceViewRange(t *testing.T) {
 		t.Errorf("got %q; want %q", got, want)
 	}
 }
+
+func TestMapIter(t *testing.T) {
+	m := MapOf(map[string]int{"foo": 1, "bar": 2})
+	var got []string
+	for k, v := range m.All() {
+		got = append(got, fmt.Sprintf("%s-%d", k, v))
+	}
+	slices.Sort(got)
+	want := []string{"bar-2", "foo-1"}
+	if !slices.Equal(got, want) {
+		t.Errorf("got %q; want %q", got, want)
+	}
+}
+
+func TestMapSliceIter(t *testing.T) {
+	m := MapSliceOf(map[string][]int{"foo": {3, 4}, "bar": {1, 2}})
+	var got []string
+	for k, v := range m.All() {
+		got = append(got, fmt.Sprintf("%s-%d", k, v))
+	}
+	slices.Sort(got)
+	want := []string{"bar-{[1 2]}", "foo-{[3 4]}"}
+	if !slices.Equal(got, want) {
+		t.Errorf("got %q; want %q", got, want)
+	}
+}
+
+func TestMapFnIter(t *testing.T) {
+	m := MapFnOf[string, *testStruct, testStructView](map[string]*testStruct{
+		"foo": {value: "fooVal"},
+		"bar": {value: "barVal"},
+	}, func(p *testStruct) testStructView { return testStructView{p} })
+	var got []string
+	for k, v := range m.All() {
+		got = append(got, fmt.Sprintf("%v-%v", k, v.ValueForTest()))
+	}
+	slices.Sort(got)
+	want := []string{"bar-barVal", "foo-fooVal"}
+	if !slices.Equal(got, want) {
+		t.Errorf("got %q; want %q", got, want)
+	}
+}