cmd/compile,runtime: add indirect key/elem to swissmap

We use the same heuristics as existing maps.

For #54766.

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

7 files changed, 286 insertions, 9 deletions

diff --git a/src/internal/runtime/maps/export_test.go b/src/internal/runtime/maps/export_test.go
index 151c11fba8..c9c1da6a1c 100644
--- a/src/internal/runtime/maps/export_test.go
+++ b/src/internal/runtime/maps/export_test.go
@@ -86,7 +86,11 @@ func (m *Map) KeyFromFullGroup(typ *abi.SwissMapType) unsafe.Pointer {
 				if g.ctrls().get(j) == ctrlDeleted {
 					continue
 				}
-				return g.key(typ, j)
+				slotKey := g.key(typ, j)
+				if typ.IndirectKey() {
+					slotKey = *((*unsafe.Pointer)(slotKey))
+				}
+				return slotKey
 			}
 		}
 	}
diff --git a/src/internal/runtime/maps/map.go b/src/internal/runtime/maps/map.go
index 543340f10c..c2c7c41805 100644
--- a/src/internal/runtime/maps/map.go
+++ b/src/internal/runtime/maps/map.go
@@ -442,8 +442,15 @@ func (m *Map) getWithKeySmall(typ *abi.SwissMapType, hash uintptr, key unsafe.Po
 		}
 
 		slotKey := g.key(typ, i)
+		if typ.IndirectKey() {
+			slotKey = *((*unsafe.Pointer)(slotKey))
+		}
 		if typ.Key.Equal(key, slotKey) {
-			return slotKey, g.elem(typ, i), true
+			slotElem := g.elem(typ, i)
+			if typ.IndirectElem() {
+				slotElem = *((*unsafe.Pointer)(slotElem))
+			}
+			return slotKey, slotElem, true
 		}
 	}
 
