Add an `Ordered` interface for the Cmp method, have rbtree use it

4 files changed, 56 insertions, 38 deletions

diff --git a/lib/btrfs/btrfsvol/lvm.go b/lib/btrfs/btrfsvol/lvm.go
index 65636f6..ca46ad3 100644
--- a/lib/btrfs/btrfsvol/lvm.go
+++ b/lib/btrfs/btrfsvol/lvm.go
@@ -19,8 +19,8 @@ type LogicalVolume[PhysicalVolume util.File[PhysicalAddr]] struct {
 
 	id2pv map[DeviceID]PhysicalVolume
 
-	logical2physical *rbtree.Tree[LogicalAddr, chunkMapping]
-	physical2logical map[DeviceID]*rbtree.Tree[PhysicalAddr, devextMapping]
+	logical2physical *rbtree.Tree[util.NativeOrdered[LogicalAddr], chunkMapping]
+	physical2logical map[DeviceID]*rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping]
 }
 
 var _ util.File[LogicalAddr] = (*LogicalVolume[util.File[PhysicalAddr]])(nil)
@@ -30,20 +30,20 @@ func (lv *LogicalVolume[PhysicalVolume]) init() {
 		lv.id2pv = make(map[DeviceID]PhysicalVolume)
 	}
 	if lv.logical2physical == nil {
-		lv.logical2physical = &rbtree.Tree[LogicalAddr, chunkMapping]{
-			KeyFn: func(chunk chunkMapping) LogicalAddr {
-				return chunk.LAddr
+		lv.logical2physical = &rbtree.Tree[util.NativeOrdered[LogicalAddr], chunkMapping]{
+			KeyFn: func(chunk chunkMapping) util.NativeOrdered[LogicalAddr] {
+				return util.NativeOrdered[LogicalAddr]{Val: chunk.LAddr}
 			},
 		}
 	}
 	if lv.physical2logical == nil {
-		lv.physical2logical = make(map[DeviceID]*rbtree.Tree[PhysicalAddr, devextMapping], len(lv.id2pv))
+		lv.physical2logical = make(map[DeviceID]*rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping], len(lv.id2pv))
 	}
 	for devid := range lv.id2pv {
 		if _, ok := lv.physical2logical[devid]; !ok {
-			lv.physical2logical[devid] = &rbtree.Tree[PhysicalAddr, devextMapping]{
-				KeyFn: func(ext devextMapping) PhysicalAddr {
-					return ext.PAddr
+			lv.physical2logical[devid] = &rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping]{
+				KeyFn: func(ext devextMapping) util.NativeOrdered[PhysicalAddr] {
+					return util.NativeOrdered[PhysicalAddr]{Val: ext.PAddr}
 				},
 			}
 		}
@@ -74,9 +74,9 @@ func (lv *LogicalVolume[PhysicalVolume]) AddPhysicalVolume(id DeviceID, dev Phys
 			lv, dev.Name(), other.Name(), id)
 	}
 	lv.id2pv[id] = dev
-	lv.physical2logical[id] = &rbtree.Tree[PhysicalAddr, devextMapping]{
-		KeyFn: func(ext devextMapping) PhysicalAddr {
-			return ext.PAddr
+	lv.physical2logical[id] = &rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping]{
+		KeyFn: func(ext devextMapping) util.NativeOrdered[PhysicalAddr] {
+			return util.NativeOrdered[PhysicalAddr]{Val: ext.PAddr}
 		},
 	}
 	return nil
@@ -148,13 +148,13 @@ func (lv *LogicalVolume[PhysicalVolume]) AddMapping(m Mapping) error {
 
 	// logical2physical
 	for _, chunk := range logicalOverlaps {
-		lv.logical2physical.Delete(chunk.LAddr)
+		lv.logical2physical.Delete(util.NativeOrdered[LogicalAddr]{Val: chunk.LAddr})
 	}
 	lv.logical2physical.Insert(newChunk)
 
 	// physical2logical
 	for _, ext := range physicalOverlaps {
-		lv.physical2logical[m.PAddr.Dev].Delete(ext.PAddr)
+		lv.physical2logical[m.PAddr.Dev].Delete(util.NativeOrdered[PhysicalAddr]{Val: ext.PAddr})
 	}
 	lv.physical2logical[m.PAddr.Dev].Insert(newExt)
 
@@ -173,7 +173,7 @@ func (lv *LogicalVolume[PhysicalVolume]) AddMapping(m Mapping) error {
 }
 
 func (lv *LogicalVolume[PhysicalVolume]) fsck() error {
-	physical2logical := make(map[DeviceID]*rbtree.Tree[PhysicalAddr, devextMapping])
+	physical2logical := make(map[DeviceID]*rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping])
 	if err := lv.logical2physical.Walk(func(node *rbtree.Node[chunkMapping]) error {
 		chunk := node.Value
 		for _, stripe := range chunk.PAddrs {
@@ -182,9 +182,9 @@ func (lv *LogicalVolume[PhysicalVolume]) fsck() error {
 					lv, stripe.Dev)
 			}
 			if _, exists := physical2logical[stripe.Dev]; !exists {
-				physical2logical[stripe.Dev] = &rbtree.Tree[PhysicalAddr, devextMapping]{
-					KeyFn: func(ext devextMapping) PhysicalAddr {
-						return ext.PAddr
+				physical2logical[stripe.Dev] = &rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping]{
+					KeyFn: func(ext devextMapping) util.NativeOrdered[PhysicalAddr] {
+						return util.NativeOrdered[PhysicalAddr]{Val: ext.PAddr}
 					},
 				}
 			}
diff --git a/lib/rbtree/rbtree.go b/lib/rbtree/rbtree.go
index 0353e75..575a9ca 100644
--- a/lib/rbtree/rbtree.go
+++ b/lib/rbtree/rbtree.go
@@ -7,7 +7,7 @@ package rbtree
 import (
 	"fmt"
 
-	"golang.org/x/exp/constraints"
+	"git.lukeshu.com/btrfs-progs-ng/lib/util"
 )
 
 type Color bool
@@ -32,7 +32,7 @@ func (node *Node[V]) getColor() Color {
 	return node.Color
 }
 
-type Tree[K constraints.Ordered, V any] struct {
+type Tree[K util.Ordered[K], V any] struct {
 	KeyFn func(V) K
 	root  *Node[V]
 }
@@ -104,14 +104,7 @@ func (node *Node[V]) search(fn func(V) int) (exact, nearest *Node[V]) {
 func (t *Tree[K, V]) exactKey(key K) func(V) int {
 	return func(val V) int {
 		valKey := t.KeyFn(val)
-		switch {
-		case key < valKey:
-			return -1
-		case key > valKey:
-			return 1
-		default: // key == valKey:
-			return 0
-		}
+		return key.Cmp(valKey)
 	}
 }
 
@@ -282,7 +275,7 @@ func (t *Tree[K, V]) Insert(val V) {
 	}
 	if parent == nil {
 		t.root = node
-	} else if key < t.KeyFn(parent.Value) {
+	} else if key.Cmp(t.KeyFn(parent.Value)) < 0 {
 		parent.Left = node
 	} else {
 		parent.Right = node
diff --git a/lib/rbtree/rbtree_test.go b/lib/rbtree/rbtree_test.go
index 9e9a734..999eb6f 100644
--- a/lib/rbtree/rbtree_test.go
+++ b/lib/rbtree/rbtree_test.go
@@ -10,9 +10,11 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/constraints"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/util"
 )
 
-func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]struct{}, tree *Tree[K, V]) {
+func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]struct{}, tree *Tree[util.NativeOrdered[K], V]) {
 	// 1. Every node is either red or black
 
 	// 2. The root is black.
@@ -62,16 +64,16 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 	expectedOrder := make([]K, 0, len(expectedSet))
 	for k := range expectedSet {
 		expectedOrder = append(expectedOrder, k)
-		node := tree.Lookup(k)
+		node := tree.Lookup(util.NativeOrdered[K]{Val: k})
 		require.NotNil(t, tree, node)
-		require.Equal(t, k, tree.KeyFn(node.Value))
+		require.Equal(t, k, tree.KeyFn(node.Value).Val)
 	}
 	sort.Slice(expectedOrder, func(i, j int) bool {
 		return expectedOrder[i] < expectedOrder[j]
 	})
 	actOrder := make([]K, 0, len(expectedSet))
 	require.NoError(t, tree.Walk(func(node *Node[V]) error {
-		actOrder = append(actOrder, tree.KeyFn(node.Value))
+		actOrder = append(actOrder, tree.KeyFn(node.Value).Val)
 		return nil
 	}))
 	require.Equal(t, expectedOrder, actOrder)
@@ -100,8 +102,10 @@ func FuzzTree(f *testing.F) {
 	})
 
 	f.Fuzz(func(t *testing.T, dat []uint8) {
-		tree := &Tree[uint8, uint8]{
-			KeyFn: func(x uint8) uint8 { return x },
+		tree := &Tree[util.NativeOrdered[uint8], uint8]{
+			KeyFn: func(x uint8) util.NativeOrdered[uint8] {
+				return util.NativeOrdered[uint8]{Val: x}
+			},
 		}
 		set := make(map[uint8]struct{})
 		checkTree(t, set, tree)
@@ -114,15 +118,15 @@ func FuzzTree(f *testing.F) {
 				tree.Insert(val)
 				set[val] = struct{}{}
 				t.Logf("\n%s\n", tree.ASCIIArt())
-				node := tree.Lookup(val)
+				node := tree.Lookup(util.NativeOrdered[uint8]{Val: val})
 				require.NotNil(t, node)
 				require.Equal(t, val, node.Value)
 			} else {
 				t.Logf("Delete(%v)", val)
-				tree.Delete(val)
+				tree.Delete(util.NativeOrdered[uint8]{Val: val})
 				delete(set, val)
 				t.Logf("\n%s\n", tree.ASCIIArt())
-				require.Nil(t, tree.Lookup(val))
+				require.Nil(t, tree.Lookup(util.NativeOrdered[uint8]{Val: val}))
 			}
 			checkTree(t, set, tree)
 		}
diff --git a/lib/util/generic.go b/lib/util/generic.go
index 61e045b..a63a814 100644
--- a/lib/util/generic.go
+++ b/lib/util/generic.go
@@ -124,3 +124,24 @@ func (m *SyncMap[K, V]) Range(f func(key K, value V) bool) {
 	})
 }
 func (m *SyncMap[K, V]) Store(key K, value V) { m.inner.Store(key, value) }
+
+type Ordered[T interface{ Cmp(T) int }] interface {
+	Cmp(T) int
+}
+
+type NativeOrdered[T constraints.Ordered] struct {
+	Val T
+}
+
+func (a NativeOrdered[T]) Cmp(b NativeOrdered[T]) int {
+	switch {
+	case a.Val < b.Val:
+		return -1
+	case a.Val > b.Val:
+		return 1
+	default:
+		return 0
+	}
+}
+
+var _ Ordered[NativeOrdered[int]] = NativeOrdered[int]{}