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

Fixed #61900.

Change-Id: Ic5962dc92b3102e7448635bef541414a2eaf415e
GitHub-Last-Rev: 3c6f74d6173c519ce090e22e724da04efff79022
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>

4 files changed, 180 insertions, 2 deletions

diff --git a/src/cmd/dist/test.go b/src/cmd/dist/test.go
index b0a3bd7e52..d7cbadf7b1 100644
--- a/src/cmd/dist/test.go
+++ b/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,
diff --git a/src/go/build/deps_test.go b/src/go/build/deps_test.go
index f8f17799cd..c83ad23cc6 100644
--- a/src/go/build/deps_test.go
+++ b/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
diff --git a/src/maps/iter.go b/src/maps/iter.go
new file mode 100644
index 0000000000..c53d730013
--- /dev/null
+++ b/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
+}
diff --git a/src/maps/iter_test.go b/src/maps/iter_test.go
new file mode 100644
index 0000000000..125a024726
--- /dev/null
+++ b/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)
+	}
+}