runtime/internal/maps: optimize long string keys for small maps

For large strings, do a quick equality check on all the slots.
Only if more than one passes the quick equality check do we
resort to hashing.

                               │    baseline    │             experiment              │
                               │     sec/op     │   sec/op     vs base                │
MegMap-24                        16609.50n ± 1%   13.91n ± 3%  -99.92% (p=0.000 n=10)
MegOneMap-24                     16655.00n ± 0%   12.27n ± 1%  -99.93% (p=0.000 n=10)
MegEqMap-24                         41.31µ ± 1%   25.03µ ± 1%  -39.40% (p=0.000 n=10)
MegEmptyMap-24                      2.034n ± 0%   2.027n ± 2%        ~ (p=0.541 n=10)
MegEmptyMapWithInterfaceKey-24      5.931n ± 2%   5.599n ± 1%   -5.60% (p=0.000 n=10)
MapStringKeysEight_16-24            8.473n ± 7%   8.224n ± 5%        ~ (p=0.315 n=10)
MapStringKeysEight_32-24            8.441n ± 2%   8.147n ± 1%   -3.48% (p=0.002 n=10)
MapStringKeysEight_64-24            8.769n ± 1%   8.517n ± 1%   -2.87% (p=0.000 n=10)
MapStringKeysEight_128-24           10.73n ± 4%   13.57n ± 8%  +26.57% (p=0.000 n=10)
MapStringKeysEight_256-24           12.97n ± 2%   14.35n ± 4%  +10.64% (p=0.001 n=10)
MapStringKeysEight_1M-24         17359.50n ± 3%   13.92n ± 4%  -99.92% (p=0.000 n=10)

Change-Id: I4cc2ea4edab12a4b03236de626c7bcf0f96b6cc0
Reviewed-on: https://go-review.googlesource.com/c/go/+/625905
Reviewed-by: Keith Randall <khr@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Pratt <mpratt@google.com>

1 files changed, 79 insertions, 22 deletions

diff --git a/src/internal/runtime/maps/runtime_faststr_swiss.go b/src/internal/runtime/maps/runtime_faststr_swiss.go
index 08172334e7..38170a1821 100644
--- a/src/internal/runtime/maps/runtime_faststr_swiss.go
+++ b/src/internal/runtime/maps/runtime_faststr_swiss.go
@@ -13,32 +13,89 @@ import (
 	"unsafe"
 )
 
