internal/runtime/maps: proper capacity hint handling

When given a hint size, set the initial capacity large enough to avoid
requiring growth in the average case.

When not given a hint (or given 0), don't allocate anything at all.

For #54766.

Cq-Include-Trybots: luci.golang.try:gotip-linux-amd64-longtest-swissmap
Change-Id: I8844fc652b8d2d4e5136cd56f7e78999a07fe381
Reviewed-on: https://go-review.googlesource.com/c/go/+/616457
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Keith Randall <khr@golang.org>
Reviewed-by: Keith Randall <khr@google.com>
Auto-Submit: Michael Pratt <mpratt@google.com>

7 files changed, 128 insertions, 57 deletions

diff --git a/src/internal/runtime/maps/export_test.go b/src/internal/runtime/maps/export_test.go
index 0cc78b954f..151c11fba8 100644
--- a/src/internal/runtime/maps/export_test.go
+++ b/src/internal/runtime/maps/export_test.go
@@ -18,9 +18,13 @@ var AlignUpPow2 = alignUpPow2
 const MaxTableCapacity = maxTableCapacity
 const MaxAvgGroupLoad = maxAvgGroupLoad
 
-func NewTestMap[K comparable, V any](length uint64) (*Map, *abi.SwissMapType) {
+// This isn't equivalent to runtime.maxAlloc. It is fine for basic testing but
+// we can't properly test hint alloc overflows with this.
+const maxAllocTest = 1 << 30
+
+func NewTestMap[K comparable, V any](hint uintptr) (*Map, *abi.SwissMapType) {
 	mt := newTestMapType[K, V]()
-	return NewMap(mt, length), mt
+	return NewMap(mt, hint, maxAllocTest), mt
 }
 
 func (m *Map) TableCount() int {
diff --git a/src/internal/runtime/maps/map.go b/src/internal/runtime/maps/map.go
index 2bfc5b7fb7..ae8afc3ea7 100644
--- a/src/internal/runtime/maps/map.go
+++ b/src/internal/runtime/maps/map.go
@@ -8,6 +8,7 @@ package maps
 import (
 	"internal/abi"
 	"internal/goarch"
+	"internal/runtime/math"
 	"internal/runtime/sys"
 	"unsafe"
 )
@@ -240,17 +241,50 @@ func depthToShift(depth uint8) uint8 {
 	return 64 - depth
 }
 
-func NewMap(mt *abi.SwissMapType, capacity uint64) *Map {
-	if capacity < abi.SwissMapGroupSlots {
-		// TODO: temporary to simplify initial implementation.
-		capacity = abi.SwissMapGroupSlots
+// maxAlloc should be runtime.maxAlloc.
+//
+// TODO(prattmic): Put maxAlloc somewhere accessible.
+func NewMap(mt *abi.SwissMapType, hint, maxAlloc uintptr) *Map {
+	// Set initial capacity to hold hint entries without growing in the
+	// average case.
+	var targetCapacity uintptr
+	if hint <= abi.SwissMapGroupSlots {
+		// Small map can fill all 8 slots. We set the target to 0 here
+		// because an 8 slot small map is what the first assignment to
+		// an empty map will allocate anyway. Whether we allocate here
+		// or in the first assignment makes no difference. And if there
+		// is a chance that the caller won't write at all then it is
+		// better to delay.
+		targetCapacity = 0
+	} else {
+		targetCapacity = (hint * abi.SwissMapGroupSlots) / maxAvgGroupLoad
+		if targetCapacity < hint { // overflow
+			targetCapacity = 0
+		}
 	}
-	dirSize := (capacity + maxTableCapacity - 1) / maxTableCapacity
+
+	dirSize := (uint64(targetCapacity) + maxTableCapacity - 1) / maxTableCapacity
 	dirSize, overflow := alignUpPow2(dirSize)
+	if overflow || dirSize > uint64(math.MaxUintptr) {
+		targetCapacity = 0
+	}
+
+	// Reject hints that are obviously too large.
+	groups, overflow := math.MulUintptr(uintptr(dirSize), maxTableCapacity)
 	if overflow {
-		panic("rounded-up capacity overflows uint64")
+		targetCapacity = 0
+	} else {
+		mem, overflow := math.MulUintptr(groups, mt.Group.Size_)
+		if overflow || mem > maxAlloc {
+			targetCapacity = 0
+		}
 	}
+
 	globalDepth := uint8(sys.TrailingZeros64(dirSize))
+	if targetCapacity == 0 {
+		// TrailingZeros64 returns 64 for 0.
+		globalDepth = 0
+	}
 
 	m := &Map{
 		//TODO
@@ -262,25 +296,17 @@ func NewMap(mt *abi.SwissMapType, capacity uint64) *Map {
 		globalShift: depthToShift(globalDepth),
 	}
 
-	if capacity > abi.SwissMapGroupSlots {
+	if targetCapacity > 0 {
+		// Full map.
 		directory := make([]*table, dirSize)
 
 		for i := range directory {
 			// TODO: Think more about initial table capacity.
-			directory[i] = newTable(mt, capacity/dirSize, i, globalDepth)
+			directory[i] = newTable(mt, uint64(targetCapacity)/dirSize, i, globalDepth)
 		}
 
 		m.dirPtr = unsafe.Pointer(&directory[0])
 		m.dirLen = len(directory)
-	} else {
-		grp := newGroups(mt, 1)
-		m.dirPtr = grp.data
-		m.dirLen = 0
-
-		g := groupReference{
-			data: m.dirPtr,
-		}
-		g.ctrls().setEmpty()
 	}
 
 	return m
@@ -356,6 +382,10 @@ func (m *Map) Get(typ *abi.SwissMapType, key unsafe.Pointer) (unsafe.Pointer, bo
 }
 
 func (m *Map) getWithKey(typ *abi.SwissMapType, key unsafe.Pointer) (unsafe.Pointer, unsafe.Pointer, bool) {
+	if m.Used() == 0 {
+		return nil, nil, false
+	}
+
 	hash := typ.Hasher(key, m.seed)
 
 	if m.dirLen == 0 {
@@ -367,6 +397,10 @@ func (m *Map) getWithKey(typ *abi.SwissMapType, key unsafe.Pointer) (unsafe.Poin
 }
 
 func (m *Map) getWithoutKey(typ *abi.SwissMapType, key unsafe.Pointer) (unsafe.Pointer, bool) {
+	if m.Used() == 0 {
+		return nil, false
+	}
+
 	hash := typ.Hasher(key, m.seed)
 
 	if m.dirLen == 0 {
@@ -414,6 +448,10 @@ func (m *Map) Put(typ *abi.SwissMapType, key, elem unsafe.Pointer) {
 func (m *Map) PutSlot(typ *abi.SwissMapType, key unsafe.Pointer) unsafe.Pointer {
 	hash := typ.Hasher(key, m.seed)
 
+	if m.dirPtr == nil {
+		m.growToSmall(typ)
+	}
+
 	if m.dirLen == 0 {
 		if m.used < abi.SwissMapGroupSlots {
 			return m.putSlotSmall(typ, hash, key)
@@ -464,7 +502,7 @@ func (m *Map) putSlotSmall(typ *abi.SwissMapType, hash uintptr, key unsafe.Point
 	// deleteSmall).
 	match = g.ctrls().matchEmpty()
 	if match == 0 {
-		panic("small map with no empty slot")
+		fatal("small map with no empty slot (concurrent map writes?)")
 	}
 
 	i := match.first()
@@ -479,6 +517,16 @@ func (m *Map) putSlotSmall(typ *abi.SwissMapType, hash uintptr, key unsafe.Point
 	return slotElem
 }
 
+func (m *Map) growToSmall(typ *abi.SwissMapType) {
+	grp := newGroups(typ, 1)
+	m.dirPtr = grp.data
+
+	g := groupReference{
+		data: m.dirPtr,
+	}
+	g.ctrls().setEmpty()
+}
+
 func (m *Map) growToTable(typ *abi.SwissMapType) {
 	tab := newTable(typ, 2*abi.SwissMapGroupSlots, 0, 0)
 
@@ -508,6 +556,13 @@ func (m *Map) growToTable(typ *abi.SwissMapType) {
 }
 
 func (m *Map) Delete(typ *abi.SwissMapType, key unsafe.Pointer) {
+	if m == nil || m.Used() == 0 {
+		if err := mapKeyError(typ, key); err != nil {
+			panic(err) // see issue 23734
+		}
+		return
+	}
+
 	hash := typ.Hasher(key, m.seed)
 
 	if m.dirLen == 0 {
@@ -546,6 +601,10 @@ func (m *Map) deleteSmall(typ *abi.SwissMapType, hash uintptr, key unsafe.Pointe
 
 // Clear deletes all entries from the map resulting in an empty map.
 func (m *Map) Clear(typ *abi.SwissMapType) {
+	if m == nil || m.Used() == 0 {
+		return
+	}
+
 	if m.dirLen == 0 {
 		m.clearSmall(typ)
 		return
diff --git a/src/internal/runtime/maps/map_swiss_test.go b/src/internal/runtime/maps/map_swiss_test.go
index 7c6b426f6d..4e02f3e660 100644
--- a/src/internal/runtime/maps/map_swiss_test.go
+++ b/src/internal/runtime/maps/map_swiss_test.go
@@ -55,53 +55,47 @@ func TestTableGroupCount(t *testing.T) {
 		{
 			n: -(1 << 30),
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
-				initialHint: mapCount{0, 1},
-				after:       mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
+				initialHint: mapCount{0, 0},
+				after:       mapCount{0, 0},
 			},
 		},
 		{
 			n: -1,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
-				initialHint: mapCount{0, 1},
-				after:       mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
+				initialHint: mapCount{0, 0},
+				after:       mapCount{0, 0},
 			},
 		},
 		{
 			n: 0,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
-				initialHint: mapCount{0, 1},
-				after:       mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
+				initialHint: mapCount{0, 0},
+				after:       mapCount{0, 0},
 			},
 		},
 		{
 			n: 1,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
-				initialHint: mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
+				initialHint: mapCount{0, 0},
 				after:       mapCount{0, 1},
 			},
 		},
 		{
 			n: abi.SwissMapGroupSlots,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
-				initialHint: mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
+				initialHint: mapCount{0, 0},
 				after:       mapCount{0, 1},
 			},
 		},
 		{
 			n: abi.SwissMapGroupSlots + 1,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
 				initialHint: mapCount{1, 2},
 				after:       mapCount{1, 2},
 			},
@@ -109,8 +103,7 @@ func TestTableGroupCount(t *testing.T) {
 		{
 			n: belowMax, // 1.5 group max = 2 groups @ 75%
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
 				initialHint: mapCount{1, 2},
 				after:       mapCount{1, 2},
 			},
@@ -118,8 +111,7 @@ func TestTableGroupCount(t *testing.T) {
 		{
 			n: atMax, // 2 groups at max
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
 				initialHint: mapCount{1, 2},
 				after:       mapCount{1, 2},
 			},
@@ -127,18 +119,15 @@ func TestTableGroupCount(t *testing.T) {
 		{
 			n: atMax + 1, // 2 groups at max + 1 -> grow to 4 groups
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit: mapCount{0, 1},
-				// TODO(go.dev/issue/54766): Initial capacity should round hint up to avoid grow.
-				initialHint: mapCount{1, 2},
+				initialLit:  mapCount{0, 0},
+				initialHint: mapCount{1, 4},
 				after:       mapCount{1, 4},
 			},
 		},
 		{
 			n: 2 * belowMax, // 3 * group max = 4 groups @75%
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit:  mapCount{0, 1},
+				initialLit:  mapCount{0, 0},
 				initialHint: mapCount{1, 4},
 				after:       mapCount{1, 4},
 			},
@@ -146,10 +135,8 @@ func TestTableGroupCount(t *testing.T) {
 		{
 			n: 2*atMax + 1, // 4 groups at max + 1 -> grow to 8 groups
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
-				initialLit: mapCount{0, 1},
-				// TODO(go.dev/issue/54766): Initial capacity should round hint up to avoid grow.
-				initialHint: mapCount{1, 4},
+				initialLit:  mapCount{0, 0},
+				initialHint: mapCount{1, 8},
 				after:       mapCount{1, 8},
 			},
 		},
diff --git a/src/internal/runtime/maps/runtime.go b/src/internal/runtime/maps/runtime.go
index 9ebfb34b28..0d569de214 100644
--- a/src/internal/runtime/maps/runtime.go
+++ b/src/internal/runtime/maps/runtime.go
@@ -11,6 +11,9 @@ import (
 
 // Functions below pushed from runtime.
 
+//go:linkname fatal
+func fatal(s string)
+
 //go:linkname rand
 func rand() uint64
 
diff --git a/src/internal/runtime/maps/runtime_noswiss.go b/src/internal/runtime/maps/runtime_noswiss.go
new file mode 100644
index 0000000000..c9342e08dd
--- /dev/null
+++ b/src/internal/runtime/maps/runtime_noswiss.go
@@ -0,0 +1,17 @@
+// 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.
+
+//go:build !goexperiment.swissmap
+
+package maps
+
+import (
+	"internal/abi"
+	"unsafe"
+)
+
+// For testing, we don't ever need key errors.
+func mapKeyError(typ *abi.SwissMapType, p unsafe.Pointer) error {
+	return nil
+}
diff --git a/src/internal/runtime/maps/runtime_swiss.go b/src/internal/runtime/maps/runtime_swiss.go
index 7a694f4f0e..1cf1dd21e5 100644
--- a/src/internal/runtime/maps/runtime_swiss.go
+++ b/src/internal/runtime/maps/runtime_swiss.go
@@ -122,6 +122,10 @@ func runtime_mapassign(typ *abi.SwissMapType, m *Map, key unsafe.Pointer) unsafe
 
 	hash := typ.Hasher(key, m.seed)
 
+	if m.dirPtr == nil {
+		m.growToSmall(typ)
+	}
+
 	if m.dirLen == 0 {
 		if m.used < abi.SwissMapGroupSlots {
 			return m.putSlotSmall(typ, hash, key)
diff --git a/src/internal/runtime/maps/table.go b/src/internal/runtime/maps/table.go
index 9b7e43837f..797d510269 100644
--- a/src/internal/runtime/maps/table.go
+++ b/src/internal/runtime/maps/table.go
@@ -812,9 +812,6 @@ func (t *table) rehash(typ *abi.SwissMapType, m *Map) {
 	// new allocation, so the existing grow support in iteration would
 	// continue to work.
 
-	// TODO(prattmic): split table
-	// TODO(prattmic): Avoid overflow (splitting the table will achieve this)
-
 	newCapacity := 2 * t.capacity
 	if newCapacity <= maxTableCapacity {
 		t.grow(typ, m, newCapacity)