@@ -521,12 +528,18 @@ func (m *Map) putSlotSmall(typ *abi.SwissMapType, hash uintptr, key unsafe.Point
 		i := match.first()
 
 		slotKey := g.key(typ, i)
+		if typ.IndirectKey() {
+			slotKey = *((*unsafe.Pointer)(slotKey))
+		}
 		if typ.Key.Equal(key, slotKey) {
 			if typ.NeedKeyUpdate() {
 				typedmemmove(typ.Key, slotKey, key)
 			}
 
 			slotElem := g.elem(typ, i)
+			if typ.IndirectElem() {
+				slotElem = *((*unsafe.Pointer)(slotElem))
+			}
 
 			return slotElem
 		}
@@ -543,8 +556,19 @@ func (m *Map) putSlotSmall(typ *abi.SwissMapType, hash uintptr, key unsafe.Point
 	i := match.first()
 
 	slotKey := g.key(typ, i)
+	if typ.IndirectKey() {
+		kmem := newobject(typ.Key)
+		*(*unsafe.Pointer)(slotKey) = kmem
+		slotKey = kmem
+	}
 	typedmemmove(typ.Key, slotKey, key)
+
 	slotElem := g.elem(typ, i)
+	if typ.IndirectElem() {
+		emem := newobject(typ.Elem)
+		*(*unsafe.Pointer)(slotElem) = emem
+		slotElem = emem
+	}
 
 	g.ctrls().set(i, ctrl(h2(hash)))
 	m.used++
@@ -574,9 +598,23 @@ func (m *Map) growToTable(typ *abi.SwissMapType) {
 			// Empty
 			continue
 		}
+
 		key := g.key(typ, i)
+		if typ.IndirectKey() {
+			key = *((*unsafe.Pointer)(key))
+		}
+
 		elem := g.elem(typ, i)
+		if typ.IndirectElem() {
+			elem = *((*unsafe.Pointer)(elem))
+		}
+
 		hash := typ.Hasher(key, m.seed)
+
+		// TODO(prattmic): For indirect key/elem, this is
+		// allocating new objects for key/elem. That is
+		// unnecessary; the new table could simply point to the
+		// existing object.
 		slotElem := tab.uncheckedPutSlot(typ, hash, key)
 		typedmemmove(typ.Elem, slotElem, elem)
 		tab.used++
@@ -631,11 +669,33 @@ func (m *Map) deleteSmall(typ *abi.SwissMapType, hash uintptr, key unsafe.Pointe
 	for match != 0 {
 		i := match.first()
 		slotKey := g.key(typ, i)
+		origSlotKey := slotKey
+		if typ.IndirectKey() {
+			slotKey = *((*unsafe.Pointer)(slotKey))
+		}
 		if typ.Key.Equal(key, slotKey) {
 			m.used--
 
-			typedmemclr(typ.Key, slotKey)
-			typedmemclr(typ.Elem, g.elem(typ, i))
+			if typ.IndirectKey() {
+				// Clearing the pointer is sufficient.
+				*(*unsafe.Pointer)(origSlotKey) = nil
+			} else if typ.Key.Pointers() {
+				// Only bother clearing if there are pointers.
+				typedmemclr(typ.Key, slotKey)
+			}
+
+			slotElem := g.elem(typ, i)
+			if typ.IndirectElem() {
+				// Clearing the pointer is sufficient.
+				*(*unsafe.Pointer)(slotElem) = nil
+			} else {
+				// Unlike keys, always clear the elem (even if
+				// it contains no pointers), as compound
+				// assignment operations depend on cleared
+				// deleted values. See
+				// https://go.dev/issue/25936.
+				typedmemclr(typ.Elem, slotElem)
+			}
 
 			// We only have 1 group, so it is OK to immediately
 			// reuse deleted slots.
diff --git a/src/internal/runtime/maps/map_test.go b/src/internal/runtime/maps/map_test.go
index cd40db8712..42db55c6a4 100644
--- a/src/internal/runtime/maps/map_test.go
+++ b/src/internal/runtime/maps/map_test.go
@@ -628,3 +628,74 @@ func TestMapZeroSizeSlot(t *testing.T) {
 		t.Errorf("elem address outside groups allocation; got %p want [%p, %p]", got, start, end)
 	}
 }
+
+func TestMapIndirect(t *testing.T) {
+	type big [abi.SwissMapMaxKeyBytes + abi.SwissMapMaxElemBytes]byte
+
+	m, typ := maps.NewTestMap[big, big](8)
+
+	key := big{}
+	elem := big{}
+	elem[0] = 128
+
+	for i := 0; i < 31; i++ {
+		key[0] += 1
+		elem[0] += 1
+		m.Put(typ, unsafe.Pointer(&key), unsafe.Pointer(&elem))
+
+		if maps.DebugLog {
+			fmt.Printf("After put %v: %v\n", key, m)
+		}
+	}
+
+	if m.Used() != 31 {
+		t.Errorf("Used() used got %d want 31", m.Used())
+	}
+
+	key = big{}
+	elem = big{}
+	elem[0] = 128
+
+	for i := 0; i < 31; i++ {
+		key[0] += 1
+		elem[0] += 1
+		got, ok := m.Get(typ, unsafe.Pointer(&key))
+		if !ok {
+			t.Errorf("Get(%v) got ok false want true", key)
+		}
+		gotElem := *(*big)(got)
+		if gotElem != elem {
+			t.Errorf("Get(%v) got elem %v want %v", key, gotElem, elem)
+		}
+	}
+}
+
+// Delete should clear element. See https://go.dev/issue/25936.
+func TestMapDeleteClear(t *testing.T) {
+	m, typ := maps.NewTestMap[int64, int64](8)
+
+	key := int64(0)
+	elem := int64(128)
+
+	m.Put(typ, unsafe.Pointer(&key), unsafe.Pointer(&elem))
+
+	if maps.DebugLog {
+		fmt.Printf("After put %d: %v\n", key, m)
+	}
+
+	got, ok := m.Get(typ, unsafe.Pointer(&key))
+	if !ok {
+		t.Errorf("Get(%d) got ok false want true", key)
+	}
+	gotElem := *(*int64)(got)
+	if gotElem != elem {
+		t.Errorf("Get(%d) got elem %d want %d", key, gotElem, elem)
+	}
+
+	m.Delete(typ, unsafe.Pointer(&key))
+
+	gotElem = *(*int64)(got)
+	if gotElem != 0 {
+		t.Errorf("Delete(%d) failed to clear element. got %d want 0", key, gotElem)
+	}
+}
diff --git a/src/internal/runtime/maps/runtime.go b/src/internal/runtime/maps/runtime.go
index 0d569de214..3d06f54f4d 100644
--- a/src/internal/runtime/maps/runtime.go
+++ b/src/internal/runtime/maps/runtime.go
@@ -25,3 +25,6 @@ func typedmemclr(typ *abi.Type, ptr unsafe.Pointer)
 
 //go:linkname newarray
 func newarray(typ *abi.Type, n int) unsafe.Pointer
+
+//go:linkname newobject
+func newobject(typ *abi.Type) unsafe.Pointer
diff --git a/src/internal/runtime/maps/runtime_swiss.go b/src/internal/runtime/maps/runtime_swiss.go
index 88042500bc..401c69a224 100644
--- a/src/internal/runtime/maps/runtime_swiss.go
+++ b/src/internal/runtime/maps/runtime_swiss.go
@@ -90,8 +90,15 @@ func runtime_mapaccess1(typ *abi.SwissMapType, m *Map, key unsafe.Pointer) unsaf
 			i := match.first()
 
 			slotKey := g.key(typ, i)
+			if typ.IndirectKey() {
+				slotKey = *((*unsafe.Pointer)(slotKey))
+			}
 			if typ.Key.Equal(key, slotKey) {
-				return g.elem(typ, i)
+				slotElem := g.elem(typ, i)
+				if typ.IndirectElem() {
+					slotElem = *((*unsafe.Pointer)(slotElem))
+				}
+				return slotElem
 			}
 			match = match.removeFirst()
 		}
@@ -176,12 +183,18 @@ outer:
 				i := match.first()
 
 				slotKey := g.key(typ, i)
+				if typ.IndirectKey() {
+					slotKey = *((*unsafe.Pointer)(slotKey))
+				}
 				if typ.Key.Equal(key, slotKey) {
 					if typ.NeedKeyUpdate() {
 						typedmemmove(typ.Key, slotKey, key)
 					}
 
 					slotElem = g.elem(typ, i)
+					if typ.IndirectElem() {
+						slotElem = *((*unsafe.Pointer)(slotElem))
+					}
 
 					t.checkInvariants(typ)
 					break outer
@@ -212,8 +225,19 @@ outer:
 				// If there is room left to grow, just insert the new entry.
 				if t.growthLeft > 0 {
 					slotKey := g.key(typ, i)
+					if typ.IndirectKey() {
+						kmem := newobject(typ.Key)
+						*(*unsafe.Pointer)(slotKey) = kmem
+						slotKey = kmem
+					}
 					typedmemmove(typ.Key, slotKey, key)
+
 					slotElem = g.elem(typ, i)
+					if typ.IndirectElem() {
+						emem := newobject(typ.Elem)
+						*(*unsafe.Pointer)(slotElem) = emem
+						slotElem = emem
+					}
 
 					g.ctrls().set(i, ctrl(h2(hash)))
 					t.growthLeft--
diff --git a/src/internal/runtime/maps/table.go b/src/internal/runtime/maps/table.go
index ac5271ea06..bb3006bfa2 100644
--- a/src/internal/runtime/maps/table.go
+++ b/src/internal/runtime/maps/table.go
@@ -207,8 +207,15 @@ func (t *table) getWithKey(typ *abi.SwissMapType, hash uintptr, key unsafe.Point
 			i := match.first()
 
 			slotKey := g.key(typ, i)
+			if typ.IndirectKey() {
+				slotKey = *((*unsafe.Pointer)(slotKey))
+			}
 			if typ.Key.Equal(key, slotKey) {
-				return slotKey, g.elem(typ, i), true
+				slotElem := g.elem(typ, i)
+				if typ.IndirectElem() {
+					slotElem = *((*unsafe.Pointer)(slotElem))
+				}
+				return slotKey, slotElem, true
 			}
 			match = match.removeFirst()
 		}
@@ -233,8 +240,15 @@ func (t *table) getWithoutKey(typ *abi.SwissMapType, hash uintptr, key unsafe.Po
 			i := match.first()
 
 			slotKey := g.key(typ, i)
+			if typ.IndirectKey() {
+				slotKey = *((*unsafe.Pointer)(slotKey))
+			}
 			if typ.Key.Equal(key, slotKey) {
-				return g.elem(typ, i), true
+				slotElem := g.elem(typ, i)
+				if typ.IndirectElem() {
+					slotElem = *((*unsafe.Pointer)(slotElem))
+				}
+				return slotElem, true
 			}
 			match = match.removeFirst()
 		}
@@ -272,12 +286,18 @@ func (t *table) PutSlot(typ *abi.SwissMapType, m *Map, hash uintptr, key unsafe.
 			i := match.first()
 
 			slotKey := g.key(typ, i)
+			if typ.IndirectKey() {
+				slotKey = *((*unsafe.Pointer)(slotKey))
+			}
 			if typ.Key.Equal(key, slotKey) {
 				if typ.NeedKeyUpdate() {
 					typedmemmove(typ.Key, slotKey, key)
 				}
 
 				slotElem := g.elem(typ, i)
+				if typ.IndirectElem() {
+					slotElem = *((*unsafe.Pointer)(slotElem))
+				}
 
 				t.checkInvariants(typ)
 				return slotElem, true
@@ -308,8 +328,19 @@ func (t *table) PutSlot(typ *abi.SwissMapType, m *Map, hash uintptr, key unsafe.
 			// If there is room left to grow, just insert the new entry.
 			if t.growthLeft > 0 {
 				slotKey := g.key(typ, i)
+				if typ.IndirectKey() {
+					kmem := newobject(typ.Key)
+					*(*unsafe.Pointer)(slotKey) = kmem
+					slotKey = kmem
+				}
 				typedmemmove(typ.Key, slotKey, key)
+
 				slotElem := g.elem(typ, i)
+				if typ.IndirectElem() {
+					emem := newobject(typ.Elem)
+					*(*unsafe.Pointer)(slotElem) = emem
+					slotElem = emem
+				}
 
 				g.ctrls().set(i, ctrl(h2(hash)))
 				t.growthLeft--
@@ -370,8 +401,19 @@ func (t *table) uncheckedPutSlot(typ *abi.SwissMapType, hash uintptr, key unsafe
 			i := match.first()
 
 			slotKey := g.key(typ, i)
+			if typ.IndirectKey() {
+				kmem := newobject(typ.Key)
+				*(*unsafe.Pointer)(slotKey) = kmem
+				slotKey = kmem
+			}
 			typedmemmove(typ.Key, slotKey, key)
+
 			slotElem := g.elem(typ, i)
+			if typ.IndirectElem() {
+				emem := newobject(typ.Elem)
+				*(*unsafe.Pointer)(slotElem) = emem
+				slotElem = emem
+			}
 
 			if g.ctrls().get(i) == ctrlEmpty {
 				t.growthLeft--
@@ -392,13 +434,38 @@ func (t *table) Delete(typ *abi.SwissMapType, m *Map, key unsafe.Pointer) {
 
 		for match != 0 {
 			i := match.first()
+
 			slotKey := g.key(typ, i)
+			origSlotKey := slotKey
+			if typ.IndirectKey() {
+				slotKey = *((*unsafe.Pointer)(slotKey))
+			}
+
 			if typ.Key.Equal(key, slotKey) {
 				t.used--
 				m.used--
 
-				typedmemclr(typ.Key, slotKey)
-				typedmemclr(typ.Elem, g.elem(typ, i))
+				if typ.IndirectKey() {
+					// Clearing the pointer is sufficient.
+					*(*unsafe.Pointer)(origSlotKey) = nil
+				} else if typ.Key.Pointers() {
+					// Only bothing clear the key if there
+					// are pointers in it.
+					typedmemclr(typ.Key, slotKey)
+				}
+
+				slotElem := g.elem(typ, i)
+				if typ.IndirectElem() {
+					// Clearing the pointer is sufficient.
+					*(*unsafe.Pointer)(slotElem) = nil
+				} else {
+					// Unlike keys, always clear the elem (even if
+					// it contains no pointers), as compound
+					// assignment operations depend on cleared
+					// deleted values. See
+					// https://go.dev/issue/25936.
+					typedmemclr(typ.Elem, slotElem)
+				}
 
 				// Only a full group can appear in the middle
 				// of a probe sequence (a group with at least
@@ -569,6 +636,9 @@ func (it *Iter) Next() {
 			}
 
 			key := g.key(it.typ, k)
+			if it.typ.IndirectKey() {
+				key = *((*unsafe.Pointer)(key))
+			}
 
 			// As below, if we have grown to a full map since Init,
 			// we continue to use the old group to decide the keys
@@ -583,6 +653,9 @@ func (it *Iter) Next() {
 					// See comment below.
 					if it.clearSeq == it.m.clearSeq && !it.typ.Key.Equal(key, key) {
 						elem = g.elem(it.typ, k)
+						if it.typ.IndirectElem() {
+							elem = *((*unsafe.Pointer)(elem))
+						}
 					} else {
 						continue
 					}
@@ -592,6 +665,9 @@ func (it *Iter) Next() {
 				}
 			} else {
 				elem = g.elem(it.typ, k)
+				if it.typ.IndirectElem() {
+					elem = *((*unsafe.Pointer)(elem))
+				}
 			}
 
 			it.entryIdx++
@@ -700,6 +776,9 @@ func (it *Iter) Next() {
 			}
 
 			key := g.key(it.typ, slotIdx)
+			if it.typ.IndirectKey() {
+				key = *((*unsafe.Pointer)(key))
+			}
 
 			// If the table has changed since the last
 			// call, then it has grown or split. In this
@@ -743,6 +822,9 @@ func (it *Iter) Next() {
 					// clear.
 					if it.clearSeq == it.m.clearSeq && !it.typ.Key.Equal(key, key) {
 						elem = g.elem(it.typ, slotIdx)
+						if it.typ.IndirectElem() {
+							elem = *((*unsafe.Pointer)(elem))
+						}
 					} else {
 						continue
 					}
@@ -752,6 +834,9 @@ func (it *Iter) Next() {
 				}
 			} else {
 				elem = g.elem(it.typ, slotIdx)
+				if it.typ.IndirectElem() {
+					elem = *((*unsafe.Pointer)(elem))
+				}
 			}
 
 			it.entryIdx++
@@ -852,8 +937,17 @@ func (t *table) split(typ *abi.SwissMapType, m *Map) {
 				// Empty or deleted
 				continue
 			}
+
 			key := g.key(typ, j)
+			if typ.IndirectKey() {
+				key = *((*unsafe.Pointer)(key))
+			}
+
 			elem := g.elem(typ, j)
+			if typ.IndirectElem() {
+				elem = *((*unsafe.Pointer)(elem))
+			}
+
 			hash := typ.Hasher(key, t.seed)
 			var newTable *table
 			if hash&mask == 0 {
@@ -861,6 +955,10 @@ func (t *table) split(typ *abi.SwissMapType, m *Map) {
 			} else {
 				newTable = right
 			}
+			// TODO(prattmic): For indirect key/elem, this is
+			// allocating new objects for key/elem. That is
+			// unnecessary; the new table could simply point to the
+			// existing object.
 			slotElem := newTable.uncheckedPutSlot(typ, hash, key)
 			typedmemmove(typ.Elem, slotElem, elem)
 			newTable.used++
@@ -885,9 +983,23 @@ func (t *table) grow(typ *abi.SwissMapType, m *Map, newCapacity uint16) {
 					// Empty or deleted
 					continue
 				}
+
 				key := g.key(typ, j)
+				if typ.IndirectKey() {
+					key = *((*unsafe.Pointer)(key))
+				}
+
 				elem := g.elem(typ, j)
+				if typ.IndirectElem() {
+					elem = *((*unsafe.Pointer)(elem))
+				}
+
 				hash := typ.Hasher(key, t.seed)
+
+				// TODO(prattmic): For indirect key/elem, this is
+				// allocating new objects for key/elem. That is
+				// unnecessary; the new table could simply point to the
+				// existing object.
 				slotElem := newTable.uncheckedPutSlot(typ, hash, key)
 				typedmemmove(typ.Elem, slotElem, elem)
 				newTable.used++
diff --git a/src/internal/runtime/maps/table_debug.go b/src/internal/runtime/maps/table_debug.go
index 27ae611ec3..345f1feb6e 100644
--- a/src/internal/runtime/maps/table_debug.go
+++ b/src/internal/runtime/maps/table_debug.go
@@ -35,6 +35,9 @@ func (t *table) checkInvariants(typ *abi.SwissMapType) {
 				used++
 
 				key := g.key(typ, j)
+				if typ.IndirectKey() {
+					key = *((*unsafe.Pointer)(key))
+				}
 
 				// Can't lookup keys that don't compare equal
 				// to themselves (e.g., NaN).