1 files changed, 97 insertions, 0 deletions
diff --git a/lib/btrfsutil/walk.go b/lib/btrfsutil/walk.go
new file mode 100644
index 0000000..355976a
--- /dev/null
+++ b/lib/btrfsutil/walk.go
@@ -0,0 +1,97 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfsutil
+
+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"
+)
+
+type WalkError struct {
+	TreeName string
+	Err      *btrfstree.TreeError
+}
+
+func (e *WalkError) Unwrap() error { return e.Err }
+
+func (e *WalkError) Error() string {
+	return fmt.Sprintf("%v: %v", e.TreeName, e.Err)
+}
+
+type WalkAllTreesHandler struct {
+	Err func(*WalkError)
+	// Callbacks for entire trees
+	PreTree  func(name string, id btrfsprim.ObjID)
+	PostTree func(name string, id btrfsprim.ObjID)
+	// Callbacks for nodes or smaller
+	btrfstree.TreeWalkHandler
+}
+
+// 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 WalkError.
+func WalkAllTrees(ctx context.Context, fs btrfstree.TreeOperator, cbs WalkAllTreesHandler) {
+	var treeName string
+
+	trees := []struct {
+		Name string
+		ID   btrfsprim.ObjID
+	}{
+		{
+			Name: "root tree",
+			ID:   btrfsprim.ROOT_TREE_OBJECTID,
+		},
+		{
+			Name: "chunk tree",
+			ID:   btrfsprim.CHUNK_TREE_OBJECTID,
+		},
+		{
+			Name: "log tree",
+			ID:   btrfsprim.TREE_LOG_OBJECTID,
+		},
+		{
+			Name: "block group tree",
+			ID:   btrfsprim.BLOCK_GROUP_TREE_OBJECTID,
+		},
+	}
+	origItem := cbs.Item
+	cbs.Item = func(path btrfstree.TreePath, item btrfstree.Item) error {
+		if item.Key.ItemType == btrfsitem.ROOT_ITEM_KEY {
+			trees = append(trees, struct {
+				Name string
+				ID   btrfsprim.ObjID
+			}{
+				Name: fmt.Sprintf("tree %v (via %v %v)",
+					item.Key.ObjectID.Format(0), treeName, path),
+				ID: item.Key.ObjectID,
+			})
+		}
+		if origItem != nil {
+			return origItem(path, item)
+		}
+		return nil
+	}
+
+	for i := 0; i < len(trees); i++ {
+		tree := trees[i]
+		treeName = tree.Name
+		if cbs.PreTree != nil {
+			cbs.PreTree(treeName, tree.ID)
+		}
+		fs.TreeWalk(
+			ctx,
+			tree.ID,
+			func(err *btrfstree.TreeError) { cbs.Err(&WalkError{TreeName: treeName, Err: err}) },
+			cbs.TreeWalkHandler,
+		)
+		if cbs.PostTree != nil {
+			cbs.PostTree(treeName, tree.ID)
+		}
+	}
+}