maps: add All, Keys, Values, Insert, Collect

Fixed #61900.

Change-Id: Ic5962dc92b3102e7448635bef541414a2eaf415e
GitHub-Last-Rev: 3c6f74d617
GitHub-Pull-Request: golang/go#67521
Reviewed-on: https://go-review.googlesource.com/c/go/+/586716
Auto-Submit: Ian Lance Taylor <iant@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Commit-Queue: Ian Lance Taylor <iant@google.com>
Reviewed-by: Russ Cox <rsc@golang.org>
Auto-Submit: Russ Cox <rsc@golang.org>
Reviewed-by: Ian Lance Taylor <iant@google.com>

api/next/61900.txt

@@ -0,0 +1,5 @@
+pkg maps, func All[$0 interface{ ~map[$1]$2 }, $1 comparable, $2 interface{}]($0) iter.Seq2[$1, $2] #61900
+pkg maps, func Collect[$0 comparable, $1 interface{}](iter.Seq2[$0, $1]) map[$0]$1 #61900
+pkg maps, func Insert[$0 interface{ ~map[$1]$2 }, $1 comparable, $2 interface{}]($0, iter.Seq2[$1, $2]) #61900
+pkg maps, func Keys[$0 interface{ ~map[$1]$2 }, $1 comparable, $2 interface{}]($0) iter.Seq[$1] #61900
+pkg maps, func Values[$0 interface{ ~map[$1]$2 }, $1 comparable, $2 interface{}]($0) iter.Seq[$2] #61900

doc/next/6-stdlib/3-iter.md

@@ -21,3 +21,11 @@ with iterators:
   but uses a stable sort algorithm.
 - [Chunk](/pkg/slices#Chunk) returns an iterator over consecutive
   sub-slices of up to n elements of a slice.
+
+The [`maps` package](/pkg/maps/) adds several functions that work
+with iterators:
+- [All](/pkg/maps#All) returns an iterator over key-value pairs from m.
+- [Keys](/pkg/maps#Keys) returns an iterator over keys in m.
+- [Values](/pkg/maps#Values) returns an iterator over values in m.
+- [Insert](/pkg/maps#Insert) adds the key-value pairs from seq to m.
+- [Collect](/pkg/maps#Collect) collects key-value pairs from seq into a new map and returns it.

doc/next/6-stdlib/99-minor/maps/61900.md

@@ -0,0 +1 @@
+<!-- see ../../3-iter.md -->

src/cmd/dist/test.go

@@ -713,7 +713,7 @@ func (t *tester) registerTests() {
 
 	// GOEXPERIMENT=rangefunc tests
 	if !t.compileOnly {
-		for _, pkg := range []string{"iter", "slices"} {
+		for _, pkg := range []string{"iter", "slices", "maps"} {
 			t.registerTest("GOEXPERIMENT=rangefunc",
 				&goTest{
 					variant: pkg,

src/go/build/deps_test.go

@@ -55,7 +55,7 @@ var depsRules = `
 
 	internal/byteorder, internal/goarch, unsafe < internal/chacha8rand;
 
-	unsafe < internal/cpu, maps;
+	unsafe < internal/cpu;
 
 	# RUNTIME is the core runtime group of packages, all of them very light-weight.
 	internal/abi,
@@ -89,6 +89,9 @@ var depsRules = `
 	< iter
 	< RUNTIME;
 
+	RUNTIME, unsafe
+	< maps;
+
 	# slices depends on unsafe for overlapping check, cmp for comparison
 	# semantics, and math/bits for # calculating bitlength of numbers.
 	RUNTIME, unsafe, cmp, math/bits

src/maps/iter.go

@@ -0,0 +1,55 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package maps
+
+import "iter"
+
+// All returns an iterator over key-value pairs from m.
+func All[Map ~map[K]V, K comparable, V any](m Map) iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
+		for k, v := range m {
+			if !yield(k, v) {
+				return
+			}
+		}
+	}
+}
+
+// Keys returns an iterator over keys in m.
+func Keys[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[K] {
+	return func(yield func(K) bool) {
+		for k := range m {
+			if !yield(k) {
+				return
+			}
+		}
+	}
+}
+
+// Values returns an iterator over values in m.
+func Values[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[V] {
+	return func(yield func(V) bool) {
+		for _, v := range m {
+			if !yield(v) {
+				return
+			}
+		}
+	}
+}
+
+// Insert adds the key-value pairs from seq to m.
+func Insert[Map ~map[K]V, K comparable, V any](m Map, seq iter.Seq2[K, V]) {
+	for k, v := range seq {
+		m[k] = v
+	}
+}
+
+// Collect collects key-value pairs from seq into a new map
+// and returns it.
+func Collect[K comparable, V any](seq iter.Seq2[K, V]) map[K]V {
+	m := make(map[K]V)
+	Insert(m, seq)
+	return m
+}

src/maps/iter_test.go

@@ -0,0 +1,120 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package maps
+
+import (
+	"slices"
+	"testing"
+)
+
+func TestAll(t *testing.T) {
+	for size := 0; size < 10; size++ {
+		m := make(map[int]int)
+		for i := range size {
+			m[i] = i
+		}
+		cnt := 0
+		for i, v := range All(m) {
+			v1, ok := m[i]
+			if !ok || v != v1 {
+				t.Errorf("at iteration %d got %d, %d want %d, %d", cnt, i, v, i, v1)
+			}
+			cnt++
+		}
+		if cnt != size {
+			t.Errorf("read %d values expected %d", cnt, size)
+		}
+	}
+}
+
+func TestKeys(t *testing.T) {
+	for size := 0; size < 10; size++ {
+		var want []int
+		m := make(map[int]int)
+		for i := range size {
+			m[i] = i
+			want = append(want, i)
+		}
+
+		var got1 []int
+		for k := range Keys(m) {
+			got1 = append(got1, k)
+		}
+		slices.Sort(got1)
+		if !slices.Equal(got1, want) {
+			t.Errorf("Keys(%v) = %v, want %v", m, got1, want)
+		}
+	}
+}
+
+func TestValues(t *testing.T) {
+	for size := 0; size < 10; size++ {
+		var want []int
+		m := make(map[int]int)
+		for i := range size {
+			m[i] = i
+			want = append(want, i)
+		}
+
+		var got1 []int
+		for v := range Values(m) {
+			got1 = append(got1, v)
+		}
+		slices.Sort(got1)
+		if !slices.Equal(got1, want) {
+			t.Errorf("Values(%v) = %v, want %v", m, got1, want)
+		}
+	}
+}
+
+func testSeq(yield func(int, int) bool) {
+	for i := 0; i < 10; i += 2 {
+		if !yield(i, i+1) {
+			return
+		}
+	}
+}
+
+var testSeqResult = map[int]int{
+	0: 1,
+	2: 3,
+	4: 5,
+	6: 7,
+	8: 9,
+}
+
+func TestInsert(t *testing.T) {
+	got := map[int]int{
+		1: 1,
+		2: 1,
+	}
+	Insert(got, testSeq)
+
+	want := map[int]int{
+		1: 1,
+		2: 1,
+	}
+	for i, v := range testSeqResult {
+		want[i] = v
+	}
+
+	if !Equal(got, want) {
+		t.Errorf("got %v, want %v", got, want)
+	}
+}
+
+func TestCollect(t *testing.T) {
+	m := map[int]int{
+		0: 1,
+		2: 3,
+		4: 5,
+		6: 7,
+		8: 9,
+	}
+	got := Collect(All(m))
+	if !Equal(got, m) {
+		t.Errorf("got %v, want %v", got, m)
+	}
+}