containers: Rethink the RBTree interface to be simpler

3 files changed, 43 insertions, 33 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go b/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go
index 69d14c7..7c02d05 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go
@@ -53,11 +53,7 @@ func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevice
 	// "AAAAAAA" shouldn't be present, and if we just discard "BBBBBBBB"
 	// because it conflicts with "CCCCCCC", then we would erroneously
 	// include "AAAAAAA".
-	addrspace := &containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]{
-		KeyFn: func(item btrfsinspect.SysExtentCSum) containers.NativeOrdered[btrfsvol.LogicalAddr] {
-			return containers.NativeOrdered[btrfsvol.LogicalAddr]{Val: item.Sums.Addr}
-		},
-	}
+	addrspace := new(containers.RBTree[btrfsinspect.SysExtentCSum])
 	for _, newRecord := range records {
 		for {
 			conflict := addrspace.Search(func(oldRecord btrfsinspect.SysExtentCSum) int {
@@ -85,7 +81,7 @@ func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevice
 			}
 			if oldRecord.Generation < newRecord.Generation {
 				// Newer generation wins.
-				addrspace.Delete(containers.NativeOrdered[btrfsvol.LogicalAddr]{Val: oldRecord.Sums.Addr})
+				addrspace.Delete(conflict)
 				// loop around to check for more conflicts
 				continue
 			}
@@ -142,7 +138,7 @@ func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevice
 					},
 				},
 			}
-			addrspace.Delete(containers.NativeOrdered[btrfsvol.LogicalAddr]{Val: oldRecord.Sums.Addr})
+			addrspace.Delete(conflict)
 			newRecord = unionRecord
 			// loop around to check for more conflicts
 		}
@@ -152,7 +148,7 @@ func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevice
 	var flattened SumRunWithGaps[btrfsvol.LogicalAddr]
 	var curAddr btrfsvol.LogicalAddr
 	var curSums strings.Builder
-	_ = addrspace.Walk(func(node *containers.RBNode[btrfsinspect.SysExtentCSum]) error {
+	addrspace.Range(func(node *containers.RBNode[btrfsinspect.SysExtentCSum]) bool {
 		curEnd := curAddr + (btrfsvol.LogicalAddr(curSums.Len()/sumSize) * btrfssum.BlockSize)
 		if node.Value.Sums.Addr != curEnd {
 			if curSums.Len() > 0 {
@@ -166,7 +162,7 @@ func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevice
 			curSums.Reset()
 		}
 		curSums.WriteString(string(node.Value.Sums.Sums))
-		return nil
+		return true
 	})
 	if curSums.Len() > 0 {
 		flattened.Runs = append(flattened.Runs, btrfssum.SumRun[btrfsvol.LogicalAddr]{
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index ee5950d..bd29278 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -770,6 +770,16 @@ func (o *rebuilder) _walkRange(
 		})
 }
 
+type gap struct {
+	// range is [Beg,End)
+	Beg, End uint64
+}
+
+// Compare implements containers.Ordered.
+func (a gap) Compare(b gap) int {
+	return containers.NativeCompare(a.Beg, b.Beg)
+}
+
 func (o *rebuilder) _wantRange(
 	ctx context.Context,
 	treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType,
@@ -787,15 +797,7 @@ func (o *rebuilder) _wantRange(
 	//
 	// Start with a gap of the whole range, then subtract each run
 	// from it.
-	type gap struct {
-		// range is [Beg,End)
-		Beg, End uint64
-	}
-	gaps := &containers.RBTree[containers.NativeOrdered[uint64], gap]{
-		KeyFn: func(gap gap) containers.NativeOrdered[uint64] {
-			return containers.NativeOrdered[uint64]{Val: gap.Beg}
-		},
-	}
+	gaps := new(containers.RBTree[gap])
 	gaps.Insert(gap{
 		Beg: beg,
 		End: end,
@@ -805,23 +807,29 @@ func (o *rebuilder) _wantRange(
 		o.rebuilt.Tree(ctx, treeID).Items(ctx),
 		treeID, objID, typ, beg, end,
 		func(runKey btrfsprim.Key, _ keyio.ItemPtr, runBeg, runEnd uint64) {
-			overlappingGaps := gaps.SearchRange(func(gap gap) int {
-				switch {
-				case gap.End <= runBeg:
-					return 1
-				case runEnd <= gap.Beg:
-					return -1
-				default:
-					return 0
-				}
-			})
+			var overlappingGaps []*containers.RBNode[gap]
+			gaps.Subrange(
+				func(gap gap) int {
+					switch {
+					case gap.End <= runBeg:
+						return 1
+					case runEnd <= gap.Beg:
+						return -1
+					default:
+						return 0
+					}
+				},
+				func(node *containers.RBNode[gap]) bool {
+					overlappingGaps = append(overlappingGaps, node)
+					return true
+				})
 			if len(overlappingGaps) == 0 {
 				return
 			}
-			gapsBeg := overlappingGaps[0].Beg
-			gapsEnd := overlappingGaps[len(overlappingGaps)-1].End
+			gapsBeg := overlappingGaps[0].Value.Beg
+			gapsEnd := overlappingGaps[len(overlappingGaps)-1].Value.End
 			for _, gap := range overlappingGaps {
-				gaps.Delete(containers.NativeOrdered[uint64]{Val: gap.Beg})
+				gaps.Delete(gap)
 			}
 			if gapsBeg < runBeg {
 				gaps.Insert(gap{
@@ -842,7 +850,7 @@ func (o *rebuilder) _wantRange(
 		return
 	}
 	potentialItems := o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx)
-	_ = gaps.Walk(func(rbNode *containers.RBNode[gap]) error {
+	gaps.Range(func(rbNode *containers.RBNode[gap]) bool {
 		gap := rbNode.Value
 		last := gap.Beg
 		o._walkRange(
@@ -874,7 +882,7 @@ func (o *rebuilder) _wantRange(
 				o.wantAugment(wantCtx, treeID, wantKey, nil)
 			}
 		}
-		return nil
+		return true
 	})
 }
 
diff --git a/lib/btrfsprogs/btrfsinspect/scandevices.go b/lib/btrfsprogs/btrfsinspect/scandevices.go
index 7668a83..9b8360c 100644
--- a/lib/btrfsprogs/btrfsinspect/scandevices.go
+++ b/lib/btrfsprogs/btrfsinspect/scandevices.go
@@ -22,6 +22,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfssum"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
@@ -79,6 +80,11 @@ type SysExtentCSum struct {
 	Sums       btrfsitem.ExtentCSum
 }
 
+// Compare implements containers.Ordered.
+func (a SysExtentCSum) Compare(b SysExtentCSum) int {
+	return containers.NativeCompare(a.Sums.Addr, b.Sums.Addr)
+}
+
 type scanStats struct {
 	textui.Portion[btrfsvol.PhysicalAddr]