Have WalkTree include path information

3 files changed, 103 insertions, 43 deletions

diff --git a/pkg/btrfs/io2_fs.go b/pkg/btrfs/io2_fs.go
index 1a3a9df..588a60a 100644
--- a/pkg/btrfs/io2_fs.go
+++ b/pkg/btrfs/io2_fs.go
@@ -128,13 +128,13 @@ func (fs *FS) Init() error {
 			fs.chunks = append(fs.chunks, chunk)
 		}
 		if err := fs.WalkTree(sb.Data.ChunkTree, WalkTreeHandler{
-			Item: func(key Key, body btrfsitem.Item) error {
-				if key.ItemType != btrfsitem.CHUNK_ITEM_KEY {
+			Item: func(_ WalkTreePath, item Item) error {
+				if item.Head.Key.ItemType != btrfsitem.CHUNK_ITEM_KEY {
 					return nil
 				}
 				fs.chunks = append(fs.chunks, SysChunk{
-					Key:   key,
-					Chunk: body.(btrfsitem.Chunk),
+					Key:   item.Head.Key,
+					Chunk: item.Body.(btrfsitem.Chunk),
 				})
 				return nil
 			},
diff --git a/pkg/btrfs/types_btree.go b/pkg/btrfs/types_btree.go
index 8a32d9d..10babfd 100644
--- a/pkg/btrfs/types_btree.go
+++ b/pkg/btrfs/types_btree.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	iofs "io/fs"
+	"strings"
 
 	"lukeshu.com/btrfs-tools/pkg/binstruct"
 	"lukeshu.com/btrfs-tools/pkg/btrfs/btrfsitem"
@@ -381,33 +382,89 @@ func (fs *FS) ReadNode(addr LogicalAddr) (*util.Ref[LogicalAddr, Node], error) {
 	})
 }
 
+// A WalkTreePathElem essentially represents a KeyPointer.
+type WalkTreePathElem struct {
+	// ItemIdx is the index of this KeyPointer in the parent Node;
+	// or -1 if this is the root and there is no KeyPointer.
+	ItemIdx int
+	// NodeAddr is the address of the node that the KeyPointer
+	// points at, or 0 if this is a leaf item and nothing is
+	// being pointed at.
+	NodeAddr LogicalAddr
+}
+
+// - The first element will always have an ItemIdx of -1.
+//
+// - For .Item() callbacks, the last element will always have a
+//   NodeAddr of 0.
+type WalkTreePath []WalkTreePathElem
+
+func (path WalkTreePath) String() string {
+	if len(path) == 0 {
+		return "(empty-path)"
+	}
+	var ret strings.Builder
+	fmt.Fprintf(&ret, "node@%v", path[0].NodeAddr)
+	for _, elem := range path[1:] {
+		fmt.Fprintf(&ret, "[%v]", elem.ItemIdx)
+		if elem.NodeAddr != 0 {
+			fmt.Fprintf(&ret, "->node@%v", elem.NodeAddr)
+		}
+	}
+	return ret.String()
+}
+
 type WalkTreeHandler struct {
 	// Callbacks for entire nodes
-	PreNode  func(LogicalAddr) error
-	Node     func(*util.Ref[LogicalAddr, Node], error) error
-	PostNode func(*util.Ref[LogicalAddr, Node]) error
+	PreNode  func(WalkTreePath) error
+	Node     func(WalkTreePath, *util.Ref[LogicalAddr, Node], error) error
+	PostNode func(WalkTreePath, *util.Ref[LogicalAddr, Node]) error
 	// Callbacks for items on internal nodes
-	PreKeyPointer  func(KeyPointer) error
-	PostKeyPointer func(KeyPointer) error
+	PreKeyPointer  func(WalkTreePath, KeyPointer) error
+	PostKeyPointer func(WalkTreePath, KeyPointer) error
 	// Callbacks for items on leaf nodes
-	Item func(Key, btrfsitem.Item) error
+	Item func(WalkTreePath, Item) error
 }
 
+// The lifecycle of callbacks is:
+//
+//     001 .PreNode()
+//     002 (read node)
+//     003 .Node()
+//         for item in node.items:
+//           if internal:
+//     004     .PreKeyPointer()
+//     005     (recurse)
+//     006     .PostKeyPointer()
+//           else:
+//     004     .Item()
+//     007 .PostNode()
 func (fs *FS) WalkTree(nodeAddr LogicalAddr, cbs WalkTreeHandler) error {
-	if nodeAddr == 0 {
+	path := WalkTreePath{
+		WalkTreePathElem{
+			ItemIdx:  -1,
+			NodeAddr: nodeAddr,
+		},
+	}
+	return fs.walkTree(path, cbs)
+}
+
+func (fs *FS) walkTree(path WalkTreePath, cbs WalkTreeHandler) error {
+	if path[len(path)-1].NodeAddr == 0 {
 		return nil
 	}
+
 	if cbs.PreNode != nil {
-		if err := cbs.PreNode(nodeAddr); err != nil {
+		if err := cbs.PreNode(path); err != nil {
 			if errors.Is(err, iofs.SkipDir) {
 				return nil
 			}
 			return err
 		}
 	}
-	node, err := fs.ReadNode(nodeAddr)
+	node, err := fs.ReadNode(path[len(path)-1].NodeAddr)
 	if cbs.Node != nil {
-		err = cbs.Node(node, err)
+		err = cbs.Node(path, node, err)
 	}
 	if err != nil {
 		if errors.Is(err, iofs.SkipDir) {
@@ -416,20 +473,24 @@ func (fs *FS) WalkTree(nodeAddr LogicalAddr, cbs WalkTreeHandler) error {
 		return fmt.Errorf("btrfs.FS.WalkTree: %w", err)
 	}
 	if node != nil {
-		for _, item := range node.Data.BodyInternal {
+		for i, item := range node.Data.BodyInternal {
+			itemPath := append(path, WalkTreePathElem{
+				ItemIdx:  i,
+				NodeAddr: item.BlockPtr,
+			})
 			if cbs.PreKeyPointer != nil {
-				if err := cbs.PreKeyPointer(item); err != nil {
+				if err := cbs.PreKeyPointer(itemPath, item); err != nil {
 					if errors.Is(err, iofs.SkipDir) {
 						continue
 					}
 					return err
 				}
 			}
-			if err := fs.WalkTree(item.BlockPtr, cbs); err != nil {
+			if err := fs.walkTree(itemPath, cbs); err != nil {
 				return err
 			}
 			if cbs.PostKeyPointer != nil {
-				if err := cbs.PostKeyPointer(item); err != nil {
+				if err := cbs.PostKeyPointer(itemPath, item); err != nil {
 					if errors.Is(err, iofs.SkipDir) {
 						continue
 					}
@@ -437,9 +498,12 @@ func (fs *FS) WalkTree(nodeAddr LogicalAddr, cbs WalkTreeHandler) error {
 				}
 			}
 		}
-		for _, item := range node.Data.BodyLeaf {
+		for i, item := range node.Data.BodyLeaf {
 			if cbs.Item != nil {
-				if err := cbs.Item(item.Head.Key, item.Body); err != nil {
+				itemPath := append(path, WalkTreePathElem{
+					ItemIdx: i,
+				})
+				if err := cbs.Item(itemPath, item); err != nil {
 					if errors.Is(err, iofs.SkipDir) {
 						continue
 					}
@@ -449,7 +513,7 @@ func (fs *FS) WalkTree(nodeAddr LogicalAddr, cbs WalkTreeHandler) error {
 		}
 	}
 	if cbs.PostNode != nil {
-		if err := cbs.PostNode(node); err != nil {
+		if err := cbs.PostNode(path, node); err != nil {
 			if errors.Is(err, iofs.SkipDir) {
 				return nil
 			}
diff --git a/pkg/btrfsmisc/print_tree.go b/pkg/btrfsmisc/print_tree.go
index 41e86a0..fa864a5 100644
--- a/pkg/btrfsmisc/print_tree.go
+++ b/pkg/btrfsmisc/print_tree.go
@@ -14,29 +14,25 @@ import (
 // kernel-shared/print-tree.c:btrfs_print_tree() and
 // kernel-shared/print-tree.c:btrfs_print_leaf()
 func PrintTree(fs *btrfs.FS, root btrfs.LogicalAddr) error {
-	nodeRef, err := fs.ReadNode(root)
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "error: %v\n", err)
-	}
-	if nodeRef == nil {
-		return nil
-	}
-	node := nodeRef.Data
-	printHeaderInfo(node)
-	if node.Head.Level > 0 { // internal
-		for _, item := range node.BodyInternal {
+	return fs.WalkTree(root, btrfs.WalkTreeHandler{
+		Node: func(path btrfs.WalkTreePath, nodeRef *util.Ref[btrfs.LogicalAddr, btrfs.Node], err error) error {
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "error: %v: %v\n", path, err)
+			}
+			if nodeRef != nil {
+				printHeaderInfo(nodeRef.Data)
+			}
+			return nil
+		},
+		PreKeyPointer: func(_ btrfs.WalkTreePath, item btrfs.KeyPointer) error {
 			fmt.Printf("\t%v block %v gen %v\n",
 				FmtKey(item.Key),
 				item.BlockPtr,
 				item.Generation)
-		}
-		for _, item := range node.BodyInternal {
-			if err := PrintTree(fs, item.BlockPtr); err != nil {
-				return err
-			}
-		}
-	} else { // leaf
-		for i, item := range node.BodyLeaf {
+			return nil
+		},
+		Item: func(path btrfs.WalkTreePath, item btrfs.Item) error {
+			i := path[len(path)-1].ItemIdx
 			fmt.Printf("\titem %v %v itemoff %v itemsize %v\n",
 				i,
 				FmtKey(item.Head.Key),
@@ -245,9 +241,9 @@ func PrintTree(fs *btrfs.FS, root btrfs.LogicalAddr) error {
 			default:
 				fmt.Printf("\t\t(error) unhandled item type: %T\n", body)
 			}
-		}
-	}
-	return nil
+			return nil
+		},
+	})
 }
 
 // printHeaderInfo mimics btrfs-progs kernel-shared/print-tree.c:print_header_info()