btrfstree: Rethink the API, but leave the old API in place

6 files changed, 256 insertions, 29 deletions

diff --git a/lib/btrfs/btrfstree/btree.go b/lib/btrfs/btrfstree/btree.go
index 19c7c68..cbaa6d0 100644
--- a/lib/btrfs/btrfstree/btree.go
+++ b/lib/btrfs/btrfstree/btree.go
@@ -14,7 +14,62 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 )
 
+type Forrest interface {
+	// ForrestLookup returns (nil, ErrNoTree) if the tree does not
+	// exist but there is otherwise no error.
+	ForrestLookup(ctx context.Context, treeID btrfsprim.ObjID) (Tree, error)
+}
+
 type Tree interface {
+	// TreeLookup looks up the Item for a given key.
+	//
+	// If no such Item exists, but there is otherwise no error,
+	// then ErrNoItem is returned.
+	TreeLookup(ctx context.Context, key btrfsprim.Key) (Item, error)
+
+	// TreeSearch searches the Tree for a value for which
+	// `search.Search(itemKey, itemSize) == 0`.
+	//
+	//	: + + + 0 0 0 - - -
+	//	:       ^ ^ ^
+	//	:       any of
+	//
+	// You can conceptualize `search.Search` as subtraction:
+	//
+	//	func(strawKey btrfsprim.Key, strawSize uint32) int {
+	//	    return needle - straw
+	//	}
+	//
+	// `search.Search` may be called with key/size pairs that do not
+	// correspond to an existing Item (for key-pointers in
+	// interior nodes in the tree); in which case the size
+	// math.MaxUint32.
+	//
+	// If no such Item exists, but there is otherwise no error,
+	// then an error that is ErrNoItem is returned.
+	TreeSearch(ctx context.Context, search TreeSearcher) (Item, error)
+
+	// TreeRange iterates over the Tree in order, calling
+	// `handleFn` for each Item.
+	TreeRange(ctx context.Context, handleFn func(Item) bool) error
+
+	// TreeSubrange iterates over the Tree in order, calling
+	// `handleFn` for all Items for which `search.Search` returns
+	// 0.
+	//
+	// `search.Search` may be called with key/size pairs that do
+	// not correspond to an existing Item (for key-pointers in
+	// interior nodes in the tree); in which case the size
+	// math.MaxUint32.
+	//
+	// If handleFn is called for fewer than `min` items, an error
+	// that is ErrNoItem is returned.
+	TreeSubrange(ctx context.Context,
+		min int,
+		search TreeSearcher,
+		handleFn func(Item) bool,
+	) error
+
 	// CheckOwner returns whether it is permissible for a node
 	// with .Head.Owner=owner and .Head.Generation=gen to be in
 	// this tree.
@@ -23,6 +78,33 @@ type Tree interface {
 	// specifies whether it should return an error (false) or nil
 	// (true).
 	TreeCheckOwner(ctx context.Context, failOpen bool, owner btrfsprim.ObjID, gen btrfsprim.Generation) error
+
+	// TreeWalk is a lower-level call than TreeSubrange.  Use with
+	// hesitancy.
+	//
+	// It walks a Tree, triggering callbacks for every node, key-pointer,
+	// and item; as well as for any errors encountered.
+	//
+	// If the Tree is valid, then everything is walked in key-order; but
+	// if the Tree is broken, then ordering is not guaranteed.
+	//
+	// Canceling the Context causes TreeWalk to return early; no values
+	// from the Context are used.
+	//
+	// The lifecycle of callbacks is:
+	//
+	//	000  (read superblock) (maybe cbs.BadSuperblock())
+	//
+	//	001  (read node)
+	//	002  cbs.Node() or cbs.BadNode()
+	//	     if interior:
+	//	       for kp in node.items:
+	//	003a     if cbs.PreKeyPointer == nil || cbs.PreKeyPointer() {
+	//	004b       (recurse)
+	//	     else:
+	//	       for item in node.items:
+	//	003b     cbs.Item() or cbs.BadItem()
+	TreeWalk(ctx context.Context, cbs TreeWalkHandler)
 }
 
 type TreeSearcher interface {
@@ -55,6 +137,8 @@ type TreeWalkHandler struct {
 	BadItem    func(Path, Item)
 }
 
+// Compat //////////////////////////////////////////////////////////////////////
+
 // TreeOperator is an interface for performing basic btree operations.
 type TreeOperator interface {
 	// TreeWalk walks a tree, triggering callbacks for every node,
diff --git a/lib/btrfs/btrfstree/btree_forrest.go b/lib/btrfs/btrfstree/btree_forrest.go
index 1875c2f..4a4a4c0 100644
--- a/lib/btrfs/btrfstree/btree_forrest.go
+++ b/lib/btrfs/btrfstree/btree_forrest.go
@@ -14,6 +14,8 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 )
 
+// LookupTreeRoot //////////////////////////////////////////////////////////////
+
 // A TreeRoot is more-or-less a btrfsitem.Root, but simpler; returned by
 // LookupTreeRoot.
 type TreeRoot struct {
@@ -28,8 +30,37 @@ type TreeRoot struct {
 }
 
 // LookupTreeRoot is a utility function to help with implementing the
-// 'TreeOperator' interface.
-func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*TreeRoot, error) {
+// 'Forrest' interface.
+//
+// The 'forrest' passed to LookupTreeRoot must handle:
+//
+//	forrest.ForrestLookup(ctx, btrfsprim.ROOT_TREE_OBJECTID)
+//
+// It is OK for forrest.ForrestLookup to recurse and call
+// LookupTreeRoot, as LookupTreeRoot will not call ForrestLookup for
+// ROOT_TREE_OBJECTID; so it will not be an infinite recursion.
+func LookupTreeRoot(ctx context.Context, forrest Forrest, sb Superblock, treeID btrfsprim.ObjID) (*TreeRoot, error) {
+	return OldLookupTreeRoot(
+		ctx,
+		func(rootTreeID btrfsprim.ObjID, searcher TreeSearcher) (Item, error) {
+			rootTree, err := forrest.ForrestLookup(ctx, rootTreeID)
+			if err != nil {
+				return Item{}, err
+			}
+			rootItem, err := rootTree.TreeSearch(ctx, searcher)
+			if err != nil {
+				return Item{}, err
+			}
+			return rootItem, nil
+		},
+		sb,
+		treeID,
+	)
+}
+
+// OldLookupTreeRoot is a utility function to help with implementing
+// the old 'TreeOperator' interface.
+func OldLookupTreeRoot(_ context.Context, treeSearch func(treeID btrfsprim.ObjID, _ TreeSearcher) (Item, error), sb Superblock, treeID btrfsprim.ObjID) (*TreeRoot, error) {
 	switch treeID {
 	case btrfsprim.ROOT_TREE_OBJECTID:
 		return &TreeRoot{
@@ -60,7 +91,7 @@ func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID bt
 			Generation: sb.BlockGroupRootGeneration,
 		}, nil
 	default:
-		rootItem, err := fs.TreeSearch(btrfsprim.ROOT_TREE_OBJECTID, SearchRootItem(treeID))
+		rootItem, err := treeSearch(btrfsprim.ROOT_TREE_OBJECTID, SearchRootItem(treeID))
 		if err != nil {
 			if errors.Is(err, ErrNoItem) {
 				err = fmt.Errorf("%w: %s", ErrNoTree, err)
@@ -87,28 +118,50 @@ func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID bt
 	}
 }
 
-type TreeOperatorImpl struct {
+// RawForrest //////////////////////////////////////////////////////////////////
+
+// RawForrest implements Forrest.
+type RawForrest struct {
 	NodeSource NodeSource
 }
 
-func (fs TreeOperatorImpl) RawTree(ctx context.Context, treeID btrfsprim.ObjID) (*RawTree, error) {
-	sb, err := fs.NodeSource.Superblock()
+var _ Forrest = RawForrest{}
+
+// RawTree is a variant of ForrestLookup that returns a concrete type
+// instead of an interface.
+func (forrest RawForrest) RawTree(ctx context.Context, treeID btrfsprim.ObjID) (*RawTree, error) {
+	sb, err := forrest.NodeSource.Superblock()
 	if err != nil {
 		return nil, err
 	}
-	rootInfo, err := LookupTreeRoot(ctx, fs, *sb, treeID)
+	rootInfo, err := LookupTreeRoot(ctx, forrest, *sb, treeID)
 	if err != nil {
 		return nil, err
 	}
 	return &RawTree{
-		Forrest:  fs,
+		Forrest:  forrest,
 		TreeRoot: *rootInfo,
 	}, nil
 }
 
+// ForrestLookup implements Forrest.
+func (forrest RawForrest) ForrestLookup(ctx context.Context, treeID btrfsprim.ObjID) (Tree, error) {
+	tree, err := forrest.RawTree(ctx, treeID)
+	if err != nil {
+		return nil, err
+	}
+	return tree, nil
+}
+
+// Compat //////////////////////////////////////////////////////////////////////
+
+var _ TreeOperator = RawForrest{}
+
+type TreeOperatorImpl = RawForrest
+
 // TreeWalk implements the 'TreeOperator' interface.
-func (fs TreeOperatorImpl) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, errHandle func(*TreeError), cbs TreeWalkHandler) {
-	tree, err := fs.RawTree(ctx, treeID)
+func (forrest RawForrest) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, errHandle func(*TreeError), cbs TreeWalkHandler) {
+	tree, err := forrest.RawTree(ctx, treeID)
 	if err != nil {
 		errHandle(&TreeError{Path: Path{PathRoot{TreeID: treeID}}, Err: err})
 		return
@@ -117,9 +170,9 @@ func (fs TreeOperatorImpl) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID,
 }
 
 // TreeLookup implements the 'TreeOperator' interface.
-func (fs TreeOperatorImpl) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key) (Item, error) {
+func (forrest RawForrest) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key) (Item, error) {
 	ctx := context.TODO()
-	tree, err := fs.RawTree(ctx, treeID)
+	tree, err := forrest.RawTree(ctx, treeID)
 	if err != nil {
 		return Item{}, err
 	}
@@ -127,9 +180,9 @@ func (fs TreeOperatorImpl) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key)
 }
 
 // TreeSearch implements the 'TreeOperator' interface.
-func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, searcher TreeSearcher) (Item, error) {
+func (forrest RawForrest) TreeSearch(treeID btrfsprim.ObjID, searcher TreeSearcher) (Item, error) {
 	ctx := context.TODO()
-	tree, err := fs.RawTree(ctx, treeID)
+	tree, err := forrest.RawTree(ctx, treeID)
 	if err != nil {
 		return Item{}, err
 	}
@@ -137,9 +190,9 @@ func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, searcher TreeSearc
 }
 
 // TreeSearchAll implements the 'TreeOperator' interface.
-func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, searcher TreeSearcher) ([]Item, error) {
+func (forrest RawForrest) TreeSearchAll(treeID btrfsprim.ObjID, searcher TreeSearcher) ([]Item, error) {
 	ctx := context.TODO()
-	tree, err := fs.RawTree(ctx, treeID)
+	tree, err := forrest.RawTree(ctx, treeID)
 	if err != nil {
 		return nil, err
 	}
diff --git a/lib/btrfs/btrfstree/btree_tree.go b/lib/btrfs/btrfstree/btree_tree.go
index 27d6548..7cfa257 100644
--- a/lib/btrfs/btrfstree/btree_tree.go
+++ b/lib/btrfs/btrfstree/btree_tree.go
@@ -16,11 +16,15 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
+// RawTree implements Tree.
 type RawTree struct {
-	Forrest TreeOperatorImpl
+	Forrest RawForrest
 	TreeRoot
 }
 
+var _ Tree = (*RawTree)(nil)
+
+// TreeWalk implements the 'Tree' interface.
 func (tree *RawTree) TreeWalk(ctx context.Context, cbs TreeWalkHandler) {
 	sb, err := tree.Forrest.NodeSource.Superblock()
 	if err != nil {
@@ -161,6 +165,7 @@ func searchKP(haystack []KeyPointer, searchFn func(key btrfsprim.Key, size uint3
 	return KeyPointer{}, false
 }
 
+// TreeSearch implements the 'Tree' interface.
 func (tree *RawTree) TreeSearch(ctx context.Context, searcher TreeSearcher) (Item, error) {
 	ctx, cancel := context.WithCancel(ctx)
 	var retErr error
@@ -206,10 +211,40 @@ func (tree *RawTree) TreeSearch(ctx context.Context, searcher TreeSearcher) (Ite
 	return ret, retErr
 }
 
+// TreeLookup implements the 'Tree' interface.
 func (tree *RawTree) TreeLookup(ctx context.Context, key btrfsprim.Key) (Item, error) {
 	return tree.TreeSearch(ctx, SearchExactKey(key))
 }
 
+// TreeRange implements the 'Tree' interface.
+func (tree *RawTree) TreeRange(ctx context.Context, handleFn func(Item) bool) error {
+	ctx, cancel := context.WithCancel(ctx)
+	var errs derror.MultiError
+
+	tree.TreeWalk(ctx, TreeWalkHandler{
+		BadNode: func(path Path, _ *Node, err error) bool {
+			errs = append(errs, fmt.Errorf("%v: %w", path, err))
+			return false
+		},
+		Item: func(_ Path, item Item) {
+			if !handleFn(item) {
+				cancel()
+			}
+		},
+		BadItem: func(_ Path, item Item) {
+			if !handleFn(item) {
+				cancel()
+			}
+		},
+	})
+
+	if len(errs) > 0 {
+		return errs
+	}
+	return nil
+}
+
+// TreeSubrange implements the 'Tree' interface.
 func (tree *RawTree) TreeSubrange(ctx context.Context, min int, searcher TreeSearcher, handleFn func(Item) bool) error {
 	ctx, cancel := context.WithCancel(ctx)
 	var errs derror.MultiError
diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index 030ea41..59e4c1d 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -12,6 +12,7 @@ import (
 	"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/diskio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
@@ -73,6 +74,21 @@ func (fs *FS) readNode(_ context.Context, addr btrfsvol.LogicalAddr, nodeEntry *
 
 var _ btrfstree.NodeSource = (*FS)(nil)
 
+// btrfstree.Forrest ///////////////////////////////////////////////////////////
+
+// RawTree is a variant of ForrestLookup that returns a concrete type
+// instead of an interface.
+func (fs *FS) RawTree(ctx context.Context, treeID btrfsprim.ObjID) (*btrfstree.RawTree, error) {
+	return btrfstree.RawForrest{NodeSource: fs}.RawTree(ctx, treeID)
+}
+
+// ForrestLookup implements btree.Forrest.
+func (fs *FS) ForrestLookup(ctx context.Context, treeID btrfsprim.ObjID) (btrfstree.Tree, error) {
+	return btrfstree.RawForrest{NodeSource: fs}.ForrestLookup(ctx, treeID)
+}
+
+var _ btrfstree.Forrest = (*FS)(nil)
+
 // btrfstree.TreeOperator //////////////////////////////////////////////////////
 
 // TreeWalk implements btrfstree.TreeOperator.
@@ -96,3 +112,18 @@ func (fs *FS) TreeSearchAll(treeID btrfsprim.ObjID, searcher btrfstree.TreeSearc
 }
 
 var _ btrfstree.TreeOperator = (*FS)(nil)
+
+// ReadableFS //////////////////////////////////////////////////////////////////
+
+type ReadableFS interface {
+	btrfstree.Forrest
+
+	Superblock() (*btrfstree.Superblock, error)
+
+	Name() string
+
+	// For reading file contents.
+	diskio.ReaderAt[btrfsvol.LogicalAddr]
+}
+
+var _ ReadableFS = (*FS)(nil)
diff --git a/lib/btrfs/io4_fs.go b/lib/btrfs/io4_fs.go
index a611f35..57eca57 100644
--- a/lib/btrfs/io4_fs.go
+++ b/lib/btrfs/io4_fs.go
@@ -102,7 +102,7 @@ func NewSubvolume(
 		sv.rootErr = err
 		return sv
 	}
-	rootInfo, err := btrfstree.LookupTreeRoot(ctx, sv.fs, *sb, sv.TreeID)
+	rootInfo, err := btrfstree.OldLookupTreeRoot(ctx, sv.fs.TreeSearch, *sb, sv.TreeID)
 	if err != nil {
 		sv.rootErr = err
 		return sv
diff --git a/lib/btrfsutil/old_rebuilt_forrest.go b/lib/btrfsutil/old_rebuilt_forrest.go
index d17aa4a..e853be7 100644
--- a/lib/btrfsutil/old_rebuilt_forrest.go
+++ b/lib/btrfsutil/old_rebuilt_forrest.go
@@ -166,7 +166,7 @@ func (bt *OldRebuiltForrest) rawTreeWalk(ctx context.Context, treeID btrfsprim.O
 		cacheEntry.RootErr = err
 		return
 	}
-	root, err := btrfstree.LookupTreeRoot(ctx, bt, *sb, treeID)
+	root, err := btrfstree.OldLookupTreeRoot(ctx, bt.TreeSearch, *sb, treeID)
 	if err != nil {
 		cacheEntry.RootErr = err
 		return
@@ -263,20 +263,21 @@ func (bt *OldRebuiltForrest) readNode(nodeInfo nodeInfo) *btrfstree.Node {
 
 // TreeLookup implements btrfstree.TreeOperator.
 func (bt *OldRebuiltForrest) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key) (btrfstree.Item, error) {
-	return bt.RebuiltTree(bt.ctx, treeID).treeLookup(bt.ctx, key)
+	return bt.RebuiltTree(bt.ctx, treeID).TreeLookup(bt.ctx, key)
 }
 
-func (tree oldRebuiltTree) treeLookup(ctx context.Context, key btrfsprim.Key) (btrfstree.Item, error) {
-	return tree.treeSearch(ctx, btrfstree.SearchExactKey(key))
+// TreeLookup implements btrfstree.Tree.
+func (tree oldRebuiltTree) TreeLookup(ctx context.Context, key btrfsprim.Key) (btrfstree.Item, error) {
+	return tree.TreeSearch(ctx, btrfstree.SearchExactKey(key))
 }
 
 // TreeSearch implements btrfstree.TreeOperator.
 func (bt *OldRebuiltForrest) TreeSearch(treeID btrfsprim.ObjID, searcher btrfstree.TreeSearcher) (btrfstree.Item, error) {
-	return bt.RebuiltTree(bt.ctx, treeID).treeSearch(bt.ctx, searcher)
+	return bt.RebuiltTree(bt.ctx, treeID).TreeSearch(bt.ctx, searcher)
 }
 
 // TreeSearch implements btrfstree.Tree.
-func (tree oldRebuiltTree) treeSearch(_ context.Context, searcher btrfstree.TreeSearcher) (btrfstree.Item, error) {
+func (tree oldRebuiltTree) TreeSearch(_ context.Context, searcher btrfstree.TreeSearcher) (btrfstree.Item, error) {
 	if tree.RootErr != nil {
 		return btrfstree.Item{}, tree.RootErr
 	}
@@ -299,6 +300,26 @@ func (tree oldRebuiltTree) treeSearch(_ context.Context, searcher btrfstree.Tree
 	return item, nil
 }
 
+// TreeRange implements btrfstree.Tree.
+func (tree oldRebuiltTree) TreeRange(_ context.Context, handleFn func(btrfstree.Item) bool) error {
+	if tree.RootErr != nil {
+		return tree.RootErr
+	}
+
+	var node *btrfstree.Node
+	tree.Items.Range(
+		func(rbnode *containers.RBNode[oldRebuiltTreeValue]) bool {
+			if node == nil || node.Head.Addr != rbnode.Value.Node.LAddr {
+				tree.forrest.inner.ReleaseNode(node)
+				node = tree.forrest.readNode(rbnode.Value.Node)
+			}
+			return handleFn(node.BodyLeaf[rbnode.Value.Slot])
+		})
+	tree.forrest.inner.ReleaseNode(node)
+
+	return tree.addErrs(func(btrfsprim.Key, uint32) int { return 0 }, nil)
+}
+
 // TreeSearchAll implements btrfstree.TreeOperator.
 func (bt *OldRebuiltForrest) TreeSearchAll(treeID btrfsprim.ObjID, searcher btrfstree.TreeSearcher) ([]btrfstree.Item, error) {
 	tree := bt.RebuiltTree(bt.ctx, treeID)
@@ -307,7 +328,7 @@ func (bt *OldRebuiltForrest) TreeSearchAll(treeID btrfsprim.ObjID, searcher btrf
 	}
 
 	var ret []btrfstree.Item
-	err := tree.treeSubrange(bt.ctx, 1, searcher, func(item btrfstree.Item) bool {
+	err := tree.TreeSubrange(bt.ctx, 1, searcher, func(item btrfstree.Item) bool {
 		item.Body = item.Body.CloneItem()
 		ret = append(ret, item)
 		return true
@@ -316,7 +337,8 @@ func (bt *OldRebuiltForrest) TreeSearchAll(treeID btrfsprim.ObjID, searcher btrf
 	return ret, err
 }
 
-func (tree oldRebuiltTree) treeSubrange(_ context.Context, min int, searcher btrfstree.TreeSearcher, handleFn func(btrfstree.Item) bool) error {
+// TreeSubrange implements btrfstree.Tree.
+func (tree oldRebuiltTree) TreeSubrange(_ context.Context, min int, searcher btrfstree.TreeSearcher, handleFn func(btrfstree.Item) bool) error {
 	var node *btrfstree.Node
 	var cnt int
 	tree.Items.Subrange(
@@ -355,10 +377,12 @@ func (bt *OldRebuiltForrest) TreeWalk(ctx context.Context, treeID btrfsprim.ObjI
 		})
 		return
 	}
-	tree.treeWalk(ctx, cbs)
+	tree.TreeWalk(ctx, cbs)
 }
 
-func (tree oldRebuiltTree) treeWalk(ctx context.Context, cbs btrfstree.TreeWalkHandler) {
+// TreeWalk implements btrfstree.Tree.  It doesn't actually visit
+// nodes or keypointers (just items).
+func (tree oldRebuiltTree) TreeWalk(ctx context.Context, cbs btrfstree.TreeWalkHandler) {
 	if cbs.Item == nil && cbs.BadItem == nil {
 		return
 	}
@@ -439,7 +463,7 @@ func (tree oldRebuiltTree) TreeCheckOwner(ctx context.Context, failOpen bool, ow
 				return nil //nolint:nilerr // fail open
 			}
 		}
-		parentIDItem, err := uuidTree.treeLookup(ctx, btrfsitem.UUIDToKey(tree.ParentUUID))
+		parentIDItem, err := uuidTree.TreeLookup(ctx, btrfsitem.UUIDToKey(tree.ParentUUID))
 		if err != nil {
 			if failOpen {
 				return nil