6 files changed, 151 insertions, 9 deletions

diff --git a/lib/containers/intervaltree_test.go b/lib/containers/intervaltree_test.go
index 7a4689b..88c3003 100644
--- a/lib/containers/intervaltree_test.go
+++ b/lib/containers/intervaltree_test.go
@@ -56,7 +56,7 @@ func TestIntervalTree(t *testing.T) {
 	tree.Insert(SimpleInterval{6, 10})
 	tree.Insert(SimpleInterval{19, 20})
 
-	t.Log(tree.ASCIIArt())
+	t.Log("\n" + tree.ASCIIArt())
 
 	// find intervals that touch [9,20]
 	intervals := tree.SearchAll(func(k NativeOrdered[int]) int {
diff --git a/lib/containers/ordered.go b/lib/containers/ordered.go
index 038edf8..d918149 100644
--- a/lib/containers/ordered.go
+++ b/lib/containers/ordered.go
@@ -8,10 +8,12 @@ import (
 	"golang.org/x/exp/constraints"
 )
 
-type Ordered[T interface{ Cmp(T) int }] interface {
+type _Ordered[T any] interface {
 	Cmp(T) int
 }
 
+type Ordered[T _Ordered[T]] _Ordered[T]
+
 type NativeOrdered[T constraints.Ordered] struct {
 	Val T
 }
diff --git a/lib/containers/rbtree.go b/lib/containers/rbtree.go
index ab79a26..e2d8f8a 100644
--- a/lib/containers/rbtree.go
+++ b/lib/containers/rbtree.go
@@ -37,6 +37,11 @@ type RBTree[K Ordered[K], V any] struct {
 	KeyFn  func(V) K
 	AttrFn func(*RBNode[V])
 	root   *RBNode[V]
+	len    int
+}
+
+func (t *RBTree[K, V]) Len() int {
+	return t.len
 }
 
 func (t *RBTree[K, V]) Walk(fn func(*RBNode[V]) error) error {
@@ -324,6 +329,7 @@ func (t *RBTree[K, V]) Insert(val V) {
 		exact.Value = val
 		return
 	}
+	t.len++
 
 	node := &RBNode[V]{
 		Color:  Red,
@@ -394,6 +400,7 @@ func (t *RBTree[K, V]) Delete(key K) {
 	if nodeToDelete == nil {
 		return
 	}
+	t.len--
 
 	// This is closely based on the algorithm presented in CLRS
 	// 3e.
diff --git a/lib/containers/rbtree_test.go b/lib/containers/rbtree_test.go
index 9841d26..a487c52 100644
--- a/lib/containers/rbtree_test.go
+++ b/lib/containers/rbtree_test.go
@@ -112,6 +112,7 @@ func checkRBTree[K constraints.Ordered, V any](t *testing.T, expectedSet Set[K],
 		return nil
 	}))
 	require.Equal(t, expectedOrder, actOrder)
+	require.Equal(t, tree.Len(), len(expectedSet))
 }
 
 func FuzzRBTree(f *testing.F) {
diff --git a/lib/containers/set.go b/lib/containers/set.go
index 1c525ca..67ba7ac 100644
--- a/lib/containers/set.go
+++ b/lib/containers/set.go
@@ -5,27 +5,70 @@
 package containers
 
 import (
+	"fmt"
 	"io"
+	"sort"
 
 	"git.lukeshu.com/go/lowmemjson"
-	"golang.org/x/exp/constraints"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
 // Set[T] is an unordered set of T.
-//
-// Despite Set[T] being unordered, T is required to be an ordered type
-// in order that a Set[T] have a deterministic JSON representation.
-type Set[T constraints.Ordered] map[T]struct{}
+type Set[T comparable] map[T]struct{}
 
 var (
 	_ lowmemjson.Encodable = Set[int]{}
 	_ lowmemjson.Decodable = (*Set[int])(nil)
 )
 
+func cast[T any](x any) T { return x.(T) }
+
 func (o Set[T]) EncodeJSON(w io.Writer) error {
-	return lowmemjson.Encode(w, maps.SortedKeys(o))
+	var less func(a, b T) bool
+	var zero T
+	switch (any(zero)).(type) {
+	case _Ordered[T]:
+		less = func(a, b T) bool { return cast[_Ordered[T]](a).Cmp(b) < 0 }
+	// This is the constraints.Ordered list
+	case string:
+		less = func(a, b T) bool { return cast[string](a) < cast[string](b) }
+	case int:
+		less = func(a, b T) bool { return cast[int](a) < cast[int](b) }
+	case int8:
+		less = func(a, b T) bool { return cast[int8](a) < cast[int8](b) }
+	case int16:
+		less = func(a, b T) bool { return cast[int16](a) < cast[int16](b) }
+	case int32:
+		less = func(a, b T) bool { return cast[int32](a) < cast[int32](b) }
+	case int64:
+		less = func(a, b T) bool { return cast[int64](a) < cast[int64](b) }
+	case uint:
+		less = func(a, b T) bool { return cast[uint](a) < cast[uint](b) }
+	case uint8:
+		less = func(a, b T) bool { return cast[uint8](a) < cast[uint8](b) }
+	case uint16:
+		less = func(a, b T) bool { return cast[uint16](a) < cast[uint16](b) }
+	case uint32:
+		less = func(a, b T) bool { return cast[uint32](a) < cast[uint32](b) }
+	case uint64:
+		less = func(a, b T) bool { return cast[uint64](a) < cast[uint64](b) }
+	case uintptr:
+		less = func(a, b T) bool { return cast[uintptr](a) < cast[uintptr](b) }
+	case float32:
+		less = func(a, b T) bool { return cast[float32](a) < cast[float32](b) }
+	case float64:
+		less = func(a, b T) bool { return cast[float64](a) < cast[float64](b) }
+	default:
+		less = func(a, b T) bool { return fmt.Sprint(a) < fmt.Sprint(b) }
+	}
+
+	keys := maps.Keys(o)
+	sort.Slice(keys, func(i, j int) bool {
+		return less(keys[i], keys[j])
+	})
+
+	return lowmemjson.Encode(w, keys)
 }
 
 func (o *Set[T]) DecodeJSON(r io.RuneScanner) error {
@@ -49,7 +92,7 @@ func (o *Set[T]) DecodeJSON(r io.RuneScanner) error {
 	})
 }
 
-func NewSet[T constraints.Ordered](values ...T) Set[T] {
+func NewSet[T comparable](values ...T) Set[T] {
 	ret := make(Set[T], len(values))
 	for _, value := range values {
 		ret.Insert(value)
@@ -107,3 +150,16 @@ func (small Set[T]) HasAny(big Set[T]) bool {
 	}
 	return false
 }
+
+func (small Set[T]) Intersection(big Set[T]) Set[T] {
+	if len(big) < len(small) {
+		small, big = big, small
+	}
+	ret := make(Set[T])
+	for v := range small {
+		if _, ok := big[v]; ok {
+			ret.Insert(v)
+		}
+	}
+	return ret
+}
diff --git a/lib/containers/sortedmap.go b/lib/containers/sortedmap.go
new file mode 100644
index 0000000..b759fa3
--- /dev/null
+++ b/lib/containers/sortedmap.go
@@ -0,0 +1,76 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package containers
+
+import (
+	"errors"
+)
+
+type orderedKV[K Ordered[K], V any] struct {
+	K K
+	V V
+}
+
+type SortedMap[K Ordered[K], V any] struct {
+	inner RBTree[K, orderedKV[K, V]]
+}
+
+func (m *SortedMap[K, V]) init() {
+	if m.inner.KeyFn == nil {
+		m.inner.KeyFn = m.keyFn
+	}
+}
+
+func (m *SortedMap[K, V]) keyFn(kv orderedKV[K, V]) K {
+	return kv.K
+}
+
+func (m *SortedMap[K, V]) Delete(key K) {
+	m.init()
+	m.inner.Delete(key)
+}
+
+func (m *SortedMap[K, V]) Load(key K) (value V, ok bool) {
+	m.init()
+	node := m.inner.Lookup(key)
+	if node == nil {
+		var zero V
+		return zero, false
+	}
+	return node.Value.V, true
+}
+
+var errStop = errors.New("stop")
+
+func (m *SortedMap[K, V]) Range(f func(key K, value V) bool) {
+	m.init()
+	_ = m.inner.Walk(func(node *RBNode[orderedKV[K, V]]) error {
+		if f(node.Value.K, node.Value.V) {
+			return nil
+		} else {
+			return errStop
+		}
+	})
+}
+
+func (m *SortedMap[K, V]) Subrange(rangeFn func(K, V) int, handleFn func(K, V) bool) {
+	m.init()
+	kvs := m.inner.SearchRange(func(kv orderedKV[K, V]) int {
+		return rangeFn(kv.K, kv.V)
+	})
+	for _, kv := range kvs {
+		if !handleFn(kv.K, kv.V) {
+			break
+		}
+	}
+}
+
+func (m *SortedMap[K, V]) Store(key K, value V) {
+	m.init()
+	m.inner.Insert(orderedKV[K, V]{
+		K: key,
+		V: value,
+	})
+}