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>
This commit is contained in:
Brad Fitzpatrick 2024-10-09 08:02:45 -07:00 committed by Brad Fitzpatrick
parent f6d4d03355
commit 94c79659fa
2 changed files with 77 additions and 0 deletions

View File

@ -440,6 +440,17 @@ func (m MapSlice[K, V]) AsMap() map[K][]V {
return out 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 // Map provides a read-only view of a map. It is the caller's responsibility to
// make sure V is immutable. // make sure V is immutable.
type Map[K comparable, V any] struct { 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. // 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] { 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]{ 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, // ContainsPointers reports whether T contains any pointers,
// either explicitly or implicitly. // either explicitly or implicitly.
// It has special handling for some types that contain pointers // It has special handling for some types that contain pointers

View File

@ -446,6 +446,7 @@ func (v testStructView) AsStruct() *testStruct {
} }
return v.p.Clone() return v.p.Clone()
} }
func (v testStructView) ValueForTest() string { return v.p.value }
func TestSliceViewRange(t *testing.T) { func TestSliceViewRange(t *testing.T) {
vs := SliceOfViews([]*testStruct{{value: "foo"}, {value: "bar"}}) vs := SliceOfViews([]*testStruct{{value: "foo"}, {value: "bar"}})
@ -458,3 +459,45 @@ func TestSliceViewRange(t *testing.T) {
t.Errorf("got %q; want %q", got, want) 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)
}
}