Disambiguate hash collisions in internal/hashmap via Hasher.Equal

internal/hashmap/hashmap.go

@@ -20,40 +20,65 @@ type entry[K, V any] struct {
 // Map is a mapping from keys of type K to values of type V,
 // using key-equivalence relation H.
 type Map[K, V any] struct {
-	entries map[uint64]entry[K, V]
-	h       Hasher[K]
+	// entries maps each key hash to the bucket of entries sharing that hash.
+	// Buckets disambiguate hash collisions via the Hasher's Equal method.
+	entries map[uint64][]entry[K, V]
+	// count is the total number of entries across all buckets, kept so Len is
+	// O(1) rather than summing bucket lengths.
+	count int
+	h     Hasher[K]
 }
 
 // NewMap returns a new mapping.
 func NewMap[K, V any](h Hasher[K]) *Map[K, V] {
 	return &Map[K, V]{
-		entries: make(map[uint64]entry[K, V]),
+		entries: make(map[uint64][]entry[K, V]),
 		h:       h,
 	}
 }
 
 // All returns an iterator over the key/value entries of the map in undefined order.
 func (m *Map[K, V]) All() iter.Seq2[K, V] {
 	return func(yield func(K, V) bool) {
-		for _, entry := range m.entries {
-			if !yield(entry.key, entry.value) {
-				return
+		for _, bucket := range m.entries {
+			for _, entry := range bucket {
+				if !yield(entry.key, entry.value) {
+					return
+				}
 			}
 		}
 	}
 }
 
 // Load returns the map entry for the given key.
 func (m *Map[K, V]) Load(key K) (V, bool) {
-	entry, ok := m.entries[m.h.Hash(key)]
-	return entry.value, ok
+	for _, entry := range m.entries[m.h.Hash(key)] {
+		if m.h.Equal(entry.key, key) {
+			return entry.value, true
+		}
+	}
+	var zero V
+	return zero, false
 }
 
-// Delete removes th//e entry with the given key, if any. It returns true if the entry was found.
+// Delete removes the entry with the given key, if any. It returns true if the entry was found.
 func (m *Map[K, V]) Delete(key K) bool {
 	hash := m.h.Hash(key)
-	if _, exists := m.entries[hash]; exists {
-		delete(m.entries, hash)
+	bucket := m.entries[hash]
+	for i, e := range bucket {
+		if !m.h.Equal(e.key, key) {
+			continue
+		}
+		if len(bucket) == 1 {
+			delete(m.entries, hash)
+		} else {
+			// Shift the tail down and clear the vacated slot so the removed
+			// entry's key/value are not retained by the backing array.
+			copy(bucket[i:], bucket[i+1:])
+			bucket[len(bucket)-1] = entry[K, V]{}
+			m.entries[hash] = bucket[:len(bucket)-1]
+		}
+		m.count--
 		return true
 	}
 	return false
@@ -62,9 +87,11 @@ func (m *Map[K, V]) Delete(key K) bool {
 // Keys returns an iterator over the map keys in unspecified order.
 func (m *Map[K, V]) Keys() iter.Seq[K] {
 	return func(yield func(K) bool) {
-		for _, entry := range m.entries {
-			if !yield(entry.key) {
-				return
+		for _, bucket := range m.entries {
+			for _, entry := range bucket {
+				if !yield(entry.key) {
+					return
+				}
 			}
 		}
 	}
@@ -73,44 +100,55 @@ func (m *Map[K, V]) Keys() iter.Seq[K] {
 // Values returns an iterator over the map values in unspecified order.
 func (m *Map[K, V]) Values() iter.Seq[V] {
 	return func(yield func(V) bool) {
-		for _, entry := range m.entries {
-			if !yield(entry.value) {
-				return
+		for _, bucket := range m.entries {
+			for _, entry := range bucket {
+				if !yield(entry.value) {
+					return
+				}
 			}
 		}
 	}
 }
 
 // Len returns the number of map entries.
 func (m *Map[K, V]) Len() int {
-	return len(m.entries)
+	return m.count
 }
 
 // Set updates the map entry for key to value, and returns the previous entry, if any.
 func (m *Map[K, V]) Set(key K, value V) (prev V) {
 	hash := m.h.Hash(key)
-	if entry, exists := m.entries[hash]; exists {
-		prev = entry.value
+	bucket := m.entries[hash]
+	for i, entry := range bucket {
+		if m.h.Equal(entry.key, key) {
+			prev = entry.value
+			bucket[i].value = value
+			return prev
+		}
 	}
-	m.entries[hash] = entry[K, V]{key, value}
+	m.entries[hash] = append(bucket, entry[K, V]{key, value})
+	m.count++
 	return prev
 }
 
 func (m *Map[K, V]) Clear() {
 	clear(m.entries)
+	m.count = 0
 }
 
 // String returns a string representation of the map's entries in unspecified order.
 // Values are printed as if by fmt.Sprint.
 func (m *Map[K, V]) String() string {
 	s := "{"
 	first := true
-	for _, entry := range m.entries {
-		if !first {
-			s += " "
+	for _, bucket := range m.entries {
+		for _, entry := range bucket {
+			if !first {
+				s += " "
+			}
+			s += fmt.Sprintf("%v: %v", entry.key, entry.value)
+			first = false
 		}
-		s += fmt.Sprintf("%v: %v", entry.key, entry.value)
-		first = false
 	}
 	s += "}"
 	return s

internal/hashmap/hashmap_test.go

@@ -0,0 +1,80 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+package hashmap_test
+
+import (
+	"testing"
+
+	"github.com/microsoft/agent-framework-go/internal/hashmap"
+)
+
+// collidingHasher forces every key to the same hash so distinct keys land in the
+// same bucket, exercising the hash-collision path. Equal still distinguishes keys.
+type collidingHasher struct{}
+
+func (collidingHasher) Hash(string) uint64     { return 1 }
+func (collidingHasher) Equal(a, b string) bool { return a == b }
+
+func TestMap_DistinctKeysWithCollidingHashCoexist(t *testing.T) {
+	m := hashmap.NewMap[string, string](collidingHasher{})
+	m.Set("alice", "a")
+	m.Set("bob", "b")
+
+	if got, ok := m.Load("alice"); !ok || got != "a" {
+		t.Errorf(`Load("alice") = (%q, %v), want ("a", true)`, got, ok)
+	}
+	if got, ok := m.Load("bob"); !ok || got != "b" {
+		t.Errorf(`Load("bob") = (%q, %v), want ("b", true)`, got, ok)
+	}
+	if got := m.Len(); got != 2 {
+		t.Errorf("Len() = %d, want 2", got)
+	}
+
+	// A missing key that collides with stored keys must not resolve to a stored value.
+	if got, ok := m.Load("carol"); ok {
+		t.Errorf(`Load("carol") = (%q, %v), want ("", false)`, got, ok)
+	}
+
+	// Deleting one colliding key must not remove the other.
+	if !m.Delete("alice") {
+		t.Error(`Delete("alice") = false, want true`)
+	}
+	if _, ok := m.Load("alice"); ok {
+		t.Error(`Load("alice") after delete still found, want not found`)
+	}
+	if got, ok := m.Load("bob"); !ok || got != "b" {
+		t.Errorf(`Load("bob") after deleting alice = (%q, %v), want ("b", true)`, got, ok)
+	}
+}
+
+func TestMap_LenTracksSetUpdateDeleteClear(t *testing.T) {
+	m := hashmap.NewMap[string, int](collidingHasher{})
+	if got := m.Len(); got != 0 {
+		t.Fatalf("Len() = %d, want 0", got)
+	}
+	m.Set("a", 1)
+	m.Set("b", 2)
+	if got := m.Len(); got != 2 {
+		t.Fatalf("after two inserts Len() = %d, want 2", got)
+	}
+	// Updating an existing key must not change the count.
+	m.Set("a", 10)
+	if got := m.Len(); got != 2 {
+		t.Fatalf("after updating an existing key Len() = %d, want 2", got)
+	}
+	if !m.Delete("a") {
+		t.Fatal(`Delete("a") = false, want true`)
+	}
+	if got := m.Len(); got != 1 {
+		t.Fatalf("after delete Len() = %d, want 1", got)
+	}
+	// Deleting a missing key must not change the count.
+	m.Delete("missing")
+	if got := m.Len(); got != 1 {
+		t.Fatalf("after deleting a missing key Len() = %d, want 1", got)
+	}
+	m.Clear()
+	if got := m.Len(); got != 0 {
+		t.Fatalf("after Clear Len() = %d, want 0", got)
+	}
+}