Merge pull request #29 from laher/map-keys-comparable

generics: widen map key constraint to 'comparable'

1 files changed, 81 insertions, 47 deletions

diff --git a/immutable.go b/immutable.go
index 68e9580..2b91c9c 100644
--- a/immutable.go
+++ b/immutable.go
@@ -1,6 +1,6 @@
 // Package immutable provides immutable collection types.
 //
-// Introduction
+// # Introduction
 //
 // Immutable collections provide an efficient, safe way to share collections
 // of data while minimizing locks. The collections in this package provide
@@ -14,7 +14,7 @@
 // with Go's built-in collection types so please evaluate for your specific
 // use.
 //
-// Collection Types
+// # Collection Types
 //
 // The List type provides an API similar to Go slices. They allow appending,
 // prepending, and updating of elements. Elements can also be fetched by index
@@ -28,7 +28,7 @@
 // provides iteration over unsorted keys. Maps improved performance and memory
 // usage as compared to SortedMaps.
 //
-// Hashing and Sorting
+// # Hashing and Sorting
 //
 // Map types require the use of a Hasher implementation to calculate hashes for
 // their keys and check for key equality. SortedMaps require the use of a
@@ -507,7 +507,7 @@ func (n *listLeafNode[T]) deleteBefore(index int, mutable bool) listNode[T] {
 		copy(other.children[idx:][:], n.children[idx:][:])
 	}
 	// Set the first idx bits to 0.
-	other.occupied &= ^((1 << idx)-1)
+	other.occupied &= ^((1 << idx) - 1)
 	return other
 }
 
@@ -530,7 +530,7 @@ func (n *listLeafNode[T]) deleteAfter(index int, mutable bool) listNode[T] {
 		copy(other.children[:idx+1][:], n.children[:idx+1][:])
 	}
 	// Set bits after idx to 0. idx < 31 because n.containsAfter(index) == true.
-	other.occupied &= (1 << (idx+1))-1
+	other.occupied &= (1 << (idx + 1)) - 1
 	return other
 }
 