-// TODO: more string-specific optimizations possible.
-
-func (m *Map) getWithoutKeySmallFastStr(typ *abi.SwissMapType, hash uintptr, key string) (unsafe.Pointer, bool) {
+func (m *Map) getWithoutKeySmallFastStr(typ *abi.SwissMapType, key string) unsafe.Pointer {
 	g := groupReference{
 		data: m.dirPtr,
 	}
 
-	h2 := uint8(h2(hash))
 	ctrls := *g.ctrls()
+	slotKey := g.key(typ, 0)
+	slotSize := typ.SlotSize
 
-	for i := uintptr(0); i < abi.SwissMapGroupSlots; i++ {
-		c := uint8(ctrls)
-		ctrls >>= 8
-		if c != h2 {
-			continue
+	// The 64 threshold was chosen based on performance of BenchmarkMapStringKeysEight,
+	// where there are 8 keys to check, all of which don't quick-match the lookup key.
+	// In that case, we can save hashing the lookup key. That savings is worth this extra code
+	// for strings that are long enough that hashing is expensive.
+	if len(key) > 64 {
+		// String hashing and equality might be expensive. Do a quick check first.
+		j := abi.SwissMapGroupSlots
+		for i := range abi.SwissMapGroupSlots {
+			if ctrls&(1<<7) == 0 && longStringQuickEqualityTest(key, *(*string)(slotKey)) {
+				if j < abi.SwissMapGroupSlots {
+					// 2 strings both passed the quick equality test.
+					// Break out of this loop and do it the slow way.
+					goto dohash
+				}
+				j = i
+			}
+			slotKey = unsafe.Pointer(uintptr(slotKey) + slotSize)
+			ctrls >>= 8
+		}
+		if j == abi.SwissMapGroupSlots {
+			// No slot passed the quick test.
+			return nil
+		}
+		// There's exactly one slot that passed the quick test. Do the single expensive comparison.
+		slotKey = g.key(typ, uintptr(j))
+		if key == *(*string)(slotKey) {
+			return unsafe.Pointer(uintptr(slotKey) + typ.ElemOff)
 		}
+		return nil
+	}
 
-		slotKey := g.key(typ, i)
+dohash:
+	// This path will cost 1 hash and 1+ε comparisons.
+	hash := typ.Hasher(abi.NoEscape(unsafe.Pointer(&key)), m.seed)
+	h2 := uint8(h2(hash))
+	ctrls = *g.ctrls()
+	slotKey = g.key(typ, 0)
 
-		if key == *(*string)(slotKey) {
-			slotElem := g.elem(typ, i)
-			return slotElem, true
+	for range abi.SwissMapGroupSlots {
+		if uint8(ctrls) == h2 && key == *(*string)(slotKey) {
+			return unsafe.Pointer(uintptr(slotKey) + typ.ElemOff)
 		}
+		slotKey = unsafe.Pointer(uintptr(slotKey) + slotSize)
+		ctrls >>= 8
 	}
+	return nil
+}
 
-	return nil, false
+// Returns true if a and b might be equal.
+// Returns false if a and b are definitely not equal.
+// Requires len(a)>=8.
+func longStringQuickEqualityTest(a, b string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	x, y := stringPtr(a), stringPtr(b)
+	// Check first 8 bytes.
+	if *(*[8]byte)(x) != *(*[8]byte)(y) {
+		return false
+	}
+	// Check last 8 bytes.
+	x = unsafe.Pointer(uintptr(x) + uintptr(len(a)) - 8)
+	y = unsafe.Pointer(uintptr(y) + uintptr(len(a)) - 8)
+	if *(*[8]byte)(x) != *(*[8]byte)(y) {
+		return false
+	}
+	return true
+}
+func stringPtr(s string) unsafe.Pointer {
+	type stringStruct struct {
+		ptr unsafe.Pointer
+		len int
+	}
+	return (*stringStruct)(unsafe.Pointer(&s)).ptr
 }
 
 //go:linkname runtime_mapaccess1_faststr runtime.mapaccess1_faststr
@@ -58,16 +115,16 @@ func runtime_mapaccess1_faststr(typ *abi.SwissMapType, m *Map, key string) unsaf
 		return nil
 	}
 
-	hash := typ.Hasher(abi.NoEscape(unsafe.Pointer(&key)), m.seed)
-
 	if m.dirLen <= 0 {
-		elem, ok := m.getWithoutKeySmallFastStr(typ, hash, key)
-		if !ok {
+		elem := m.getWithoutKeySmallFastStr(typ, key)
+		if elem == nil {
 			return unsafe.Pointer(&zeroVal[0])
 		}
 		return elem
 	}
 
+	hash := typ.Hasher(abi.NoEscape(unsafe.Pointer(&key)), m.seed)
+
 	// Select table.
 	idx := m.directoryIndex(hash)
 	t := m.directoryAt(idx)
@@ -116,16 +173,16 @@ func runtime_mapaccess2_faststr(typ *abi.SwissMapType, m *Map, key string) (unsa
 		return nil, false
 	}
 
-	hash := typ.Hasher(abi.NoEscape(unsafe.Pointer(&key)), m.seed)
-
 	if m.dirLen <= 0 {
-		elem, ok := m.getWithoutKeySmallFastStr(typ, hash, key)
-		if !ok {
+		elem := m.getWithoutKeySmallFastStr(typ, key)
+		if elem == nil {
 			return unsafe.Pointer(&zeroVal[0]), false
 		}
 		return elem, true
 	}
 
+	hash := typ.Hasher(abi.NoEscape(unsafe.Pointer(&key)), m.seed)
+
 	// Select table.
 	idx := m.directoryIndex(hash)
 	t := m.directoryAt(idx)