internal/runtime/maps: reuse deleted slots on insert

While walking the probe sequence, Put keeps track of the first deleted
slot it encountered. If it reaches the end of the probe sequence without
finding a match, then it will prefer to use the deleted slot rather than
a new empty slot.

For #54766.

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

src/internal/runtime/maps/export_test.go

@@ -36,12 +36,48 @@ func (m *Map) GroupCount() uint64 {
 	return n
 }
 
+// Return a key from a group containing no empty slots, or nil if there are no
+// full groups.
+//
+// Also returns nil if a group is full but contains entirely deleted slots.
+func (m *Map) KeyFromFullGroup() unsafe.Pointer {
+	var lastTab *table
+	for _, t := range m.directory {
+		if t == lastTab {
+			continue
+		}
+		lastTab = t
+
+		for i := uint64(0); i <= t.groups.lengthMask; i++ {
+			g := t.groups.group(i)
+			match := g.ctrls().matchEmpty()
+			if match != 0 {
+				continue
+			}
+
+			// All full or deleted slots.
+			for j := uint32(0); j < abi.SwissMapGroupSlots; j++ {
+				if g.ctrls().get(j) == ctrlDeleted {
+					continue
+				}
+				return g.key(j)
+			}
+		}
+	}
+
+	return nil
+}
+
 func (m *Map) TableFor(key unsafe.Pointer) *table {
 	hash := m.typ.Hasher(key, m.seed)
 	idx := m.directoryIndex(hash)
 	return m.directory[idx]
 }
 
