runtime: reuse map string key instead of allocating

[dsnet]

Consider the following benchmark:

func Benchmark(b *testing.B) {
	m := map[string]int{}
	k := []byte("hello, world")

	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		m[string(k)] = 1
	}
}

This currently prints:

Benchmark    	37319388	        29.83 ns/op	      16 B/op	       1 allocs/op

I expect this to not allocate since it is reusing the same map entry every time.
The very first mapassign will need to allocate the string on heap, but subsequent operations should not need to.
This is related to #45045, which made this optimization harder.

\cc @randall77 @cuonglm

[dsnet]

One way to implement this is to have a variation of mapassign_faststr called mapassign_faststr_noescape that is only called if s does not escape. This implementation would avoid the overwrite functionality introduced in #45045, and clone the string upon first insertion.

[zigo101]

Doesn't the allocation happen during the []byte->string conversion instead?

[dsnet]

I do not believe that is what is happening since switching to map access results in zero allocations:

func Benchmark(b *testing.B) {
	m := map[string]int{}
	k := []byte("hello, world")

	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		sink += m[string(k)]
	}
}

prints:

Benchmark-24    	602860267	         1.980 ns/op	       0 B/op	       0 allocs/op

[randall77]

There is a special case in the compiler for m[string(b)] not on the LHS of an assignment. The string->[]byte conversion can be done without allocation because the resulting string is guaranteed to be dead afterwards.
That's not the case when m[string(b)] is on the LHS. @dsnet has a reasonable plan, have a special map write that clones the "string" (which would be just the header of a byte slice reinterpreted as a string) if it needs to keep it in the map.

That said, where does this come up? You can always do

  if _, ok := m[string(b)]; !ok {
    m[string(b)] = v
  }

That saves an allocation at the cost of a map lookup.

[dsnet]

@randall77. Unfortunately, the suggested workaround doesn't work if you want to update the value for an entry at the same key, which is the situation that I'm running into.

[randall77]

I see, you're getting bitten by the key overwrite we decided was correct in #11088.

[dsnet]

In #11088 (comment):

In Russ' example, suppose you had a map from strings to something. Suppose the strings were small but potentially keeping large buffers live. You could do:

for k, v := range m {
   m[string([]byte(k))] = v
}

To allocate only the backing store required for the existing strings. If the behavior was the other way around (first write wins) there's no way to fix the map. You'd have to allocate a new map.

Personally, I've never seen that pattern ever before. Usually, if someone were concerned about string keys pinning a large amount of memory, they do the equivalent of strings.Clone(k) on the first insertion, rather than trying trim the map after the fact.

On the other hand, I'd argue that this issue would benefit patterns that are more common:

1. It is common to have a []byte on hand since that's usually what you read off the network or a file. You slice out a key from that []byte and do map operations (some of which may be updating an existing entry).
2. In our case, we have a map[string]V where we want the map to be case-insensitive. I have a helper function that does the equivalent of:

   func CaseInsensitiveSet[V any](m map[string]V, k string, v V) {
   	if isLowerASCII(k) {
   		m[k] = v
   	} else {
   		var a [32]byte
   		m[string(appendToLower(a[:0], k))] = v
   	}
   }

   In the event that the string contains uppercase letters, we convert it to lowercase in stack memory.
   Unfortunately, this pattern always allocates even if the entry exists.

[dsnet]

Also, the -0 and +0 issue makes sense for general purpose maps, but since we already have a fast-path for map[string]V it seems acceptable that we can differ in behavior here.

[cuonglm]

@dsnet Could we do something like this:

func BenchmarkAlloc(b *testing.B) {
	m := map[string]int{}
	k := []byte("hello, world")
	ks := unsafe.String(&k[0], len(k))

	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		if _, ok := m[string(k)]; !ok {
			m[string(k)] = 1
			continue
		}
		m[ks] = 1
	}
}

[dsnet]

@cuonglm. Yuck. Must we use "unsafe"? 😛

Isn't this also dangerous given that the map implementation stores a reference to the last string key per #45045. This would mean that mutations to k would be observable in the map keys.

https://go.dev/play/p/GZXB25nOcPR?v=gotip

[cuonglm]

@cuonglm. Yuck. Must we use "unsafe"? 😛

Isn't this also dangerous given that the map implementation stores a reference to the last string key per #45045. This would mean that mutations to k would be observable in the map keys.

https://go.dev/play/p/GZXB25nOcPR?v=gotip

Yeah, but I see in your example that appendToLower is controlled by us, so it should still be safe, no 🤔

[dsnet]

It appends into a buffer that is allocated on the stack (backed by the a array). Wouldn't the stack memory become invalid when the function returns? That seems wrong to me.

[cuonglm]

It appends into a buffer that is allocated on the stack (backed by the a array). Wouldn't the stack memory become invalid when the function returns? That seems wrong to me.

So a would be moved to heap:

package p

import (
	"testing"
	"unsafe"
)

var m1 = map[string]int{}
var m2 = map[string]int{}

func CaseInsensitiveSet[V any](m map[string]V, k string, v V) {
	var a [32]byte
	m[string(appendToLower(a[:0], k))] = v
}

func CaseInsensitiveSetUnsafe[V any](m map[string]V, k string, v V) {
	var a [32]byte
	ks := appendToLower(a[:0], k)
	if _, ok := m[string(ks)]; !ok {
		m[string(ks)] = v
		return
	}
	key := unsafe.String(&a[0], len(a))
	m[key] = v
}

func appendToLower(s []byte, k string) []byte {
	return append(s, k...)
}

func BenchmarkAlloc(b *testing.B) {
	k := "hello, world"

	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		CaseInsensitiveSet(m1, k, 1)
	}
}

func BenchmarkAllocUnsafe(b *testing.B) {
	k := "hello, world"

	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		CaseInsensitiveSetUnsafe(m2, k, 1)
	}
}

And that's count as allocating:

goos: darwin
goarch: arm64
pkg: p
BenchmarkAlloc-8         	53314322	        22.18 ns/op	      16 B/op	       1 allocs/op
BenchmarkAllocUnsafe-8   	39893672	        29.14 ns/op	      32 B/op	       1 allocs/op
PASS
ok  	p	3.518s

My original example works because m is also not escaped to heap.