@@ -680,7 +680,7 @@ const (
 // to generate hashes and check for equality of key values.
 //
 // It is implemented as an Hash Array Mapped Trie.
-type Map[K constraints.Ordered, V any] struct {
+type Map[K comparable, V any] struct {
 	size   int           // total number of key/value pairs
 	root   mapNode[K, V] // root node of trie
 	hasher Hasher[K]     // hasher implementation
@@ -689,7 +689,7 @@ type Map[K constraints.Ordered, V any] struct {
 // NewMap returns a new instance of Map. If hasher is nil, a default hasher
 // implementation will automatically be chosen based on the first key added.
 // Default hasher implementations only exist for int, string, and byte slice types.
-func NewMap[K constraints.Ordered, V any](hasher Hasher[K]) *Map[K, V] {
+func NewMap[K comparable, V any](hasher Hasher[K]) *Map[K, V] {
 	return &Map[K, V]{
 		hasher: hasher,
 	}
@@ -796,12 +796,12 @@ func (m *Map[K, V]) Iterator() *MapIterator[K, V] {
 }
 
 // MapBuilder represents an efficient builder for creating Maps.
-type MapBuilder[K constraints.Ordered, V any] struct {
+type MapBuilder[K comparable, V any] struct {
 	m *Map[K, V] // current state
 }
 
 // NewMapBuilder returns a new instance of MapBuilder.
-func NewMapBuilder[K constraints.Ordered, V any](hasher Hasher[K]) *MapBuilder[K, V] {
+func NewMapBuilder[K comparable, V any](hasher Hasher[K]) *MapBuilder[K, V] {
 	return &MapBuilder[K, V]{m: NewMap[K, V](hasher)}
 }
 
@@ -845,7 +845,7 @@ func (b *MapBuilder[K, V]) Iterator() *MapIterator[K, V] {
 }
 
 // mapNode represents any node in the map tree.
-type mapNode[K constraints.Ordered, V any] interface {
+type mapNode[K comparable, V any] interface {
 	get(key K, shift uint, keyHash uint32, h Hasher[K]) (value V, ok bool)
 	set(key K, value V, shift uint, keyHash uint32, h Hasher[K], mutable bool, resized *bool) mapNode[K, V]
 	delete(key K, shift uint, keyHash uint32, h Hasher[K], mutable bool, resized *bool) mapNode[K, V]
@@ -858,7 +858,7 @@ var _ mapNode[string, any] = (*mapValueNode[string, any])(nil)
 var _ mapNode[string, any] = (*mapHashCollisionNode[string, any])(nil)
 
 // mapLeafNode represents a node that stores a single key hash at the leaf of the map tree.
-type mapLeafNode[K constraints.Ordered, V any] interface {
+type mapLeafNode[K comparable, V any] interface {
 	mapNode[K, V]
 	keyHashValue() uint32
 }
@@ -869,7 +869,7 @@ var _ mapLeafNode[string, any] = (*mapHashCollisionNode[string, any])(nil)
 // mapArrayNode is a map node that stores key/value pairs in a slice.
 // Entries are stored in insertion order. An array node expands into a bitmap
 // indexed node once a given threshold size is crossed.
-type mapArrayNode[K constraints.Ordered, V any] struct {
+type mapArrayNode[K comparable, V any] struct {
 	entries []mapEntry[K, V]
 }
 
@@ -971,7 +971,7 @@ func (n *mapArrayNode[K, V]) delete(key K, shift uint, keyHash uint32, h Hasher[
 // mapBitmapIndexedNode represents a map branch node with a variable number of
 // node slots and indexed using a bitmap. Indexes for the node slots are
 // calculated by counting the number of set bits before the target bit using popcount.
-type mapBitmapIndexedNode[K constraints.Ordered, V any] struct {
+type mapBitmapIndexedNode[K comparable, V any] struct {
 	bitmap uint32
 	nodes  []mapNode[K, V]
 }
@@ -1117,7 +1117,7 @@ func (n *mapBitmapIndexedNode[K, V]) delete(key K, shift uint, keyHash uint32, h
 
 // mapHashArrayNode is a map branch node that stores nodes in a fixed length
 // array. Child nodes are indexed by their index bit segment for the current depth.
-type mapHashArrayNode[K constraints.Ordered, V any] struct {
+type mapHashArrayNode[K comparable, V any] struct {
 	count uint                       // number of set nodes
 	nodes [mapNodeSize]mapNode[K, V] // child node slots, may contain empties
 }
@@ -1213,14 +1213,14 @@ func (n *mapHashArrayNode[K, V]) delete(key K, shift uint, keyHash uint32, h Has
 // mapValueNode represents a leaf node with a single key/value pair.
 // A value node can be converted to a hash collision leaf node if a different
 // key with the same keyHash is inserted.
-type mapValueNode[K constraints.Ordered, V any] struct {
+type mapValueNode[K comparable, V any] struct {
 	keyHash uint32
 	key     K
 	value   V
 }
 
 // newMapValueNode returns a new instance of mapValueNode.
-func newMapValueNode[K constraints.Ordered, V any](keyHash uint32, key K, value V) *mapValueNode[K, V] {
+func newMapValueNode[K comparable, V any](keyHash uint32, key K, value V) *mapValueNode[K, V] {
 	return &mapValueNode[K, V]{
 		keyHash: keyHash,
 		key:     key,
@@ -1285,7 +1285,7 @@ func (n *mapValueNode[K, V]) delete(key K, shift uint, keyHash uint32, h Hasher[
 
 // mapHashCollisionNode represents a leaf node that contains two or more key/value
 // pairs with the same key hash. Single pairs for a hash are stored as value nodes.
-type mapHashCollisionNode[K constraints.Ordered, V any] struct {
+type mapHashCollisionNode[K comparable, V any] struct {
 	keyHash uint32 // key hash for all entries
 	entries []mapEntry[K, V]
 }
@@ -1391,7 +1391,7 @@ func (n *mapHashCollisionNode[K, V]) delete(key K, shift uint, keyHash uint32, h
 
 // mergeIntoNode merges a key/value pair into an existing node.
 // Caller must verify that node's keyHash is not equal to keyHash.
-func mergeIntoNode[K constraints.Ordered, V any](node mapLeafNode[K, V], shift uint, keyHash uint32, key K, value V) mapNode[K, V] {
+func mergeIntoNode[K comparable, V any](node mapLeafNode[K, V], shift uint, keyHash uint32, key K, value V) mapNode[K, V] {
 	idx1 := (node.keyHashValue() >> shift) & mapNodeMask
 	idx2 := (keyHash >> shift) & mapNodeMask
 
@@ -1410,14 +1410,14 @@ func mergeIntoNode[K constraints.Ordered, V any](node mapLeafNode[K, V], shift u
 }
 
 // mapEntry represents a single key/value pair.
-type mapEntry[K constraints.Ordered, V any] struct {
+type mapEntry[K comparable, V any] struct {
 	key   K
 	value V
 }
 
 // MapIterator represents an iterator over a map's key/value pairs. Although
 // map keys are not sorted, the iterator's order is deterministic.
-type MapIterator[K constraints.Ordered, V any] struct {
+type MapIterator[K comparable, V any] struct {
 	m *Map[K, V] // source map
 
 	stack [32]mapIteratorElem[K, V] // search stack
@@ -1541,7 +1541,7 @@ func (itr *MapIterator[K, V]) first() {
 }
 
 // mapIteratorElem represents a node/index pair in the MapIterator stack.
-type mapIteratorElem[K constraints.Ordered, V any] struct {
+type mapIteratorElem[K comparable, V any] struct {
 	node  mapNode[K, V]
 	index int
 }
@@ -1555,7 +1555,7 @@ const (
 // is determined by the Comparer used by the map.
 //
 // This map is implemented as a B+tree.
-type SortedMap[K constraints.Ordered, V any] struct {
+type SortedMap[K comparable, V any] struct {
 	size     int                 // total number of key/value pairs
 	root     sortedMapNode[K, V] // root of b+tree
 	comparer Comparer[K]
@@ -1564,7 +1564,7 @@ type SortedMap[K constraints.Ordered, V any] struct {
 // NewSortedMap returns a new instance of SortedMap. If comparer is nil then
 // a default comparer is set after the first key is inserted. Default comparers
 // exist for int, string, and byte slice keys.
-func NewSortedMap[K constraints.Ordered, V any](comparer Comparer[K]) *SortedMap[K, V] {
+func NewSortedMap[K comparable, V any](comparer Comparer[K]) *SortedMap[K, V] {
 	return &SortedMap[K, V]{
 		comparer: comparer,
 	}
@@ -1594,7 +1594,7 @@ func (m *SortedMap[K, V]) set(key K, value V, mutable bool) *SortedMap[K, V] {
 	// Set a comparer on the first value if one does not already exist.
 	comparer := m.comparer
 	if comparer == nil {
-		comparer = NewComparer[K](key)
+		comparer = NewComparer(key)
 	}
 
 	// Create copy, if necessary.
@@ -1673,12 +1673,12 @@ func (m *SortedMap[K, V]) Iterator() *SortedMapIterator[K, V] {
 }
 
 // SortedMapBuilder represents an efficient builder for creating sorted maps.
-type SortedMapBuilder[K constraints.Ordered, V any] struct {
+type SortedMapBuilder[K comparable, V any] struct {
 	m *SortedMap[K, V] // current state
 }
 
 // NewSortedMapBuilder returns a new instance of SortedMapBuilder.
-func NewSortedMapBuilder[K constraints.Ordered, V any](comparer Comparer[K]) *SortedMapBuilder[K, V] {
+func NewSortedMapBuilder[K comparable, V any](comparer Comparer[K]) *SortedMapBuilder[K, V] {
 	return &SortedMapBuilder[K, V]{m: NewSortedMap[K, V](comparer)}
 }
 
@@ -1722,7 +1722,7 @@ func (b *SortedMapBuilder[K, V]) Iterator() *SortedMapIterator[K, V] {
 }
 
 // sortedMapNode represents a branch or leaf node in the sorted map.
-type sortedMapNode[K constraints.Ordered, V any] interface {
+type sortedMapNode[K comparable, V any] interface {
 	minKey() K
 	indexOf(key K, c Comparer[K]) int
 	get(key K, c Comparer[K]) (value V, ok bool)
@@ -1734,12 +1734,12 @@ var _ sortedMapNode[string, any] = (*sortedMapBranchNode[string, any])(nil)
 var _ sortedMapNode[string, any] = (*sortedMapLeafNode[string, any])(nil)
 
 // sortedMapBranchNode represents a branch in the sorted map.
-type sortedMapBranchNode[K constraints.Ordered, V any] struct {
+type sortedMapBranchNode[K comparable, V any] struct {
 	elems []sortedMapBranchElem[K, V]
 }
 
 // newSortedMapBranchNode returns a new branch node with the given child nodes.
-func newSortedMapBranchNode[K constraints.Ordered, V any](children ...sortedMapNode[K, V]) *sortedMapBranchNode[K, V] {
+func newSortedMapBranchNode[K comparable, V any](children ...sortedMapNode[K, V]) *sortedMapBranchNode[K, V] {
 	// Fetch min keys for every child.
 	elems := make([]sortedMapBranchElem[K, V], len(children))
 	for i, child := range children {
@@ -1882,13 +1882,13 @@ func (n *sortedMapBranchNode[K, V]) delete(key K, c Comparer[K], mutable bool, r
 	return other
 }
 
-type sortedMapBranchElem[K constraints.Ordered, V any] struct {
+type sortedMapBranchElem[K comparable, V any] struct {
 	key  K
 	node sortedMapNode[K, V]
 }
 
 // sortedMapLeafNode represents a leaf node in the sorted map.
-type sortedMapLeafNode[K constraints.Ordered, V any] struct {
+type sortedMapLeafNode[K comparable, V any] struct {
 	entries []mapEntry[K, V]
 }
 
@@ -2003,7 +2003,7 @@ func (n *sortedMapLeafNode[K, V]) delete(key K, c Comparer[K], mutable bool, res
 
 // SortedMapIterator represents an iterator over a sorted map.
 // Iteration can occur in natural or reverse order based on use of Next() or Prev().
-type SortedMapIterator[K constraints.Ordered, V any] struct {
+type SortedMapIterator[K comparable, V any] struct {
 	m *SortedMap[K, V] // source map
 
 	stack [32]sortedMapIteratorElem[K, V] // search stack
@@ -2191,13 +2191,13 @@ func (itr *SortedMapIterator[K, V]) seek(key K) {
 }
 
 // sortedMapIteratorElem represents node/index pair in the SortedMapIterator stack.
-type sortedMapIteratorElem[K constraints.Ordered, V any] struct {
+type sortedMapIteratorElem[K comparable, V any] struct {
 	node  sortedMapNode[K, V]
 	index int
 }
 
 // Hasher hashes keys and checks them for equality.
-type Hasher[K constraints.Ordered] interface {
+type Hasher[K comparable] interface {
 	// Computes a hash for key.
 	Hash(key K) uint32
 
@@ -2206,7 +2206,7 @@ type Hasher[K constraints.Ordered] interface {
 }
 
 // NewHasher returns the built-in hasher for a given key type.
-func NewHasher[K constraints.Ordered](key K) Hasher[K] {
+func NewHasher[K comparable](key K) Hasher[K] {
 	// Attempt to use non-reflection based hasher first.
 	switch (any(key)).(type) {
 	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, uintptr, string:
@@ -2234,8 +2234,8 @@ func hashString(value string) uint32 {
 	return hash
 }
 
-// reflectIntHasher implements a reflection-based Hasher for int keys.
-type reflectHasher[K constraints.Ordered] struct{}
+// reflectIntHasher implements a reflection-based Hasher for keys.
+type reflectHasher[K comparable] struct{}
 
 // Hash returns a hash for key.
 func (h *reflectHasher[K]) Hash(key K) uint32 {
@@ -2256,7 +2256,7 @@ func (h *reflectHasher[K]) Hash(key K) uint32 {
 }
 
 // Equal returns true if a is equal to b. Otherwise returns false.
-// Panics if a and b are not ints.
+// Panics if a and b are not int-ish or string-ish.
 func (h *reflectHasher[K]) Equal(a, b K) bool {
 	switch reflect.TypeOf(a).Kind() {
 	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
@@ -2281,7 +2281,7 @@ func hashUint64(value uint64) uint32 {
 }
 
 // defaultHasher implements Hasher.
-type defaultHasher[K constraints.Ordered] struct{}
+type defaultHasher[K comparable] struct{}
 
 // Hash returns a hash for key.
 func (h *defaultHasher[K]) Hash(key K) uint32 {
@@ -2322,14 +2322,16 @@ func (h *defaultHasher[K]) Equal(a, b K) bool {
 }
 
 // Comparer allows the comparison of two keys for the purpose of sorting.
-type Comparer[K constraints.Ordered] interface {
+type Comparer[K comparable] interface {
 	// Returns -1 if a is less than b, returns 1 if a is greater than b,
 	// and returns 0 if a is equal to b.
 	Compare(a, b K) int
 }
 
 // NewComparer returns the built-in comparer for a given key type.
-func NewComparer[K constraints.Ordered](key K) Comparer[K] {
+// Note that only int-ish and string-ish types are supported, despite the 'comparable' constraint.
+// Attempts to use other types will result in a panic - users should define their own Comparers for these cases.
+func NewComparer[K comparable](key K) Comparer[K] {
 	// Attempt to use non-reflection based comparer first.
 	switch (any(key)).(type) {
 	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, uintptr, string:
@@ -2346,12 +2348,44 @@ func NewComparer[K constraints.Ordered](key K) Comparer[K] {
 	panic(fmt.Sprintf("immutable.NewComparer: must set comparer for %T type", key))
 }
 
-// defaultComparer compares two integers. Implements Comparer.
-type defaultComparer[K constraints.Ordered] struct{}
+// defaultComparer compares two values (int-ish and string-ish types are supported. Implements Comparer.
+type defaultComparer[K comparable] struct{}
 
 // Compare returns -1 if a is less than b, returns 1 if a is greater than b, and
-// returns 0 if a is equal to b. Panic if a or b is not an int.
+// returns 0 if a is equal to b. Panic if a or b is not a string or int* type
 func (c *defaultComparer[K]) Compare(i K, j K) int {
+	switch x := (any(i)).(type) {
+	case int:
+		return defaultCompare(x, (any(j)).(int))
+	case int8:
+		return defaultCompare(x, (any(j)).(int8))
+	case int16:
+		return defaultCompare(x, (any(j)).(int16))
+	case int32:
+		return defaultCompare(x, (any(j)).(int32))
+	case int64:
+		return defaultCompare(x, (any(j)).(int64))
+	case uint:
+		return defaultCompare(x, (any(j)).(uint))
+	case uint8:
+		return defaultCompare(x, (any(j)).(uint8))
+	case uint16:
+		return defaultCompare(x, (any(j)).(uint16))
+	case uint32:
+		return defaultCompare(x, (any(j)).(uint32))
+	case uint64:
+		return defaultCompare(x, (any(j)).(uint64))
+	case uintptr:
+		return defaultCompare(x, (any(j)).(uintptr))
+	case string:
+		return defaultCompare(x, (any(j)).(string))
+	}
+	panic(fmt.Sprintf("immutable.defaultComparer: must set comparer for %T type", i))
+}
+
+// defaultCompare only operates on constraints.Ordered.
+// For other types, users should bring their own comparers
+func defaultCompare[K constraints.Ordered](i, j K) int {
 	if i < j {
 		return -1
 	} else if i > j {
@@ -2360,11 +2394,11 @@ func (c *defaultComparer[K]) Compare(i K, j K) int {
 	return 0
 }
 
-// reflectIntComparer compares two int values using reflection. Implements Comparer.
-type reflectComparer[K constraints.Ordered] struct{}
+// reflectIntComparer compares two values using reflection. Implements Comparer.
+type reflectComparer[K comparable] struct{}
 
 // Compare returns -1 if a is less than b, returns 1 if a is greater than b, and
-// returns 0 if a is equal to b. Panic if a or b is not an int.
+// returns 0 if a is equal to b. Panic if a or b is not an int-ish or string-ish type.
 func (c *reflectComparer[K]) Compare(a, b K) int {
 	switch reflect.TypeOf(a).Kind() {
 	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: