2 files changed, 40 insertions, 81 deletions

diff --git a/pkg/btrfsmisc/print_tree.go b/pkg/btrfsmisc/print_tree.go
index 8bfcbe5..cc3055d 100644
--- a/pkg/btrfsmisc/print_tree.go
+++ b/pkg/btrfsmisc/print_tree.go
@@ -15,8 +15,8 @@ import (
 // PrintTree mimics btrfs-progs
 // kernel-shared/print-tree.c:btrfs_print_tree() and
 // kernel-shared/print-tree.c:btrfs_print_leaf()
-func PrintTree(fs *btrfs.FS, root btrfsvol.LogicalAddr) error {
-	return fs.TreeWalk(root, btrfs.TreeWalkHandler{
+func PrintTree(fs *btrfs.FS, treeID btrfs.ObjID) error {
+	return fs.TreeWalk(treeID, btrfs.TreeWalkHandler{
 		Node: func(path btrfs.TreePath, nodeRef *util.Ref[btrfsvol.LogicalAddr, btrfs.Node], err error) error {
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "error: %v: %v\n", path, err)
diff --git a/pkg/btrfsmisc/walk.go b/pkg/btrfsmisc/walk.go
index 231f46b..55adf1d 100644
--- a/pkg/btrfsmisc/walk.go
+++ b/pkg/btrfsmisc/walk.go
@@ -24,20 +24,20 @@ func (e WalkErr) Error() string {
 	return fmt.Sprintf("%v: %v: %v", e.TreeName, e.Path, e.Err)
 }
 
-type WalkFSHandler struct {
+type WalkAllTreesHandler struct {
 	Err func(error)
 	// Callbacks for entire trees
-	PreTree  func(name string, id btrfs.ObjID, laddr btrfsvol.LogicalAddr)
-	PostTree func(name string, id btrfs.ObjID, laddr btrfsvol.LogicalAddr)
+	PreTree  func(name string, id btrfs.ObjID)
+	PostTree func(name string, id btrfs.ObjID)
 	// Callbacks for nodes or smaller
 	UnsafeNodes bool
 	btrfs.TreeWalkHandler
 }
 
-// WalkFS walks all trees in a *btrfs.FS.  Rather than returning an
-// error, it calls errCb each time an error is encountered.  The error
-// will always be of type WalkErr.
-func WalkFS(fs *btrfs.FS, cbs WalkFSHandler) {
+// WalkAllTrees walks all trees in a *btrfs.FS.  Rather than returning
+// an error, it calls errCb each time an error is encountered.  The
+// error will always be of type WalkErr.
+func WalkAllTrees(fs *btrfs.FS, cbs WalkAllTreesHandler) {
 	var treeName string
 	handleErr := func(path btrfs.TreePath, err error) {
 		cbs.Err(WalkErr{
@@ -47,29 +47,38 @@ func WalkFS(fs *btrfs.FS, cbs WalkFSHandler) {
 		})
 	}
 
-	var foundTrees []struct {
+	trees := []struct {
 		Name string
 		ID   btrfs.ObjID
-		Root btrfsvol.LogicalAddr
+	}{
+		{
+			Name: "root tree",
+			ID:   btrfs.ROOT_TREE_OBJECTID,
+		},
+		{
+			Name: "chunk tree",
+			ID:   btrfs.CHUNK_TREE_OBJECTID,
+		},
+		{
+			Name: "log tree",
+			ID:   btrfs.TREE_LOG_OBJECTID,
+		},
+		{
+			Name: "block group tree",
+			ID:   btrfs.BLOCK_GROUP_TREE_OBJECTID,
+		},
 	}
 	origItem := cbs.Item
 	cbs.Item = func(path btrfs.TreePath, item btrfs.Item) error {
 		if item.Head.Key.ItemType == btrfsitem.ROOT_ITEM_KEY {
-			root, ok := item.Body.(btrfsitem.Root)
-			if !ok {
-				handleErr(path, fmt.Errorf("ROOT_ITEM_KEY is a %T, not a btrfsitem.Root", item.Body))
-			} else {
-				foundTrees = append(foundTrees, struct {
-					Name string
-					ID   btrfs.ObjID
-					Root btrfsvol.LogicalAddr
-				}{
-					Name: fmt.Sprintf("tree %v (via %v %v)",
-						item.Head.Key.ObjectID.Format(0), treeName, path),
-					ID:   item.Head.Key.ObjectID,
-					Root: root.ByteNr,
-				})
-			}
+			trees = append(trees, struct {
+				Name string
+				ID   btrfs.ObjID
+			}{
+				Name: fmt.Sprintf("tree %v (via %v %v)",
+					item.Head.Key.ObjectID.Format(0), treeName, path),
+				ID: item.Head.Key.ObjectID,
+			})
 		}
 		if origItem != nil {
 			return origItem(path, item)
@@ -90,67 +99,17 @@ func WalkFS(fs *btrfs.FS, cbs WalkFSHandler) {
 		}
 	}
 
-	treeName = "superblock"
-	superblock, err := fs.Superblock()
-	if err != nil {
-		handleErr(nil, err)
-		return
-	}
-
-	treeName = "root tree"
-	if cbs.PreTree != nil {
-		cbs.PreTree(treeName, btrfs.ROOT_TREE_OBJECTID, superblock.Data.RootTree)
-	}
-	if err := fs.TreeWalk(superblock.Data.RootTree, cbs.TreeWalkHandler); err != nil {
-		handleErr(nil, err)
-	}
-	if cbs.PostTree != nil {
-		cbs.PostTree(treeName, btrfs.ROOT_TREE_OBJECTID, superblock.Data.RootTree)
-	}
-
-	treeName = "chunk tree"
-	if cbs.PreTree != nil {
-		cbs.PreTree(treeName, btrfs.CHUNK_TREE_OBJECTID, superblock.Data.ChunkTree)
-	}
-	if err := fs.TreeWalk(superblock.Data.ChunkTree, cbs.TreeWalkHandler); err != nil {
-		handleErr(nil, err)
-	}
-	if cbs.PostTree != nil {
-		cbs.PostTree(treeName, btrfs.CHUNK_TREE_OBJECTID, superblock.Data.ChunkTree)
-	}
-
-	treeName = "log tree"
-	if cbs.PreTree != nil {
-		cbs.PreTree(treeName, btrfs.TREE_LOG_OBJECTID, superblock.Data.LogTree)
-	}
-	if err := fs.TreeWalk(superblock.Data.LogTree, cbs.TreeWalkHandler); err != nil {
-		handleErr(nil, err)
-	}
-	if cbs.PostTree != nil {
-		cbs.PostTree(treeName, btrfs.TREE_LOG_OBJECTID, superblock.Data.LogTree)
-	}
-
-	treeName = "block group tree"
-	if cbs.PreTree != nil {
-		cbs.PreTree(treeName, btrfs.BLOCK_GROUP_TREE_OBJECTID, superblock.Data.BlockGroupRoot)
-	}
-	if err := fs.TreeWalk(superblock.Data.BlockGroupRoot, cbs.TreeWalkHandler); err != nil {
-		handleErr(nil, err)
-	}
-	if cbs.PostTree != nil {
-		cbs.PostTree(treeName, btrfs.BLOCK_GROUP_TREE_OBJECTID, superblock.Data.BlockGroupRoot)
-	}
-
-	for _, tree := range foundTrees {
+	for i := 0; i < len(trees); i++ {
+		tree := trees[i]
 		treeName = tree.Name
 		if cbs.PreTree != nil {
-			cbs.PreTree(treeName, tree.ID, tree.Root)
+			cbs.PreTree(treeName, tree.ID)
 		}
-		if err := fs.TreeWalk(tree.Root, cbs.TreeWalkHandler); err != nil {
+		if err := fs.TreeWalk(tree.ID, cbs.TreeWalkHandler); err != nil {
 			handleErr(nil, err)
 		}
 		if cbs.PostTree != nil {
-			cbs.PostTree(treeName, tree.ID, tree.Root)
+			cbs.PostTree(treeName, tree.ID)
 		}
 	}
 }