1 files changed, 71 insertions, 61 deletions

diff --git a/stmutil/containers.go b/stmutil/containers.go
index c7a4a49..0cc592d 100644
--- a/stmutil/containers.go
+++ b/stmutil/containers.go
@@ -3,142 +3,154 @@ package stmutil
 import (
 	"unsafe"
 
-	"github.com/benbjohnson/immutable"
-
 	"github.com/anacrolix/missinggo/v2/iter"
+	"github.com/benbjohnson/immutable"
 )
 
-type Settish interface {
-	Add(any) Settish
-	Delete(any) Settish
-	Contains(any) bool
-	Range(func(any) bool)
+// This is the type constraint for keys passed through from github.com/benbjohnson/immutable.
+type KeyConstraint interface {
+	comparable
+}
+
+type Settish[K KeyConstraint] interface {
+	Add(K) Settish[K]
+	Delete(K) Settish[K]
+	Contains(K) bool
+	Range(func(K) bool)
 	iter.Iterable
 	Len() int
 }
 
-type mapToSet struct {
-	m Mappish
+type mapToSet[K KeyConstraint] struct {
+	m Mappish[K, struct{}]
 }
 
-type interhash struct{}
+type interhash[K KeyConstraint] struct{}
 
-func (interhash) Hash(x any) uint32 {
+func (interhash[K]) Hash(x K) uint32 {
 	return uint32(nilinterhash(unsafe.Pointer(&x), 0))
 }
 
-func (interhash) Equal(i, j any) bool {
+func (interhash[K]) Equal(i, j K) bool {
 	return i == j
 }
 
-func NewSet() Settish {
-	return mapToSet{NewMap()}
+func NewSet[K KeyConstraint]() Settish[K] {
+	return mapToSet[K]{NewMap[K, struct{}]()}
 }
 
-func NewSortedSet(lesser lessFunc) Settish {
-	return mapToSet{NewSortedMap(lesser)}
+func NewSortedSet[K KeyConstraint](lesser lessFunc[K]) Settish[K] {
+	return mapToSet[K]{NewSortedMap[K, struct{}](lesser)}
 }
 
-func (s mapToSet) Add(x any) Settish {
-	s.m = s.m.Set(x, nil)
+func (s mapToSet[K]) Add(x K) Settish[K] {
+	s.m = s.m.Set(x, struct{}{})
 	return s
 }
 
-func (s mapToSet) Delete(x any) Settish {
+func (s mapToSet[K]) Delete(x K) Settish[K] {
 	s.m = s.m.Delete(x)
 	return s
 }
 
-func (s mapToSet) Len() int {
+func (s mapToSet[K]) Len() int {
 	return s.m.Len()
 }
 
-func (s mapToSet) Contains(x any) bool {
+func (s mapToSet[K]) Contains(x K) bool {
 	_, ok := s.m.Get(x)
 	return ok
 }
 
-func (s mapToSet) Range(f func(any) bool) {
-	s.m.Range(func(k, _ any) bool {
+func (s mapToSet[K]) Range(f func(K) bool) {
+	s.m.Range(func(k K, _ struct{}) bool {
 		return f(k)
 	})
 }
 
-func (s mapToSet) Iter(cb iter.Callback) {
-	s.Range(cb)
+func (s mapToSet[K]) Iter(cb iter.Callback) {
+	s.Range(func(k K) bool {
+		return cb(k)
+	})
 }
 
-type Map struct {
-	*immutable.Map
+type Map[K KeyConstraint, V any] struct {
+	*immutable.Map[K, V]
 }
 
-func NewMap() Mappish {
-	return Map{immutable.NewMap(interhash{})}
+func NewMap[K KeyConstraint, V any]() Mappish[K, V] {
+	return Map[K, V]{immutable.NewMap[K, V](interhash[K]{})}
 }
 
-var _ Mappish = Map{}
-
-func (m Map) Delete(x any) Mappish {
+func (m Map[K, V]) Delete(x K) Mappish[K, V] {
 	m.Map = m.Map.Delete(x)
 	return m
 }
 
-func (m Map) Set(key, value any) Mappish {
+func (m Map[K, V]) Set(key K, value V) Mappish[K, V] {
 	m.Map = m.Map.Set(key, value)
 	return m
 }
 
-func (sm Map) Range(f func(key, value any) bool) {
+func (sm Map[K, V]) Range(f func(K, V) bool) {
 	iter := sm.Map.Iterator()
-	for !iter.Done() {
-		if !f(iter.Next()) {
+	for {
+		k, v, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if !f(k, v) {
 			return
 		}
 	}
 }
 
-func (sm Map) Iter(cb iter.Callback) {
-	sm.Range(func(key, _ any) bool {
+func (sm Map[K, V]) Iter(cb iter.Callback) {
+	sm.Range(func(key K, _ V) bool {
 		return cb(key)
 	})
 }
 
-type SortedMap struct {
-	*immutable.SortedMap
+type SortedMap[K KeyConstraint, V any] struct {
+	*immutable.SortedMap[K, V]
 }
 
-func (sm SortedMap) Set(key, value any) Mappish {
+func (sm SortedMap[K, V]) Set(key K, value V) Mappish[K, V] {
 	sm.SortedMap = sm.SortedMap.Set(key, value)
 	return sm
 }
 
-func (sm SortedMap) Delete(key any) Mappish {
+func (sm SortedMap[K, V]) Delete(key K) Mappish[K, V] {
 	sm.SortedMap = sm.SortedMap.Delete(key)
 	return sm
 }
 
-func (sm SortedMap) Range(f func(key, value any) bool) {
+func (sm SortedMap[K, V]) Range(f func(key K, value V) bool) {
 	iter := sm.SortedMap.Iterator()
-	for !iter.Done() {
-		if !f(iter.Next()) {
+	for {
+		k, v, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if !f(k, v) {
 			return
 		}
 	}
 }
 
-func (sm SortedMap) Iter(cb iter.Callback) {
-	sm.Range(func(key, _ any) bool {
+func (sm SortedMap[K, V]) Iter(cb iter.Callback) {
+	sm.Range(func(key K, _ V) bool {
 		return cb(key)
 	})
 }
 
-type lessFunc func(l, r any) bool
+type lessFunc[T KeyConstraint] func(l, r T) bool
 
-type comparer struct {
-	less lessFunc
+type comparer[K KeyConstraint] struct {
+	less lessFunc[K]
 }
 
-func (me comparer) Compare(i, j any) int {
+func (me comparer[K]) Compare(i, j K) int {
 	if me.less(i, j) {
 		return -1
 	} else if me.less(j, i) {
@@ -148,17 +160,17 @@ func (me comparer) Compare(i, j any) int {
 	}
 }
 
-func NewSortedMap(less lessFunc) Mappish {
-	return SortedMap{
-		SortedMap: immutable.NewSortedMap(comparer{less}),
+func NewSortedMap[K KeyConstraint, V any](less lessFunc[K]) Mappish[K, V] {
+	return SortedMap[K, V]{
+		SortedMap: immutable.NewSortedMap[K, V](comparer[K]{less}),
 	}
 }
 
-type Mappish interface {
-	Set(key, value any) Mappish
-	Delete(key any) Mappish
-	Get(key any) (any, bool)
-	Range(func(_, _ any) bool)
+type Mappish[K, V any] interface {
+	Set(K, V) Mappish[K, V]
+	Delete(key K) Mappish[K, V]
+	Get(key K) (V, bool)
+	Range(func(K, V) bool)
 	Len() int
 	iter.Iterable
 }
@@ -178,5 +190,3 @@ func interfaceHash(x any) uint32 {
 type Lenner interface {
 	Len() int
 }
-
-type List = *immutable.List