containers: Rethink the RBTree interface to be simpler

4 files changed, 72 insertions, 57 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]
 
diff --git a/lib/btrfsprogs/btrfsutil/broken_btree.go b/lib/btrfsprogs/btrfsutil/broken_btree.go
index 8261119..7ea31ce 100644
--- a/lib/btrfsprogs/btrfsutil/broken_btree.go
+++ b/lib/btrfsprogs/btrfsutil/broken_btree.go
@@ -23,7 +23,7 @@ import (
 
 type treeIndex struct {
 	TreeRootErr error
-	Items       *containers.RBTree[btrfsprim.Key, treeIndexValue]
+	Items       *containers.RBTree[treeIndexValue]
 	Errors      *containers.IntervalTree[btrfsprim.Key, treeIndexError]
 }
 
@@ -38,13 +38,14 @@ type treeIndexValue struct {
 	ItemSize uint32
 }
 
+// Compare implements containers.Ordered.
+func (a treeIndexValue) Compare(b treeIndexValue) int {
+	return a.Key.Compare(b.Key)
+}
+
 func newTreeIndex(arena *SkinnyPathArena) treeIndex {
 	return treeIndex{
-		Items: &containers.RBTree[btrfsprim.Key, treeIndexValue]{
-			KeyFn: func(iv treeIndexValue) btrfsprim.Key {
-				return iv.Key
-			},
-		},
+		Items: new(containers.RBTree[treeIndexValue]),
 		Errors: &containers.IntervalTree[btrfsprim.Key, treeIndexError]{
 			MinFn: func(err treeIndexError) btrfsprim.Key {
 				return arena.Inflate(err.Path).Node(-1).ToKey
@@ -173,7 +174,7 @@ func (bt *brokenTrees) rawTreeWalk(root btrfstree.TreeRoot, cacheEntry treeIndex
 		},
 		btrfstree.TreeWalkHandler{
 			Item: func(path btrfstree.TreePath, item btrfstree.Item) error {
-				if cacheEntry.Items.Lookup(item.Key) != nil {
+				if cacheEntry.Items.Search(func(v treeIndexValue) int { return item.Key.Compare(v.Key) }) != nil {
 					// This is a panic because I'm not really sure what the best way to
 					// handle this is, and so if this happens I want the program to crash
 					// and force me to figure out how to handle it.
@@ -203,15 +204,15 @@ func (bt *brokenTrees) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key) (bt
 
 func (bt *brokenTrees) addErrs(index treeIndex, fn func(btrfsprim.Key, uint32) int, err error) error {
 	var errs derror.MultiError
-	if _errs := index.Errors.SearchAll(func(k btrfsprim.Key) int { return fn(k, 0) }); len(_errs) > 0 {
-		errs = make(derror.MultiError, len(_errs))
-		for i := range _errs {
-			errs[i] = &btrfstree.TreeError{
-				Path: bt.arena.Inflate(_errs[i].Path),
-				Err:  _errs[i].Err,
-			}
-		}
-	}
+	index.Errors.Subrange(
+		func(k btrfsprim.Key) int { return fn(k, 0) },
+		func(v treeIndexError) bool {
+			errs = append(errs, &btrfstree.TreeError{
+				Path: bt.arena.Inflate(v.Path),
+				Err:  v.Err,
+			})
+			return true
+		})
 	if len(errs) == 0 {
 		return err
 	}
@@ -253,9 +254,13 @@ func (bt *brokenTrees) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfsprim.K
 		return nil, index.TreeRootErr
 	}
 
-	indexItems := index.Items.SearchRange(func(indexItem treeIndexValue) int {
-		return fn(indexItem.Key, indexItem.ItemSize)
-	})
+	var indexItems []treeIndexValue
+	index.Items.Subrange(
+		func(indexItem treeIndexValue) int { return fn(indexItem.Key, indexItem.ItemSize) },
+		func(node *containers.RBNode[treeIndexValue]) bool {
+			indexItems = append(indexItems, node.Value)
+			return true
+		})
 	if len(indexItems) == 0 {
 		return nil, bt.addErrs(index, fn, iofs.ErrNotExist)
 	}
@@ -290,12 +295,12 @@ func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, err
 		return
 	}
 	var node *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]
-	_ = index.Items.Walk(func(indexItem *containers.RBNode[treeIndexValue]) error {
+	index.Items.Range(func(indexItem *containers.RBNode[treeIndexValue]) bool {
 		if ctx.Err() != nil {
-			return ctx.Err()
+			return false
 		}
 		if bt.ctx.Err() != nil {
-			return bt.ctx.Err()
+			return false
 		}
 		if cbs.Item != nil {
 			itemPath := bt.arena.Inflate(indexItem.Value.Path)
@@ -304,7 +309,7 @@ func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, err
 				node, err = bt.inner.ReadNode(itemPath.Parent())
 				if err != nil {
 					errHandle(&btrfstree.TreeError{Path: itemPath, Err: err})
-					return nil //nolint:nilerr // We already called errHandle().
+					return true
 				}
 			}
 			item := node.Data.BodyLeaf[itemPath.Node(-1).FromItemIdx]
@@ -312,7 +317,7 @@ func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, err
 				errHandle(&btrfstree.TreeError{Path: itemPath, Err: err})
 			}
 		}
-		return nil
+		return true
 	})
 }