+func (t *table) GrowthLeft() uint64 {
+	return uint64(t.growthLeft)
+}
+
 // Returns the start address of the groups array.
 func (t *table) GroupsStart() unsafe.Pointer {
 	return t.groups.data

src/internal/runtime/maps/map.go

@@ -77,9 +77,10 @@ import (
 // slot as empty, as there could be more slots used later in a probe sequence
 // and this deletion would cause probing to stop too early. Instead, we mark
 // such slots as "deleted" with a tombstone. If the group still has an empty
-// slot, we don't need a tombstone and directly mark the slot empty. Currently,
+// slot, we don't need a tombstone and directly mark the slot empty. Insert
-// tombstone are only cleared during grow, as an in-place cleanup complicates
+// prioritizes reuse of tombstones over filling an empty slots. Otherwise,
-// iteration.
+// tombstones are only completely cleared during grow, as an in-place cleanup
+// complicates iteration.
 //
 // Growth
 //

src/internal/runtime/maps/map_test.go

@@ -213,6 +213,68 @@ func TestTableKeyUpdate(t *testing.T) {
 	}
 }
 
+// Put should reuse a deleted slot rather than consuming an empty slot.
+func TestTablePutDelete(t *testing.T) {
+	// Put will reuse the first deleted slot it encounters.
+	//
+	// This is awkward to test because Delete will only install ctrlDeleted
+	// if the group is full, otherwise it goes straight to empty.
+	//
+	// So first we must add to the table continuously until we happen to
+	// fill a group.
+
+	m, _ := maps.NewTestMap[uint32, uint32](8)
+
+	key := uint32(0)
+	elem := uint32(256 + 0)
+
+	for {
+		key += 1
+		elem += 1
+
+		m.Put(unsafe.Pointer(&key), unsafe.Pointer(&elem))
+
+		// Normally a Put that fills a group would fill it with the
+		// inserted key, so why search the whole map for a potentially
+		// different key in a full group?
+		//
+		// Put may grow/split a table. Initial construction of the new
+		// table(s) could result in a full group consisting of
+		// arbitrary keys.
+		fullKeyPtr := m.KeyFromFullGroup()
+		if fullKeyPtr != nil {
+			// Found a full group.
+			key = *(*uint32)(fullKeyPtr)
+			elem = 256 + key
+			break
+		}
+	}
+
+	// Key is in a full group. Deleting it will result in a ctrlDeleted
+	// slot.
+	m.Delete(unsafe.Pointer(&key))
+
+	// Re-insert key. This should reuse the deleted slot rather than
+	// consuming space.
+	tabWant := m.TableFor(unsafe.Pointer(&key))
+	growthLeftWant := tabWant.GrowthLeft()
+
+	m.Put(unsafe.Pointer(&key), unsafe.Pointer(&elem))
+
+	tabGot := m.TableFor(unsafe.Pointer(&key))
+	growthLeftGot := tabGot.GrowthLeft()
+
+	if tabGot != tabWant {
+		// There shouldn't be a grow, as replacing a deleted slot
+		// doesn't require more space.
+		t.Errorf("Put(%d) grew table got %v want %v map %v", key, tabGot, tabWant, m)
+	}
+
+	if growthLeftGot != growthLeftWant {
+		t.Errorf("GrowthLeft got %d want %d: map %v tab %v", growthLeftGot, growthLeftWant, m, tabGot)
+	}
+}
+
 func TestTableIteration(t *testing.T) {
 	m, _ := maps.NewTestMap[uint32, uint64](8)
 

src/internal/runtime/maps/table.go

@@ -161,8 +161,6 @@ func (t *table) Get(key unsafe.Pointer) (unsafe.Pointer, bool) {
 	// TODO(prattmic): We could avoid hashing in a variety of special
 	// cases.
 	//
-	// - One group maps with simple keys could iterate over all keys and
-	//   compare them directly.
 	// - One entry maps could just directly compare the single entry
 	//   without hashing.
 	// - String keys could do quick checks of a few bytes before hashing.
@@ -243,6 +241,11 @@ func (t *table) getWithKey(hash uintptr, key unsafe.Pointer) (unsafe.Pointer, un
 func (t *table) PutSlot(m *Map, hash uintptr, key unsafe.Pointer) (unsafe.Pointer, bool) {
 	seq := makeProbeSeq(h1(hash), t.groups.lengthMask)
 
+	// As we look for a match, keep track of the first deleted slot we
+	// find, which we'll use to insert the new entry if necessary.
+	var firstDeletedGroup groupReference
+	var firstDeletedSlot uint32
+
 	for ; ; seq = seq.next() {
 		g := t.groups.group(seq.offset)
 		match := g.ctrls().matchH2(h2(hash))
@@ -265,15 +268,28 @@ func (t *table) PutSlot(m *Map, hash uintptr, key unsafe.Pointer) (unsafe.Pointe
 			match = match.removeFirst()
 		}
 
+		// No existing slot for this key in this group. Is this the end
+		// of the probe sequence?
 		match = g.ctrls().matchEmpty()
 		if match != 0 {
 			// Finding an empty slot means we've reached the end of
 			// the probe sequence.
 
+			var i uint32
+
+			// If we found a deleted slot along the way, we can
+			// replace it without consuming growthLeft.
+			if firstDeletedGroup.data != nil {
+				g = firstDeletedGroup
+				i = firstDeletedSlot
+				t.growthLeft++ // will be decremented below to become a no-op.
+			} else {
+				// Otherwise, use the empty slot.
+				i = match.first()
+			}
+
 			// If there is room left to grow, just insert the new entry.
 			if t.growthLeft > 0 {
-				i := match.first()
-
 				slotKey := g.key(i)
 				typedmemmove(t.typ.Key, slotKey, key)
 				slotElem := g.elem(i)
@@ -287,24 +303,24 @@ func (t *table) PutSlot(m *Map, hash uintptr, key unsafe.Pointer) (unsafe.Pointe
 				return slotElem, true
 			}
 
-			// TODO(prattmic): While searching the probe sequence,
-			// we may have passed deleted slots which we could use
-			// for this entry.
-			//
-			// At the moment, we leave this behind for
-			// rehash to free up.
-			//
-			// cockroachlabs/swiss restarts search of the probe
-			// sequence for a deleted slot.
-			//
-			// TODO(go.dev/issue/54766): We want this optimization
-			// back. We could search for the first deleted slot
-			// during the main search, but only use it if we don't
-			// find an existing entry.
-
 			t.rehash(m)
 			return nil, false
 		}
+
+		// No empty slots in this group. Check for a deleted
+		// slot, which we'll use if we don't find a match later
+		// in the probe sequence.
+		//
+		// We only need to remember a single deleted slot.
+		if firstDeletedGroup.data == nil {
+			// Since we already checked for empty slots
+			// above, matches here must be deleted slots.
+			match = g.ctrls().matchEmptyOrDeleted()
+			if match != 0 {
+				firstDeletedGroup = g
+				firstDeletedSlot = match.first()
+			}
+		}
 	}
 }