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>

src/go/build/deps_test.go

@@ -89,8 +89,8 @@ var depsRules = `
 	< internal/runtime/syscall
 	< internal/runtime/atomic
 	< internal/runtime/exithook
-	< internal/runtime/maps
 	< internal/runtime/math
+	< internal/runtime/maps
 	< runtime
 	< internal/race
 	< sync/atomic

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 {

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 {
+// maxAlloc should be runtime.maxAlloc.
-	if capacity < abi.SwissMapGroupSlots {
+//
-		// TODO: temporary to simplify initial implementation.
+// TODO(prattmic): Put maxAlloc somewhere accessible.
-		capacity = abi.SwissMapGroupSlots
+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 {
+	if overflow || dirSize > uint64(math.MaxUintptr) {
-		panic("rounded-up capacity overflows uint64")
+		targetCapacity = 0
 	}
+
+	// Reject hints that are obviously too large.
+	groups, overflow := math.MulUintptr(uintptr(dirSize), maxTableCapacity)
+	if overflow {
+		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

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, 0},
-				initialLit:  mapCount{0, 1},
+				initialHint: mapCount{0, 0},
-				initialHint: mapCount{0, 1},
+				after:       mapCount{0, 0},
-				after:       mapCount{0, 1},
 			},
 		},
 		{
 			n: -1,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
+				initialLit:  mapCount{0, 0},
-				initialLit:  mapCount{0, 1},
+				initialHint: mapCount{0, 0},
-				initialHint: mapCount{0, 1},
+				after:       mapCount{0, 0},
-				after:       mapCount{0, 1},
 			},
 		},
 		{
 			n: 0,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
+				initialLit:  mapCount{0, 0},
-				initialLit:  mapCount{0, 1},
+				initialHint: mapCount{0, 0},
-				initialHint: mapCount{0, 1},
+				after:       mapCount{0, 0},
-				after:       mapCount{0, 1},
 			},
 		},
 		{
 			n: 1,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
+				initialLit:  mapCount{0, 0},
-				initialLit:  mapCount{0, 1},
+				initialHint: mapCount{0, 0},
-				initialHint: mapCount{0, 1},
 				after:       mapCount{0, 1},
 			},
 		},
 		{
 			n: abi.SwissMapGroupSlots,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
+				initialLit:  mapCount{0, 0},
-				initialLit:  mapCount{0, 1},
+				initialHint: mapCount{0, 0},
-				initialHint: mapCount{0, 1},
 				after:       mapCount{0, 1},
 			},
 		},
 		{
 			n: abi.SwissMapGroupSlots + 1,
 			escape: mapCase{
-				// TODO(go.dev/issue/54766): empty maps
+				initialLit:  mapCount{0, 0},
-				initialLit:  mapCount{0, 1},
 				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, 0},
-				initialLit:  mapCount{0, 1},
 				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, 0},
-				initialLit:  mapCount{0, 1},
 				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, 0},
-				initialLit: mapCount{0, 1},
+				initialHint: mapCount{1, 4},
-				// TODO(go.dev/issue/54766): Initial capacity should round hint up to avoid grow.
-				initialHint: mapCount{1, 2},
 				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, 0},
-				initialLit:  mapCount{0, 1},
 				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, 0},
-				initialLit: mapCount{0, 1},
+				initialHint: mapCount{1, 8},
-				// TODO(go.dev/issue/54766): Initial capacity should round hint up to avoid grow.
-				initialHint: mapCount{1, 4},
 				after:       mapCount{1, 8},
 			},
 		},

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
 

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
+}

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)

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)

src/runtime/alg.go

@@ -256,11 +256,6 @@ func mapKeyError(t *maptype, p unsafe.Pointer) error {
 	return mapKeyError2(t.Key, p)
 }
 
-//go:linkname maps_mapKeyError internal/runtime/maps.mapKeyError
-func maps_mapKeyError(t *maptype, p unsafe.Pointer) error {
-	return mapKeyError(t, p)
-}
-
 func mapKeyError2(t *_type, p unsafe.Pointer) error {
 	if t.TFlag&abi.TFlagRegularMemory != 0 {
 		return nil

src/runtime/crash_test.go

@@ -622,7 +622,10 @@ func TestConcurrentMapWrites(t *testing.T) {
 	testenv.MustHaveGoRun(t)
 	output := runTestProg(t, "testprog", "concurrentMapWrites")
 	want := "fatal error: concurrent map writes\n"
-	if !strings.HasPrefix(output, want) {
+	// Concurrent writes can corrupt the map in a way that we
+	// detect with a separate throw.
+	want2 := "fatal error: small map with no empty slot (concurrent map writes?)\n"
+	if !strings.HasPrefix(output, want) && !strings.HasPrefix(output, want2) {
 		t.Fatalf("output does not start with %q:\n%s", want, output)
 	}
 }
@@ -633,7 +636,10 @@ func TestConcurrentMapReadWrite(t *testing.T) {
 	testenv.MustHaveGoRun(t)
 	output := runTestProg(t, "testprog", "concurrentMapReadWrite")
 	want := "fatal error: concurrent map read and map write\n"
-	if !strings.HasPrefix(output, want) {
+	// Concurrent writes can corrupt the map in a way that we
+	// detect with a separate throw.
+	want2 := "fatal error: small map with no empty slot (concurrent map writes?)\n"
+	if !strings.HasPrefix(output, want) && !strings.HasPrefix(output, want2) {
 		t.Fatalf("output does not start with %q:\n%s", want, output)
 	}
 }
@@ -644,7 +650,10 @@ func TestConcurrentMapIterateWrite(t *testing.T) {
 	testenv.MustHaveGoRun(t)
 	output := runTestProg(t, "testprog", "concurrentMapIterateWrite")
 	want := "fatal error: concurrent map iteration and map write\n"
-	if !strings.HasPrefix(output, want) {
+	// Concurrent writes can corrupt the map in a way that we
+	// detect with a separate throw.
+	want2 := "fatal error: small map with no empty slot (concurrent map writes?)\n"
+	if !strings.HasPrefix(output, want) && !strings.HasPrefix(output, want2) {
 		t.Fatalf("output does not start with %q:\n%s", want, output)
 	}
 }
@@ -667,7 +676,10 @@ func TestConcurrentMapWritesIssue69447(t *testing.T) {
 			continue
 		}
 		want := "fatal error: concurrent map writes\n"
-		if !strings.HasPrefix(output, want) {
+		// Concurrent writes can corrupt the map in a way that we
+		// detect with a separate throw.
+		want2 := "fatal error: small map with no empty slot (concurrent map writes?)\n"
+		if !strings.HasPrefix(output, want) && !strings.HasPrefix(output, want2) {
 			t.Fatalf("output does not start with %q:\n%s", want, output)
 		}
 	}

src/runtime/map_swiss.go

@@ -9,7 +9,6 @@ package runtime
 import (
 	"internal/abi"
 	"internal/runtime/maps"
-	"internal/runtime/math"
 	"internal/runtime/sys"
 	"unsafe"
 )
@@ -25,6 +24,11 @@ type maptype = abi.SwissMapType
 //go:linkname maps_errNilAssign internal/runtime/maps.errNilAssign
 var maps_errNilAssign error = plainError("assignment to entry in nil map")
 
+//go:linkname maps_mapKeyError internal/runtime/maps.mapKeyError
+func maps_mapKeyError(t *abi.SwissMapType, p unsafe.Pointer) error {
+	return mapKeyError(t, p)
+}
+
 func makemap64(t *abi.SwissMapType, hint int64, m *maps.Map) *maps.Map {
 	if int64(int(hint)) != hint {
 		hint = 0
@@ -39,63 +43,18 @@ func makemap_small() *maps.Map {
 	panic("unimplemented")
 }
 
-// checkHint verifies that hint is reasonable, adjusting as necessary.
-func checkHint(t *abi.SwissMapType, hint int) uint64 {
-	if hint <= 0 {
-		return 0
-	}
-
-	capacity := uint64(hint)
-
-	// Ensure a groups allocation for a capacity this high doesn't exceed
-	// the maximum allocation size.
-	//
-	// TODO(prattmic): Once we split tables, a large hint will result in
-	// splitting the tables up front, which will use smaller individual
-	// allocations.
-	//
-	// TODO(prattmic): This logic is largely duplicated from maps.newTable
-	// / maps.(*table).reset.
-	capacity, overflow := alignUpPow2(capacity)
-	if !overflow {
-		groupCount := capacity / abi.SwissMapGroupSlots
-		mem, overflow := math.MulUintptr(uintptr(groupCount), t.Group.Size_)
-		if overflow || mem > maxAlloc {
-			return 0
-		}
-	} else {
-		return 0
-	}
-
-	return capacity
-}
-
 // makemap implements Go map creation for make(map[k]v, hint).
 // If the compiler has determined that the map or the first bucket
 // can be created on the stack, h and/or bucket may be non-nil.
 // If h != nil, the map can be created directly in h.
 // If h.buckets != nil, bucket pointed to can be used as the first bucket.
 func makemap(t *abi.SwissMapType, hint int, m *maps.Map) *maps.Map {
-	capacity := checkHint(t, hint)
+	if hint < 0 {
+		hint = 0
+	}
 
 	// TODO: use existing m
-	return maps.NewMap(t, capacity)
+	return maps.NewMap(t, uintptr(hint), maxAlloc)
-}
-
-// alignUpPow2 rounds n up to the next power of 2.
-//
-// Returns true if round up causes overflow.
-//
-// TODO(prattmic): deduplicate from internal/runtime/maps.
-func alignUpPow2(n uint64) (uint64, bool) {
-	if n == 0 {
-		return 0, false
-	}
-	v := (uint64(1) << sys.Len64(n-1))
-	if v == 0 {
-		return 0, true
-	}
-	return v, false
 }
 
 // mapaccess1 returns a pointer to h[key].  Never returns nil, instead
@@ -176,13 +135,6 @@ func mapdelete(t *abi.SwissMapType, m *maps.Map, key unsafe.Pointer) {
 		asanread(key, t.Key.Size_)
 	}
 
-	if m == nil || m.Used() == 0 {
-		if err := mapKeyError(t, key); err != nil {
-			panic(err) // see issue 23734
-		}
-		return
-	}
-
 	m.Delete(t, key)
 }
 
@@ -219,10 +171,6 @@ func mapclear(t *abi.SwissMapType, m *maps.Map) {
 		racewritepc(unsafe.Pointer(m), callerpc, pc)
 	}
 
-	if m == nil || m.Used() == 0 {
-		return
-	}
-
 	m.Clear(t)
 }
 

src/runtime/panic.go

@@ -1043,6 +1043,11 @@ func fips_fatal(s string) {
 	fatal(s)
 }
 
+//go:linkname maps_fatal internal/runtime/maps.fatal
+func maps_fatal(s string) {
+	fatal(s)
+}
+
 // throw triggers a fatal error that dumps a stack trace and exits.
 //
 // throw should be used for runtime-internal fatal errors where Go itself,