Skip to content

Commit a07a61f

Browse files
Adding more functions to InstanceMap
1 parent 56cec36 commit a07a61f

File tree

2 files changed

+334
-47
lines changed

2 files changed

+334
-47
lines changed

compiler/internal/typeparams/map.go

Lines changed: 113 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package typeparams
22

33
import (
4+
"fmt"
45
"go/types"
5-
"sync"
6+
"sort"
7+
"strings"
68

79
"golang.org/x/tools/go/types/typeutil"
810
)
@@ -25,40 +27,35 @@ type (
2527
// instance equality, objects are compared by pointer equality, and type
2628
// arguments with types.Identical(). To reduce access complexity, we bucket
2729
// entries by a combined hash of type args. This type is generally inspired by
28-
// typeutil.Map.
30+
// golang.org/x/tools/go/types/typeutil/map.go
2931
type InstanceMap[V any] struct {
30-
bootstrap sync.Once
31-
data map[types.Object]mapBuckets[V]
32-
len int
33-
hasher typeutil.Hasher
34-
zero V
32+
data map[types.Object]mapBuckets[V]
33+
len int
34+
hasher typeutil.Hasher
3535
}
3636

37-
func (im *InstanceMap[V]) init() {
38-
im.bootstrap.Do(func() {
39-
im.data = map[types.Object]mapBuckets[V]{}
40-
im.hasher = typeutil.MakeHasher()
41-
})
37+
// findIndex returns bucket and index of the entry with the given key.
38+
// If the given key isn't found, an empty bucket and -1 are returned.
39+
func (im *InstanceMap[V]) findIndex(key Instance) (mapBucket[V], int) {
40+
if im != nil && im.data != nil {
41+
bucket := im.data[key.Object][typeHash(im.hasher, key.TArgs...)]
42+
for i, candidate := range bucket {
43+
if candidate != nil && typeArgsEq(candidate.key.TArgs, key.TArgs) {
44+
return bucket, i
45+
}
46+
}
47+
}
48+
return nil, -1
4249
}
4350

51+
// get returns the stored value for the provided key and
52+
// a bool indicating whether the key was present in the map or not.
4453
func (im *InstanceMap[V]) get(key Instance) (V, bool) {
45-
im.init()
46-
47-
buckets, ok := im.data[key.Object]
48-
if !ok {
49-
return im.zero, false
50-
}
51-
bucket := buckets[typeHash(im.hasher, key.TArgs...)]
52-
if len(bucket) == 0 {
53-
return im.zero, false
54-
}
55-
56-
for _, candidate := range bucket {
57-
if typeArgsEq(candidate.key.TArgs, key.TArgs) {
58-
return candidate.value, true
59-
}
54+
if bucket, i := im.findIndex(key); i >= 0 {
55+
return bucket[i].value, true
6056
}
61-
return im.zero, false
57+
var zero V
58+
return zero, false
6259
}
6360

6461
// Get returns the stored value for the provided key. If the key is missing from
@@ -76,35 +73,111 @@ func (im *InstanceMap[V]) Has(key Instance) bool {
7673

7774
// Set new value for the key in the map. Returns the previous value that was
7875
// stored in the map, or zero value if the key wasn't present before.
79-
func (im *InstanceMap[V]) Set(key Instance, value V) (old V) {
80-
im.init()
76+
func (im *InstanceMap[V]) Set(key Instance, value V) V {
77+
if im.data == nil {
78+
im.data = map[types.Object]mapBuckets[V]{}
79+
im.hasher = typeutil.MakeHasher()
80+
}
8181

8282
if _, ok := im.data[key.Object]; !ok {
8383
im.data[key.Object] = mapBuckets[V]{}
8484
}
8585
bucketID := typeHash(im.hasher, key.TArgs...)
8686

8787
// If there is already an identical key in the map, override the entry value.
88-
for _, candidate := range im.data[key.Object][bucketID] {
89-
if typeArgsEq(candidate.key.TArgs, key.TArgs) {
90-
old = candidate.value
88+
hole := -1
89+
bucket := im.data[key.Object][bucketID]
90+
for i, candidate := range bucket {
91+
if candidate == nil {
92+
hole = i
93+
} else if typeArgsEq(candidate.key.TArgs, key.TArgs) {
94+
old := candidate.value
9195
candidate.value = value
9296
return old
9397
}
9498
}
9599

96-
// Otherwise append a new entry.
97-
im.data[key.Object][bucketID] = append(im.data[key.Object][bucketID], &mapEntry[V]{
98-
key: key,
99-
value: value,
100-
})
100+
// If there is a hole in the bucket, reuse it.
101+
if hole >= 0 {
102+
im.data[key.Object][bucketID][hole] = &mapEntry[V]{
103+
key: key,
104+
value: value,
105+
}
106+
} else {
107+
// Otherwise append a new entry.
108+
im.data[key.Object][bucketID] = append(bucket, &mapEntry[V]{
109+
key: key,
110+
value: value,
111+
})
112+
}
101113
im.len++
102-
return im.zero
114+
var zero V
115+
return zero
103116
}
104117

105118
// Len returns the number of elements in the map.
106119
func (im *InstanceMap[V]) Len() int {
107-
return im.len
120+
if im != nil {
121+
return im.len
122+
}
123+
return 0
124+
}
125+
126+
// Delete removes the entry with the given key, if any.
127+
// It returns true if the entry was found.
128+
func (im *InstanceMap[V]) Delete(key Instance) bool {
129+
if bucket, i := im.findIndex(key); i >= 0 {
130+
// We can't compact the bucket as it
131+
// would disturb iterators.
132+
bucket[i] = nil
133+
im.len--
134+
return true
135+
}
136+
return false
137+
}
138+
139+
// Iterate calls function f on each entry in the map in unspecified order.
140+
//
141+
// Return true from f to continue the iteration, or false to stop it.
142+
//
143+
// If f should mutate the map, Iterate provides the same guarantees as
144+
// Go maps: if f deletes a map entry that Iterate has not yet reached,
145+
// f will not be invoked for it, but if f inserts a map entry that
146+
// Iterate has not yet reached, whether or not f will be invoked for
147+
// it is unspecified.
148+
func (im *InstanceMap[V]) Iterate(f func(key Instance, value V)) {
149+
if im != nil && im.data != nil {
150+
for _, mapBucket := range im.data {
151+
for _, bucket := range mapBucket {
152+
for _, e := range bucket {
153+
if e != nil {
154+
f(e.key, e.value)
155+
}
156+
}
157+
}
158+
}
159+
}
160+
}
161+
162+
// Keys returns a new slice containing the set of map keys.
163+
// The order is unspecified.
164+
func (im *InstanceMap[V]) Keys() []Instance {
165+
keys := make([]Instance, 0, im.Len())
166+
im.Iterate(func(key Instance, _ V) {
167+
keys = append(keys, key)
168+
})
169+
return keys
170+
}
171+
172+
// String returns a string representation of the map's entries.
173+
// The entries are sorted by string representation of the entry.
174+
func (im *InstanceMap[V]) String() string {
175+
entries := make([]string, 0, im.Len())
176+
im.Iterate(func(key Instance, value V) {
177+
entries = append(entries, fmt.Sprintf("%v:%v", key, value))
178+
})
179+
sort.Strings(entries)
180+
return `{` + strings.Join(entries, `, `) + `}`
108181
}
109182

110183
// typeHash returns a combined hash of several types.

0 commit comments

Comments
 (0)