Skip to content

Commit

Permalink
maps: add All, Keys, Values, Insert, Collect
Browse files Browse the repository at this point in the history
  • Loading branch information
aimuz committed May 20, 2024
1 parent e7bf995 commit 9a8daa4
Show file tree
Hide file tree
Showing 7 changed files with 193 additions and 2 deletions.
5 changes: 5 additions & 0 deletions api/next/61900.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
pkg maps, func All[$0 interface{ ~map[$1]$2 }, $1 comparable, $2 interface{}]($0) iter.Seq2 #61900
pkg maps, func Collect[$0 comparable, $1 interface{}](iter.Seq2) map[$0]$1 #61900
pkg maps, func Insert[$0 interface{ ~map[$1]$2 }, $1 comparable, $2 interface{}]($0, iter.Seq2) #61900
pkg maps, func Keys[$0 interface{ ~map[$1]$2 }, $1 comparable, $2 interface{}]($0) iter.Seq #61900
pkg maps, func Values[$0 interface{ ~map[$1]$2 }, $1 comparable, $2 interface{}]($0) iter.Seq #61900
8 changes: 8 additions & 0 deletions doc/next/6-stdlib/3-iter.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
1 change: 1 addition & 0 deletions doc/next/6-stdlib/99-minor/slices/61900.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
<!-- see ../../3-iter.md -->
2 changes: 1 addition & 1 deletion src/cmd/dist/test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion src/go/build/deps_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,9 @@ var depsRules = `
internal/byteorder, internal/goarch, unsafe < internal/chacha8rand;
unsafe < internal/cpu, maps;
unsafe < internal/cpu;
RUNTIME, unsafe < maps;
# RUNTIME is the core runtime group of packages, all of them very light-weight.
internal/abi,
Expand Down
55 changes: 55 additions & 0 deletions src/maps/iter.go
Original file line number Diff line number Diff line change
@@ -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
}
120 changes: 120 additions & 0 deletions src/maps/iter_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
}

0 comments on commit 9a8daa4

Please sign in to comment.