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
5package unique
6
7import (
8	"internal/abi"
9	"internal/concurrent"
10	"internal/weak"
11	"runtime"
12	"sync"
13	_ "unsafe"
14)
15
16// Handle is a globally unique identity for some value of type T.
17//
18// Two handles compare equal exactly if the two values used to create the handles
19// would have also compared equal. The comparison of two handles is trivial and
20// typically much more efficient than comparing the values used to create them.
21type Handle[T comparable] struct {
22	value *T
23}
24
25// Value returns a shallow copy of the T value that produced the Handle.
26func (h Handle[T]) Value() T {
27	return *h.value
28}
29
30// Make returns a globally unique handle for a value of type T. Handles
31// are equal if and only if the values used to produce them are equal.
32func Make[T comparable](value T) Handle[T] {
33	// Find the map for type T.
34	typ := abi.TypeFor[T]()
35	ma, ok := uniqueMaps.Load(typ)
36	if !ok {
37		// This is a good time to initialize cleanup, since we must go through
38		// this path on the first use of Make, and it's not on the hot path.
39		setupMake.Do(registerCleanup)
40		ma = addUniqueMap[T](typ)
41	}
42	m := ma.(*uniqueMap[T])
43
44	// Keep around any values we allocate for insertion. There
45	// are a few different ways we can race with other threads
46	// and create values that we might discard. By keeping
47	// the first one we make around, we can avoid generating
48	// more than one per racing thread.
49	var (
50		toInsert     *T // Keep this around to keep it alive.
51		toInsertWeak weak.Pointer[T]
52	)
53	newValue := func() (T, weak.Pointer[T]) {
54		if toInsert == nil {
55			toInsert = new(T)
56			*toInsert = clone(value, &m.cloneSeq)
57			toInsertWeak = weak.Make(toInsert)
58		}
59		return *toInsert, toInsertWeak
60	}
61	var ptr *T
62	for {
63		// Check the map.
64		wp, ok := m.Load(value)
65		if !ok {
66			// Try to insert a new value into the map.
67			k, v := newValue()
68			wp, _ = m.LoadOrStore(k, v)
69		}
70		// Now that we're sure there's a value in the map, let's
71		// try to get the pointer we need out of it.
72		ptr = wp.Strong()
73		if ptr != nil {
74			break
75		}
76		// The weak pointer is nil, so the old value is truly dead.
77		// Try to remove it and start over.
78		m.CompareAndDelete(value, wp)
79	}
80	runtime.KeepAlive(toInsert)
81	return Handle[T]{ptr}
82}
83
84var (
85	// uniqueMaps is an index of type-specific concurrent maps used for unique.Make.
86	//
87	// The two-level map might seem odd at first since the HashTrieMap could have "any"
88	// as its key type, but the issue is escape analysis. We do not want to force lookups
89	// to escape the argument, and using a type-specific map allows us to avoid that where
90	// possible (for example, for strings and plain-ol'-data structs). We also get the
91	// benefit of not cramming every different type into a single map, but that's certainly
92	// not enough to outweigh the cost of two map lookups. What is worth it though, is saving
93	// on those allocations.
94	uniqueMaps = concurrent.NewHashTrieMap[*abi.Type, any]() // any is always a *uniqueMap[T].
95
96	// cleanupFuncs are functions that clean up dead weak pointers in type-specific
97	// maps in uniqueMaps. We express cleanup this way because there's no way to iterate
98	// over the sync.Map and call functions on the type-specific data structures otherwise.
99	// These cleanup funcs each close over one of these type-specific maps.
100	//
101	// cleanupMu protects cleanupNotify and is held across the entire cleanup. Used for testing.
102	// cleanupNotify is a test-only mechanism that allow tests to wait for the cleanup to run.
103	cleanupMu      sync.Mutex
104	cleanupFuncsMu sync.Mutex
105	cleanupFuncs   []func()
106	cleanupNotify  []func() // One-time notifications when cleanups finish.
107)
108
109type uniqueMap[T comparable] struct {
110	*concurrent.HashTrieMap[T, weak.Pointer[T]]
111	cloneSeq
112}
113
114func addUniqueMap[T comparable](typ *abi.Type) *uniqueMap[T] {
115	// Create a map for T and try to register it. We could
116	// race with someone else, but that's fine; it's one
117	// small, stray allocation. The number of allocations
118	// this can create is bounded by a small constant.
119	m := &uniqueMap[T]{
120		HashTrieMap: concurrent.NewHashTrieMap[T, weak.Pointer[T]](),
121		cloneSeq:    makeCloneSeq(typ),
122	}
123	a, loaded := uniqueMaps.LoadOrStore(typ, m)
124	if !loaded {
125		// Add a cleanup function for the new map.
126		cleanupFuncsMu.Lock()
127		cleanupFuncs = append(cleanupFuncs, func() {
128			// Delete all the entries whose weak references are nil and clean up
129			// deleted entries.
130			m.All()(func(key T, wp weak.Pointer[T]) bool {
131				if wp.Strong() == nil {
132					m.CompareAndDelete(key, wp)
133				}
134				return true
135			})
136		})
137		cleanupFuncsMu.Unlock()
138	}
139	return a.(*uniqueMap[T])
140}
141
142// setupMake is used to perform initial setup for unique.Make.
143var setupMake sync.Once
144
145// startBackgroundCleanup sets up a background goroutine to occasionally call cleanupFuncs.
146func registerCleanup() {
147	runtime_registerUniqueMapCleanup(func() {
148		// Lock for cleanup.
149		cleanupMu.Lock()
150
151		// Grab funcs to run.
152		cleanupFuncsMu.Lock()
153		cf := cleanupFuncs
154		cleanupFuncsMu.Unlock()
155
156		// Run cleanup.
157		for _, f := range cf {
158			f()
159		}
160
161		// Run cleanup notifications.
162		for _, f := range cleanupNotify {
163			f()
164		}
165		cleanupNotify = nil
166
167		// Finished.
168		cleanupMu.Unlock()
169	})
170}
171
172// Implemented in runtime.
173
174//go:linkname runtime_registerUniqueMapCleanup
175func runtime_registerUniqueMapCleanup(cleanup func())
176