hashtriemap.go

     1  // Copyright 2024 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package sync
     6  
     7  import (
     8  	"internal/abi"
     9  	"internal/goarch"
    10  	"sync/atomic"
    11  	"unsafe"
    12  )
    13  
    14  // HashTrieMap is an implementation of a concurrent hash-trie. The implementation
    15  // is designed around frequent loads, but offers decent performance for stores
    16  // and deletes as well, especially if the map is larger. Its primary use-case is
    17  // the unique package, but can be used elsewhere as well.
    18  //
    19  // The zero HashTrieMap is empty and ready to use.
    20  // It must not be copied after first use.
    21  type HashTrieMap[K comparable, V any] struct {
    22  	inited   atomic.Uint32
    23  	initMu   Mutex
    24  	root     atomic.Pointer[indirect[K, V]]
    25  	keyHash  hashFunc
    26  	valEqual equalFunc
    27  	seed     uintptr
    28  }
    29  
    30  func (ht *HashTrieMap[K, V]) init() {
    31  	if ht.inited.Load() == 0 {
    32  		ht.initSlow()
    33  	}
    34  }
    35  
    36  //go:noinline
    37  func (ht *HashTrieMap[K, V]) initSlow() {
    38  	ht.initMu.Lock()
    39  	defer ht.initMu.Unlock()
    40  
    41  	if ht.inited.Load() != 0 {
    42  		// Someone got to it while we were waiting.
    43  		return
    44  	}
    45  
    46  	// Set up root node, derive the hash function for the key, and the
    47  	// equal function for the value, if any.
    48  	var m map[K]V
    49  	mapType := abi.TypeOf(m).MapType()
    50  	ht.root.Store(newIndirectNode[K, V](nil))
    51  	ht.keyHash = mapType.Hasher
    52  	ht.valEqual = mapType.Elem.Equal
    53  	ht.seed = uintptr(runtime_rand())
    54  
    55  	ht.inited.Store(1)
    56  }
    57  
    58  type hashFunc func(unsafe.Pointer, uintptr) uintptr
    59  type equalFunc func(unsafe.Pointer, unsafe.Pointer) bool
    60  
    61  // Load returns the value stored in the map for a key, or nil if no
    62  // value is present.
    63  // The ok result indicates whether value was found in the map.
    64  func (ht *HashTrieMap[K, V]) Load(key K) (value V, ok bool) {
    65  	ht.init()
    66  	hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)
    67  
    68  	i := ht.root.Load()
    69  	hashShift := 8 * goarch.PtrSize
    70  	for hashShift != 0 {
    71  		hashShift -= nChildrenLog2
    72  
    73  		n := i.children[(hash>>hashShift)&nChildrenMask].Load()
    74  		if n == nil {
    75  			return *new(V), false
    76  		}
    77  		if n.isEntry {
    78  			return n.entry().lookup(key)
    79  		}
    80  		i = n.indirect()
    81  	}
    82  	panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
    83  }
    84  
    85  // LoadOrStore returns the existing value for the key if present.
    86  // Otherwise, it stores and returns the given value.
    87  // The loaded result is true if the value was loaded, false if stored.
    88  func (ht *HashTrieMap[K, V]) LoadOrStore(key K, value V) (result V, loaded bool) {
    89  	ht.init()
    90  	hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)
    91  	var i *indirect[K, V]
    92  	var hashShift uint
    93  	var slot *atomic.Pointer[node[K, V]]
    94  	var n *node[K, V]
    95  	for {
    96  		// Find the key or a candidate location for insertion.
    97  		i = ht.root.Load()
    98  		hashShift = 8 * goarch.PtrSize
    99  		haveInsertPoint := false
   100  		for hashShift != 0 {
   101  			hashShift -= nChildrenLog2
   102  
   103  			slot = &i.children[(hash>>hashShift)&nChildrenMask]
   104  			n = slot.Load()
   105  			if n == nil {
   106  				// We found a nil slot which is a candidate for insertion.
   107  				haveInsertPoint = true
   108  				break
   109  			}
   110  			if n.isEntry {
   111  				// We found an existing entry, which is as far as we can go.
   112  				// If it stays this way, we'll have to replace it with an
   113  				// indirect node.
   114  				if v, ok := n.entry().lookup(key); ok {
   115  					return v, true
   116  				}
   117  				haveInsertPoint = true
   118  				break
   119  			}
   120  			i = n.indirect()
   121  		}
   122  		if !haveInsertPoint {
   123  			panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
   124  		}
   125  
   126  		// Grab the lock and double-check what we saw.
   127  		i.mu.Lock()
   128  		n = slot.Load()
   129  		if (n == nil || n.isEntry) && !i.dead.Load() {
   130  			// What we saw is still true, so we can continue with the insert.
   131  			break
   132  		}
   133  		// We have to start over.
   134  		i.mu.Unlock()
   135  	}
   136  	// N.B. This lock is held from when we broke out of the outer loop above.
   137  	// We specifically break this out so that we can use defer here safely.
   138  	// One option is to break this out into a new function instead, but
   139  	// there's so much local iteration state used below that this turns out
   140  	// to be cleaner.
   141  	defer i.mu.Unlock()
   142  
   143  	var oldEntry *entry[K, V]
   144  	if n != nil {
   145  		oldEntry = n.entry()
   146  		if v, ok := oldEntry.lookup(key); ok {
   147  			// Easy case: by loading again, it turns out exactly what we wanted is here!
   148  			return v, true
   149  		}
   150  	}
   151  	newEntry := newEntryNode(key, value)
   152  	if oldEntry == nil {
   153  		// Easy case: create a new entry and store it.
   154  		slot.Store(&newEntry.node)
   155  	} else {
   156  		// We possibly need to expand the entry already there into one or more new nodes.
   157  		//
   158  		// Publish the node last, which will make both oldEntry and newEntry visible. We
   159  		// don't want readers to be able to observe that oldEntry isn't in the tree.
   160  		slot.Store(ht.expand(oldEntry, newEntry, hash, hashShift, i))
   161  	}
   162  	return value, false
   163  }
   164  
   165  // expand takes oldEntry and newEntry whose hashes conflict from bit 64 down to hashShift and
   166  // produces a subtree of indirect nodes to hold the two new entries.
   167  func (ht *HashTrieMap[K, V]) expand(oldEntry, newEntry *entry[K, V], newHash uintptr, hashShift uint, parent *indirect[K, V]) *node[K, V] {
   168  	// Check for a hash collision.
   169  	oldHash := ht.keyHash(unsafe.Pointer(&oldEntry.key), ht.seed)
   170  	if oldHash == newHash {
   171  		// Store the old entry in the new entry's overflow list, then store
   172  		// the new entry.
   173  		newEntry.overflow.Store(oldEntry)
   174  		return &newEntry.node
   175  	}
   176  	// We have to add an indirect node. Worse still, we may need to add more than one.
   177  	newIndirect := newIndirectNode(parent)
   178  	top := newIndirect
   179  	for {
   180  		if hashShift == 0 {
   181  			panic("internal/concurrent.HashMapTrie: ran out of hash bits while inserting")
   182  		}
   183  		hashShift -= nChildrenLog2 // hashShift is for the level parent is at. We need to go deeper.
   184  		oi := (oldHash >> hashShift) & nChildrenMask
   185  		ni := (newHash >> hashShift) & nChildrenMask
   186  		if oi != ni {
   187  			newIndirect.children[oi].Store(&oldEntry.node)
   188  			newIndirect.children[ni].Store(&newEntry.node)
   189  			break
   190  		}
   191  		nextIndirect := newIndirectNode(newIndirect)
   192  		newIndirect.children[oi].Store(&nextIndirect.node)
   193  		newIndirect = nextIndirect
   194  	}
   195  	return &top.node
   196  }
   197  
   198  // Store sets the value for a key.
   199  func (ht *HashTrieMap[K, V]) Store(key K, old V) {
   200  	_, _ = ht.Swap(key, old)
   201  }
   202  
   203  // Swap swaps the value for a key and returns the previous value if any.
   204  // The loaded result reports whether the key was present.
   205  func (ht *HashTrieMap[K, V]) Swap(key K, new V) (previous V, loaded bool) {
   206  	ht.init()
   207  	hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)
   208  	var i *indirect[K, V]
   209  	var hashShift uint
   210  	var slot *atomic.Pointer[node[K, V]]
   211  	var n *node[K, V]
   212  	for {
   213  		// Find the key or a candidate location for insertion.
   214  		i = ht.root.Load()
   215  		hashShift = 8 * goarch.PtrSize
   216  		haveInsertPoint := false
   217  		for hashShift != 0 {
   218  			hashShift -= nChildrenLog2
   219  
   220  			slot = &i.children[(hash>>hashShift)&nChildrenMask]
   221  			n = slot.Load()
   222  			if n == nil {
   223  				// We found a nil slot which is a candidate for insertion,
   224  				// or an existing entry that we'll replace.
   225  				haveInsertPoint = true
   226  				break
   227  			}
   228  			if n.isEntry {
   229  				// Swap if the keys compare.
   230  				old, swapped := n.entry().swap(key, new)
   231  				if swapped {
   232  					return old, true
   233  				}
   234  				// If we fail, that means we should try to insert.
   235  				haveInsertPoint = true
   236  				break
   237  			}
   238  			i = n.indirect()
   239  		}
   240  		if !haveInsertPoint {
   241  			panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
   242  		}
   243  
   244  		// Grab the lock and double-check what we saw.
   245  		i.mu.Lock()
   246  		n = slot.Load()
   247  		if (n == nil || n.isEntry) && !i.dead.Load() {
   248  			// What we saw is still true, so we can continue with the insert.
   249  			break
   250  		}
   251  		// We have to start over.
   252  		i.mu.Unlock()
   253  	}
   254  	// N.B. This lock is held from when we broke out of the outer loop above.
   255  	// We specifically break this out so that we can use defer here safely.
   256  	// One option is to break this out into a new function instead, but
   257  	// there's so much local iteration state used below that this turns out
   258  	// to be cleaner.
   259  	defer i.mu.Unlock()
   260  
   261  	var zero V
   262  	var oldEntry *entry[K, V]
   263  	if n != nil {
   264  		// Between before and now, something got inserted. Swap if the keys compare.
   265  		oldEntry = n.entry()
   266  		old, swapped := oldEntry.swap(key, new)
   267  		if swapped {
   268  			return old, true
   269  		}
   270  	}
   271  	// The keys didn't compare, so we're doing an insertion.
   272  	newEntry := newEntryNode(key, new)
   273  	if oldEntry == nil {
   274  		// Easy case: create a new entry and store it.
   275  		slot.Store(&newEntry.node)
   276  	} else {
   277  		// We possibly need to expand the entry already there into one or more new nodes.
   278  		//
   279  		// Publish the node last, which will make both oldEntry and newEntry visible. We
   280  		// don't want readers to be able to observe that oldEntry isn't in the tree.
   281  		slot.Store(ht.expand(oldEntry, newEntry, hash, hashShift, i))
   282  	}
   283  	return zero, false
   284  }
   285  
   286  // CompareAndSwap swaps the old and new values for key
   287  // if the value stored in the map is equal to old.
   288  // The value type must be of a comparable type, otherwise CompareAndSwap will panic.
   289  func (ht *HashTrieMap[K, V]) CompareAndSwap(key K, old, new V) (swapped bool) {
   290  	ht.init()
   291  	if ht.valEqual == nil {
   292  		panic("called CompareAndSwap when value is not of comparable type")
   293  	}
   294  	hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)
   295  	for {
   296  		// Find the key or return if it's not there.
   297  		i := ht.root.Load()
   298  		hashShift := 8 * goarch.PtrSize
   299  		found := false
   300  		for hashShift != 0 {
   301  			hashShift -= nChildrenLog2
   302  
   303  			slot := &i.children[(hash>>hashShift)&nChildrenMask]
   304  			n := slot.Load()
   305  			if n == nil {
   306  				// Nothing to compare with. Give up.
   307  				return false
   308  			}
   309  			if n.isEntry {
   310  				// We found an entry. Try to compare and swap directly.
   311  				return n.entry().compareAndSwap(key, old, new, ht.valEqual)
   312  			}
   313  			i = n.indirect()
   314  		}
   315  		if !found {
   316  			panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
   317  		}
   318  	}
   319  }
   320  
   321  // LoadAndDelete deletes the value for a key, returning the previous value if any.
   322  // The loaded result reports whether the key was present.
   323  func (ht *HashTrieMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
   324  	ht.init()
   325  	hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)
   326  
   327  	// Find a node with the key and compare with it. n != nil if we found the node.
   328  	i, hashShift, slot, n := ht.find(key, hash, nil, *new(V))
   329  	if n == nil {
   330  		if i != nil {
   331  			i.mu.Unlock()
   332  		}
   333  		return *new(V), false
   334  	}
   335  
   336  	// Try to delete the entry.
   337  	v, e, loaded := n.entry().loadAndDelete(key)
   338  	if !loaded {
   339  		// Nothing was actually deleted, which means the node is no longer there.
   340  		i.mu.Unlock()
   341  		return *new(V), false
   342  	}
   343  	if e != nil {
   344  		// We didn't actually delete the whole entry, just one entry in the chain.
   345  		// Nothing else to do, since the parent is definitely not empty.
   346  		slot.Store(&e.node)
   347  		i.mu.Unlock()
   348  		return v, true
   349  	}
   350  	// Delete the entry.
   351  	slot.Store(nil)
   352  
   353  	// Check if the node is now empty (and isn't the root), and delete it if able.
   354  	for i.parent != nil && i.empty() {
   355  		if hashShift == 8*goarch.PtrSize {
   356  			panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
   357  		}
   358  		hashShift += nChildrenLog2
   359  
   360  		// Delete the current node in the parent.
   361  		parent := i.parent
   362  		parent.mu.Lock()
   363  		i.dead.Store(true)
   364  		parent.children[(hash>>hashShift)&nChildrenMask].Store(nil)
   365  		i.mu.Unlock()
   366  		i = parent
   367  	}
   368  	i.mu.Unlock()
   369  	return v, true
   370  }
   371  
   372  // Delete deletes the value for a key.
   373  func (ht *HashTrieMap[K, V]) Delete(key K) {
   374  	_, _ = ht.LoadAndDelete(key)
   375  }
   376  
   377  // CompareAndDelete deletes the entry for key if its value is equal to old.
   378  // The value type must be comparable, otherwise this CompareAndDelete will panic.
   379  //
   380  // If there is no current value for key in the map, CompareAndDelete returns false
   381  // (even if the old value is the nil interface value).
   382  func (ht *HashTrieMap[K, V]) CompareAndDelete(key K, old V) (deleted bool) {
   383  	ht.init()
   384  	if ht.valEqual == nil {
   385  		panic("called CompareAndDelete when value is not of comparable type")
   386  	}
   387  	hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)
   388  
   389  	// Find a node with the key. n != nil if we found the node.
   390  	i, hashShift, slot, n := ht.find(key, hash, nil, *new(V))
   391  	if n == nil {
   392  		if i != nil {
   393  			i.mu.Unlock()
   394  		}
   395  		return false
   396  	}
   397  
   398  	// Try to delete the entry.
   399  	e, deleted := n.entry().compareAndDelete(key, old, ht.valEqual)
   400  	if !deleted {
   401  		// Nothing was actually deleted, which means the node is no longer there.
   402  		i.mu.Unlock()
   403  		return false
   404  	}
   405  	if e != nil {
   406  		// We didn't actually delete the whole entry, just one entry in the chain.
   407  		// Nothing else to do, since the parent is definitely not empty.
   408  		slot.Store(&e.node)
   409  		i.mu.Unlock()
   410  		return true
   411  	}
   412  	// Delete the entry.
   413  	slot.Store(nil)
   414  
   415  	// Check if the node is now empty (and isn't the root), and delete it if able.
   416  	for i.parent != nil && i.empty() {
   417  		if hashShift == 8*goarch.PtrSize {
   418  			panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
   419  		}
   420  		hashShift += nChildrenLog2
   421  
   422  		// Delete the current node in the parent.
   423  		parent := i.parent
   424  		parent.mu.Lock()
   425  		i.dead.Store(true)
   426  		parent.children[(hash>>hashShift)&nChildrenMask].Store(nil)
   427  		i.mu.Unlock()
   428  		i = parent
   429  	}
   430  	i.mu.Unlock()
   431  	return true
   432  }
   433  
   434  // find searches the tree for a node that contains key (hash must be the hash of key).
   435  // If valEqual != nil, then it will also enforce that the values are equal as well.
   436  //
   437  // Returns a non-nil node, which will always be an entry, if found.
   438  //
   439  // If i != nil then i.mu is locked, and it is the caller's responsibility to unlock it.
   440  func (ht *HashTrieMap[K, V]) find(key K, hash uintptr, valEqual equalFunc, value V) (i *indirect[K, V], hashShift uint, slot *atomic.Pointer[node[K, V]], n *node[K, V]) {
   441  	for {
   442  		// Find the key or return if it's not there.
   443  		i = ht.root.Load()
   444  		hashShift = 8 * goarch.PtrSize
   445  		found := false
   446  		for hashShift != 0 {
   447  			hashShift -= nChildrenLog2
   448  
   449  			slot = &i.children[(hash>>hashShift)&nChildrenMask]
   450  			n = slot.Load()
   451  			if n == nil {
   452  				// Nothing to compare with. Give up.
   453  				i = nil
   454  				return
   455  			}
   456  			if n.isEntry {
   457  				// We found an entry. Check if it matches.
   458  				if _, ok := n.entry().lookupWithValue(key, value, valEqual); !ok {
   459  					// No match, comparison failed.
   460  					i = nil
   461  					n = nil
   462  					return
   463  				}
   464  				// We've got a match. Prepare to perform an operation on the key.
   465  				found = true
   466  				break
   467  			}
   468  			i = n.indirect()
   469  		}
   470  		if !found {
   471  			panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
   472  		}
   473  
   474  		// Grab the lock and double-check what we saw.
   475  		i.mu.Lock()
   476  		n = slot.Load()
   477  		if !i.dead.Load() && (n == nil || n.isEntry) {
   478  			// Either we've got a valid node or the node is now nil under the lock.
   479  			// In either case, we're done here.
   480  			return
   481  		}
   482  		// We have to start over.
   483  		i.mu.Unlock()
   484  	}
   485  }
   486  
   487  // All returns an iterator over each key and value present in the map.
   488  //
   489  // The iterator does not necessarily correspond to any consistent snapshot of the
   490  // HashTrieMap's contents: no key will be visited more than once, but if the value
   491  // for any key is stored or deleted concurrently (including by yield), the iterator
   492  // may reflect any mapping for that key from any point during iteration. The iterator
   493  // does not block other methods on the receiver; even yield itself may call any
   494  // method on the HashTrieMap.
   495  func (ht *HashTrieMap[K, V]) All() func(yield func(K, V) bool) {
   496  	ht.init()
   497  	return func(yield func(key K, value V) bool) {
   498  		ht.iter(ht.root.Load(), yield)
   499  	}
   500  }
   501  
   502  // Range calls f sequentially for each key and value present in the map.
   503  // If f returns false, range stops the iteration.
   504  //
   505  // This exists for compatibility with sync.Map; All should be preferred.
   506  // It provides the same guarantees as sync.Map, and All.
   507  func (ht *HashTrieMap[K, V]) Range(yield func(K, V) bool) {
   508  	ht.init()
   509  	ht.iter(ht.root.Load(), yield)
   510  }
   511  
   512  func (ht *HashTrieMap[K, V]) iter(i *indirect[K, V], yield func(key K, value V) bool) bool {
   513  	for j := range i.children {
   514  		n := i.children[j].Load()
   515  		if n == nil {
   516  			continue
   517  		}
   518  		if !n.isEntry {
   519  			if !ht.iter(n.indirect(), yield) {
   520  				return false
   521  			}
   522  			continue
   523  		}
   524  		e := n.entry()
   525  		for e != nil {
   526  			if !yield(e.key, *e.value.Load()) {
   527  				return false
   528  			}
   529  			e = e.overflow.Load()
   530  		}
   531  	}
   532  	return true
   533  }
   534  
   535  // Clear deletes all the entries, resulting in an empty HashTrieMap.
   536  func (ht *HashTrieMap[K, V]) Clear() {
   537  	ht.init()
   538  
   539  	// It's sufficient to just drop the root on the floor, but the root
   540  	// must always be non-nil.
   541  	ht.root.Store(newIndirectNode[K, V](nil))
   542  }
   543  
   544  const (
   545  	// 16 children. This seems to be the sweet spot for
   546  	// load performance: any smaller and we lose out on
   547  	// 50% or more in CPU performance. Any larger and the
   548  	// returns are minuscule (~1% improvement for 32 children).
   549  	nChildrenLog2 = 4
   550  	nChildren     = 1 << nChildrenLog2
   551  	nChildrenMask = nChildren - 1
   552  )
   553  
   554  // indirect is an internal node in the hash-trie.
   555  type indirect[K comparable, V any] struct {
   556  	node[K, V]
   557  	dead     atomic.Bool
   558  	mu       Mutex // Protects mutation to children and any children that are entry nodes.
   559  	parent   *indirect[K, V]
   560  	children [nChildren]atomic.Pointer[node[K, V]]
   561  }
   562  
   563  func newIndirectNode[K comparable, V any](parent *indirect[K, V]) *indirect[K, V] {
   564  	return &indirect[K, V]{node: node[K, V]{isEntry: false}, parent: parent}
   565  }
   566  
   567  func (i *indirect[K, V]) empty() bool {
   568  	nc := 0
   569  	for j := range i.children {
   570  		if i.children[j].Load() != nil {
   571  			nc++
   572  		}
   573  	}
   574  	return nc == 0
   575  }
   576  
   577  // entry is a leaf node in the hash-trie.
   578  type entry[K comparable, V any] struct {
   579  	node[K, V]
   580  	overflow atomic.Pointer[entry[K, V]] // Overflow for hash collisions.
   581  	key      K
   582  	value    atomic.Pointer[V]
   583  }
   584  
   585  func newEntryNode[K comparable, V any](key K, value V) *entry[K, V] {
   586  	e := &entry[K, V]{
   587  		node: node[K, V]{isEntry: true},
   588  		key:  key,
   589  	}
   590  	e.value.Store(&value)
   591  	return e
   592  }
   593  
   594  func (e *entry[K, V]) lookup(key K) (V, bool) {
   595  	for e != nil {
   596  		if e.key == key {
   597  			return *e.value.Load(), true
   598  		}
   599  		e = e.overflow.Load()
   600  	}
   601  	return *new(V), false
   602  }
   603  
   604  func (e *entry[K, V]) lookupWithValue(key K, value V, valEqual equalFunc) (V, bool) {
   605  	for e != nil {
   606  		oldp := e.value.Load()
   607  		if e.key == key && (valEqual == nil || valEqual(unsafe.Pointer(oldp), abi.NoEscape(unsafe.Pointer(&value)))) {
   608  			return *oldp, true
   609  		}
   610  		e = e.overflow.Load()
   611  	}
   612  	return *new(V), false
   613  }
   614  
   615  // swap replaces a value in the overflow chain if keys compare equal.
   616  // Returns the old value, and whether or not anything was swapped.
   617  //
   618  // swap must be called under the mutex of the indirect node which e is a child of.
   619  func (head *entry[K, V]) swap(key K, newv V) (V, bool) {
   620  	if head.key == key {
   621  		vp := new(V)
   622  		*vp = newv
   623  		oldp := head.value.Swap(vp)
   624  		return *oldp, true
   625  	}
   626  	i := &head.overflow
   627  	e := i.Load()
   628  	for e != nil {
   629  		if e.key == key {
   630  			vp := new(V)
   631  			*vp = newv
   632  			oldp := e.value.Swap(vp)
   633  			return *oldp, true
   634  		}
   635  		i = &e.overflow
   636  		e = e.overflow.Load()
   637  	}
   638  	var zero V
   639  	return zero, false
   640  }
   641  
   642  // compareAndSwap replaces a value for a matching key and existing value in the overflow chain.
   643  // Returns whether or not anything was swapped.
   644  //
   645  // compareAndSwap must be called under the mutex of the indirect node which e is a child of.
   646  func (head *entry[K, V]) compareAndSwap(key K, oldv, newv V, valEqual equalFunc) bool {
   647  	var vbox *V
   648  outerLoop:
   649  	for {
   650  		oldvp := head.value.Load()
   651  		if head.key == key && valEqual(unsafe.Pointer(oldvp), abi.NoEscape(unsafe.Pointer(&oldv))) {
   652  			// Return the new head of the list.
   653  			if vbox == nil {
   654  				// Delay explicit creation of a new value to hold newv. If we just pass &newv
   655  				// to CompareAndSwap, then newv will unconditionally escape, even if the CAS fails.
   656  				vbox = new(V)
   657  				*vbox = newv
   658  			}
   659  			if head.value.CompareAndSwap(oldvp, vbox) {
   660  				return true
   661  			}
   662  			// We need to restart from the head of the overflow list in case, due to a removal, a node
   663  			// is moved up the list and we miss it.
   664  			continue outerLoop
   665  		}
   666  		i := &head.overflow
   667  		e := i.Load()
   668  		for e != nil {
   669  			oldvp := e.value.Load()
   670  			if e.key == key && valEqual(unsafe.Pointer(oldvp), abi.NoEscape(unsafe.Pointer(&oldv))) {
   671  				if vbox == nil {
   672  					// Delay explicit creation of a new value to hold newv. If we just pass &newv
   673  					// to CompareAndSwap, then newv will unconditionally escape, even if the CAS fails.
   674  					vbox = new(V)
   675  					*vbox = newv
   676  				}
   677  				if e.value.CompareAndSwap(oldvp, vbox) {
   678  					return true
   679  				}
   680  				continue outerLoop
   681  			}
   682  			i = &e.overflow
   683  			e = e.overflow.Load()
   684  		}
   685  		return false
   686  	}
   687  }
   688  
   689  // loadAndDelete deletes an entry in the overflow chain by key. Returns the value for the key, the new
   690  // entry chain and whether or not anything was loaded (and deleted).
   691  //
   692  // loadAndDelete must be called under the mutex of the indirect node which e is a child of.
   693  func (head *entry[K, V]) loadAndDelete(key K) (V, *entry[K, V], bool) {
   694  	if head.key == key {
   695  		// Drop the head of the list.
   696  		return *head.value.Load(), head.overflow.Load(), true
   697  	}
   698  	i := &head.overflow
   699  	e := i.Load()
   700  	for e != nil {
   701  		if e.key == key {
   702  			i.Store(e.overflow.Load())
   703  			return *e.value.Load(), head, true
   704  		}
   705  		i = &e.overflow
   706  		e = e.overflow.Load()
   707  	}
   708  	return *new(V), head, false
   709  }
   710  
   711  // compareAndDelete deletes an entry in the overflow chain if both the key and value compare
   712  // equal. Returns the new entry chain and whether or not anything was deleted.
   713  //
   714  // compareAndDelete must be called under the mutex of the indirect node which e is a child of.
   715  func (head *entry[K, V]) compareAndDelete(key K, value V, valEqual equalFunc) (*entry[K, V], bool) {
   716  	if head.key == key && valEqual(unsafe.Pointer(head.value.Load()), abi.NoEscape(unsafe.Pointer(&value))) {
   717  		// Drop the head of the list.
   718  		return head.overflow.Load(), true
   719  	}
   720  	i := &head.overflow
   721  	e := i.Load()
   722  	for e != nil {
   723  		if e.key == key && valEqual(unsafe.Pointer(e.value.Load()), abi.NoEscape(unsafe.Pointer(&value))) {
   724  			i.Store(e.overflow.Load())
   725  			return head, true
   726  		}
   727  		i = &e.overflow
   728  		e = e.overflow.Load()
   729  	}
   730  	return head, false
   731  }
   732  
   733  // node is the header for a node. It's polymorphic and
   734  // is actually either an entry or an indirect.
   735  type node[K comparable, V any] struct {
   736  	isEntry bool
   737  }
   738  
   739  func (n *node[K, V]) entry() *entry[K, V] {
   740  	if !n.isEntry {
   741  		panic("called entry on non-entry node")
   742  	}
   743  	return (*entry[K, V])(unsafe.Pointer(n))
   744  }
   745  
   746  func (n *node[K, V]) indirect() *indirect[K, V] {
   747  	if n.isEntry {
   748  		panic("called indirect on entry node")
   749  	}
   750  	return (*indirect[K, V])(unsafe.Pointer(n))
   751  }
   752  
   753  // Pull in runtime.rand so that we don't need to take a dependency
   754  // on math/rand/v2.
   755  //
   756  //go:linkname runtime_rand runtime.rand
   757  func runtime_rand() uint64
   758
View as plain text