5 files changed, 33 insertions, 21 deletions
diff --git a/lib/btrfs/btrfstree/btree.go b/lib/btrfs/btrfstree/btree.go
index bc1bac2..dde485a 100644
--- a/lib/btrfs/btrfstree/btree.go
+++ b/lib/btrfs/btrfstree/btree.go
@@ -11,6 +11,7 @@ import (
 	"fmt"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 )
 
 type TreeSearcher interface {
@@ -96,5 +97,5 @@ func (e *TreeError) Error() string {
 
 type NodeSource interface {
 	Superblock() (*Superblock, error)
-	ReadNode(Path) (*Node, error)
+	ReadNode(ctx context.Context, addr btrfsvol.LogicalAddr, exp NodeExpectations) (*Node, error)
 }
diff --git a/lib/btrfs/btrfstree/btree_forrest.go b/lib/btrfs/btrfstree/btree_forrest.go
index b017888..b04bfc0 100644
--- a/lib/btrfs/btrfstree/btree_forrest.go
+++ b/lib/btrfs/btrfstree/btree_forrest.go
@@ -83,7 +83,10 @@ func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID bt
 }
 
 type TreeOperatorImpl struct {
-	NodeSource NodeSource
+	NodeSource interface {
+		NodeSource
+		NodeFile
+	}
 }
 
 func (fs TreeOperatorImpl) RawTree(ctx context.Context, treeID btrfsprim.ObjID) (*RawTree, error) {
diff --git a/lib/btrfs/btrfstree/btree_tree.go b/lib/btrfs/btrfstree/btree_tree.go
index 89b4c52..a6311f9 100644
--- a/lib/btrfs/btrfstree/btree_tree.go
+++ b/lib/btrfs/btrfstree/btree_tree.go
@@ -22,6 +22,9 @@ type RawTree struct {
 }
 
 func (tree *RawTree) TreeWalk(ctx context.Context, cbs TreeWalkHandler) {
+	if tree.RootNode == 0 {
+		return
+	}
 	path := Path{{
 		FromTree:         tree.ID,
 		FromItemSlot:     -1,
@@ -37,12 +40,13 @@ func (tree *RawTree) walk(ctx context.Context, path Path, cbs TreeWalkHandler) {
 	if ctx.Err() != nil {
 		return
 	}
-	if path.Node(-1).ToNodeAddr == 0 {
-		return
-	}
 
 	// 001
-	node, err := tree.Forrest.NodeSource.ReadNode(path)
+	nodeAddr, nodeExp, ok := path.NodeExpectations(tree.Forrest.NodeSource)
+	if !ok {
+		return
+	}
+	node, err := tree.Forrest.NodeSource.ReadNode(ctx, nodeAddr, nodeExp)
 	defer node.Free()
 	if ctx.Err() != nil {
 		return
diff --git a/lib/btrfs/btrfstree/readnode.go b/lib/btrfs/btrfstree/readnode.go
index ac82c62..a4ccf10 100644
--- a/lib/btrfs/btrfstree/readnode.go
+++ b/lib/btrfs/btrfstree/readnode.go
@@ -14,8 +14,7 @@ import (
 )
 
 type NodeFile interface {
-	diskio.File[btrfsvol.LogicalAddr]
-	Superblock() (*Superblock, error)
+	diskio.ReaderAt[btrfsvol.LogicalAddr]
 
 	// ParentTree, given a tree ID, returns that tree's parent
 	// tree, if it has one.
@@ -29,15 +28,14 @@ type NodeFile interface {
 	ParentTree(btrfsprim.ObjID) (btrfsprim.ObjID, btrfsprim.Generation, bool)
 }
 
-// FSReadNode is a utility function to help with implementing the
-// 'NodeSource' interface.
-func FSReadNode(
-	fs NodeFile,
-	path Path,
-) (*Node, error) {
-	sb, err := fs.Superblock()
-	if err != nil {
-		return nil, fmt.Errorf("btrfs.FS.ReadNode: %w", err)
+// NodeExpectations returns the address to read and the expectations
+// to have when reading the node pointed to by this Path.
+//
+// `ok` is false if the path is empty or if this Path points to an
+// item rather than a node.
+func (path Path) NodeExpectations(fs NodeFile) (_ btrfsvol.LogicalAddr, _ NodeExpectations, ok bool) {
+	if path.Node(-1).ToNodeAddr == 0 && path.Node(-1).ToNodeGeneration == 0 && path.Node(-1).ToNodeLevel == 0 {
+		return 0, NodeExpectations{}, false
 	}
 
 	checkOwner := func(owner btrfsprim.ObjID, gen btrfsprim.Generation) error {
@@ -70,12 +68,12 @@ func FSReadNode(
 		}
 	}
 
-	return ReadNode[btrfsvol.LogicalAddr](fs, *sb, path.Node(-1).ToNodeAddr, NodeExpectations{
+	return path.Node(-1).ToNodeAddr, NodeExpectations{
 		LAddr:      containers.OptionalValue(path.Node(-1).ToNodeAddr),
 		Level:      containers.OptionalValue(path.Node(-1).ToNodeLevel),
 		Generation: containers.OptionalValue(path.Node(-1).ToNodeGeneration),
 		Owner:      checkOwner,
 		MinItem:    containers.OptionalValue(path.Node(-1).ToKey),
 		MaxItem:    containers.OptionalValue(path.Node(-1).ToMaxKey),
-	})
+	}, true
 }
diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index f1436fb..bb68d7d 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -6,10 +6,12 @@ package btrfs
 
 import (
 	"context"
+	"fmt"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 )
 
 // This file is ordered from low-level to high-level.
@@ -79,8 +81,12 @@ var _ btrfstree.NodeFile = (*FS)(nil)
 // btrfstree.NodeSource ////////////////////////////////////////////////////////
 
 // ReadNode implements btrfstree.NodeSource.
-func (fs *FS) ReadNode(path btrfstree.Path) (*btrfstree.Node, error) {
-	return btrfstree.FSReadNode(fs, path)
+func (fs *FS) ReadNode(_ context.Context, addr btrfsvol.LogicalAddr, exp btrfstree.NodeExpectations) (*btrfstree.Node, error) {
+	sb, err := fs.Superblock()
+	if err != nil {
+		return nil, fmt.Errorf("btrfs.FS.ReadNode: %w", err)
+	}
+	return btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, addr, exp)
 }
 
 var _ btrfstree.NodeSource = (*FS)(nil)