Adding more functions to InstanceMap

compiler/internal/typeparams/map.go

@@ -1,8 +1,10 @@
 package typeparams
 
 import (
+	"fmt"
 	"go/types"
-	"sync"
+	"sort"
+	"strings"
 
 	"golang.org/x/tools/go/types/typeutil"
 )
@@ -25,40 +27,35 @@ type (
 // instance equality, objects are compared by pointer equality, and type
 // arguments with types.Identical(). To reduce access complexity, we bucket
 // entries by a combined hash of type args. This type is generally inspired by
-// typeutil.Map.
+// golang.org/x/tools/go/types/typeutil/map.go
 type InstanceMap[V any] struct {
-	bootstrap sync.Once
-	data      map[types.Object]mapBuckets[V]
-	len       int
-	hasher    typeutil.Hasher
-	zero      V
+	data   map[types.Object]mapBuckets[V]
+	len    int
+	hasher typeutil.Hasher
 }
 
-func (im *InstanceMap[V]) init() {
-	im.bootstrap.Do(func() {
-		im.data = map[types.Object]mapBuckets[V]{}
-		im.hasher = typeutil.MakeHasher()
-	})
+// findIndex returns bucket and index of the entry with the given key.
+// If the given key isn't found, an empty bucket and -1 are returned.
+func (im *InstanceMap[V]) findIndex(key Instance) (mapBucket[V], int) {
+	if im != nil && im.data != nil {
+		bucket := im.data[key.Object][typeHash(im.hasher, key.TArgs...)]
+		for i, candidate := range bucket {
+			if candidate != nil && typeArgsEq(candidate.key.TArgs, key.TArgs) {
+				return bucket, i
+			}
+		}
+	}
+	return nil, -1
 }
 
+// get returns the stored value for the provided key and
+// a bool indicating whether the key was present in the map or not.
 func (im *InstanceMap[V]) get(key Instance) (V, bool) {
-	im.init()
-
-	buckets, ok := im.data[key.Object]
-	if !ok {
-		return im.zero, false
-	}
-	bucket := buckets[typeHash(im.hasher, key.TArgs...)]
-	if len(bucket) == 0 {
-		return im.zero, false
-	}
-
-	for _, candidate := range bucket {
-		if typeArgsEq(candidate.key.TArgs, key.TArgs) {
-			return candidate.value, true
-		}
+	if bucket, i := im.findIndex(key); i >= 0 {
+		return bucket[i].value, true
 	}
-	return im.zero, false
+	var zero V
+	return zero, false
 }
 
 // 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 {
 
 // Set new value for the key in the map. Returns the previous value that was
 // stored in the map, or zero value if the key wasn't present before.
-func (im *InstanceMap[V]) Set(key Instance, value V) (old V) {
-	im.init()
+func (im *InstanceMap[V]) Set(key Instance, value V) V {
+	if im.data == nil {
+		im.data = map[types.Object]mapBuckets[V]{}
+		im.hasher = typeutil.MakeHasher()
+	}
 
 	if _, ok := im.data[key.Object]; !ok {
 		im.data[key.Object] = mapBuckets[V]{}
 	}
 	bucketID := typeHash(im.hasher, key.TArgs...)
 
 	// If there is already an identical key in the map, override the entry value.
-	for _, candidate := range im.data[key.Object][bucketID] {
-		if typeArgsEq(candidate.key.TArgs, key.TArgs) {
-			old = candidate.value
+	hole := -1
+	bucket := im.data[key.Object][bucketID]
+	for i, candidate := range bucket {
+		if candidate == nil {
+			hole = i
+		} else if typeArgsEq(candidate.key.TArgs, key.TArgs) {
+			old := candidate.value
 			candidate.value = value
 			return old
 		}
 	}
 
-	// Otherwise append a new entry.
-	im.data[key.Object][bucketID] = append(im.data[key.Object][bucketID], &mapEntry[V]{
-		key:   key,
-		value: value,
-	})
+	// If there is a hole in the bucket, reuse it.
+	if hole >= 0 {
+		im.data[key.Object][bucketID][hole] = &mapEntry[V]{
+			key:   key,
+			value: value,
+		}
+	} else {
+		// Otherwise append a new entry.
+		im.data[key.Object][bucketID] = append(bucket, &mapEntry[V]{
+			key:   key,
+			value: value,
+		})
+	}
 	im.len++
-	return im.zero
+	var zero V
+	return zero
 }
 
 // Len returns the number of elements in the map.
 func (im *InstanceMap[V]) Len() int {
-	return im.len
+	if im != nil {
+		return im.len
+	}
+	return 0
+}
+
+// Delete removes the entry with the given key, if any.
+// It returns true if the entry was found.
+func (im *InstanceMap[V]) Delete(key Instance) bool {
+	if bucket, i := im.findIndex(key); i >= 0 {
+		// We can't compact the bucket as it
+		// would disturb iterators.
+		bucket[i] = nil
+		im.len--
+		return true
+	}
+	return false
+}
+
+// Iterate calls function f on each entry in the map in unspecified order.
+//
+// Return true from f to continue the iteration, or false to stop it.
+//
+// If f should mutate the map, Iterate provides the same guarantees as
+// Go maps: if f deletes a map entry that Iterate has not yet reached,
+// f will not be invoked for it, but if f inserts a map entry that
+// Iterate has not yet reached, whether or not f will be invoked for
+// it is unspecified.
+func (im *InstanceMap[V]) Iterate(f func(key Instance, value V)) {
+	if im != nil && im.data != nil {
+		for _, mapBucket := range im.data {
+			for _, bucket := range mapBucket {
+				for _, e := range bucket {
+					if e != nil {
+						f(e.key, e.value)
+					}
+				}
+			}
+		}
+	}
+}
+
+// Keys returns a new slice containing the set of map keys.
+// The order is unspecified.
+func (im *InstanceMap[V]) Keys() []Instance {
+	keys := make([]Instance, 0, im.Len())
+	im.Iterate(func(key Instance, _ V) {
+		keys = append(keys, key)
+	})
+	return keys
+}
+
+// String returns a string representation of the map's entries.
+// The entries are sorted by string representation of the entry.
+func (im *InstanceMap[V]) String() string {
+	entries := make([]string, 0, im.Len())
+	im.Iterate(func(key Instance, value V) {
+		entries = append(entries, fmt.Sprintf("%v:%v", key, value))
+	})
+	sort.Strings(entries)
+	return `{` + strings.Join(entries, `, `) + `}`
 }
 
 // typeHash returns a combined hash of several types.