maps,runtime: improve maps.Keys

name     old time/op    new time/op    delta
Keys-10    8.65ms ± 0%    7.06ms ± 0%  -18.41%  (p=0.000 n=9+9)

name     old alloc/op   new alloc/op   delta
Keys-10    58.2kB ± 1%    47.4kB ± 2%  -18.54%  (p=0.000 n=10+10)

name     old allocs/op  new allocs/op  delta
Keys-10      0.00           0.00          ~     (all equal)

Change-Id: Ia7061c37be89e906e79bc3ef3bb1ef0d7913f9f6
Reviewed-on: https://go-review.googlesource.com/c/go/+/481435
Reviewed-by: Keith Randall <khr@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Heschi Kreinick <heschi@google.com>
Run-TryBot: xie cui <523516579@qq.com>

2 files changed, 44 insertions, 3 deletions

diff --git a/src/maps/maps.go b/src/maps/maps.go
index 27eea01501..dddfb37973 100644
--- a/src/maps/maps.go
+++ b/src/maps/maps.go
@@ -5,16 +5,25 @@
 // Package maps defines various functions useful with maps of any type.
 package maps
 
+import "unsafe"
+
+// keys is implemented in the runtime package.
+//
+//go:noescape
+func keys(m any, slice unsafe.Pointer)
+
 // Keys returns the keys of the map m.
 // The keys will be in an indeterminate order.
 func Keys[M ~map[K]V, K comparable, V any](m M) []K {
 	r := make([]K, 0, len(m))
-	for k := range m {
-		r = append(r, k)
-	}
+	keys(m, unsafe.Pointer(&r))
 	return r
 }
 
+func keysForBenchmarking[M ~map[K]V, K comparable, V any](m M, s []K) {
+	keys(m, unsafe.Pointer(&s))
+}
+
 // Values returns the values of the map m.
 // The values will be in an indeterminate order.
 func Values[M ~map[K]V, K comparable, V any](m M) []V {
diff --git a/src/maps/maps_test.go b/src/maps/maps_test.go
index 1825df5b77..a7a8c10f71 100644
--- a/src/maps/maps_test.go
+++ b/src/maps/maps_test.go
@@ -29,6 +29,23 @@ func TestKeys(t *testing.T) {
 	if !slices.Equal(got2, want) {
 		t.Errorf("Keys(%v) = %v, want %v", m2, got2, want)
 	}
+
+	// test for oldbucket code path
+	// We grow from 128 to 256 buckets at size 832 (6.5 * 128).
+	// Then we have to evacuate 128 buckets, which means we'll be done evacuation at 832+128=960 elements inserted.
+	// so 840 is a good number to test for oldbucket code path.
+	var want3 []int
+	var m = make(map[int]int)
+	for i := 0; i < 840; i++ {
+		want3 = append(want3, i)
+		m[i] = i
+	}
+
+	got3 := Keys(m)
+	sort.Ints(got3)
+	if !slices.Equal(got3, want3) {
+		t.Errorf("Keys(%v) = %v, want %v", m, got3, want3)
+	}
 }
 
 func TestValues(t *testing.T) {
@@ -216,3 +233,18 @@ func TestCloneWithMapAssign(t *testing.T) {
 		}
 	}
 }
+
+var keysArr []int
+
+func BenchmarkKeys(b *testing.B) {
+	m := make(map[int]int, 1000000)
+	for i := 0; i < 1000000; i++ {
+		m[i] = i
+	}
+	b.ResetTimer()
+
+	slice := make([]int, 0, len(m))
+	for i := 0; i < b.N; i++ {
+		keysForBenchmarking(m, slice)
+	}
+}