get initial expected levels from the superblock

6 files changed, 141 insertions, 127 deletions

diff --git a/pkg/btrfs/io2_lv.go b/pkg/btrfs/io2_lv.go
index 0564066..486a4ae 100644
--- a/pkg/btrfs/io2_lv.go
+++ b/pkg/btrfs/io2_lv.go
@@ -153,7 +153,7 @@ func (fs *FS) initDev(sb *util.Ref[btrfsvol.PhysicalAddr, Superblock]) error {
 			}
 		}
 	}
-	if err := fs.TreeWalk(sb.Data.ChunkTree, TreeWalkHandler{
+	if err := fs.TreeWalk(CHUNK_TREE_OBJECTID, TreeWalkHandler{
 		Item: func(_ TreePath, item Item) error {
 			if item.Head.Key.ItemType != btrfsitem.CHUNK_ITEM_KEY {
 				return nil
diff --git a/pkg/btrfs/io3_btree.go b/pkg/btrfs/io3_btree.go
index 1bcb193..be46d21 100644
--- a/pkg/btrfs/io3_btree.go
+++ b/pkg/btrfs/io3_btree.go
@@ -5,11 +5,11 @@ import (
 	"fmt"
 	"io"
 	iofs "io/fs"
-	"math"
 	"strings"
 
 	"github.com/datawire/dlib/derror"
 
+	"lukeshu.com/btrfs-tools/pkg/btrfs/btrfsitem"
 	"lukeshu.com/btrfs-tools/pkg/btrfs/btrfsvol"
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
@@ -74,16 +74,12 @@ type TreePathElem struct {
 	// being pointed at.
 	NodeAddr btrfsvol.LogicalAddr
 	// NodeLevel is the expected or actual level of the node at
-	// NodeAddr, or 255 if there is no knowledge of the level.
+	// NodeAddr.
 	NodeLevel uint8
 }
 
 func (elem TreePathElem) writeNodeTo(w io.Writer) {
-	if elem.NodeLevel != math.MaxUint8 {
-		fmt.Fprintf(w, "node:%d@%v", elem.NodeLevel, elem.NodeAddr)
-	} else {
-		fmt.Fprintf(w, "node@%v", elem.NodeAddr)
-	}
+	fmt.Fprintf(w, "node:%d@%v", elem.NodeLevel, elem.NodeAddr)
 }
 
 func (path TreePath) String() string {
@@ -102,6 +98,75 @@ func (path TreePath) String() string {
 	return ret.String()
 }
 
+// A treeRoot is more-or-less a btrfsitem.Root, but simpler and generalized for
+type treeRoot struct {
+	TreeID     ObjID
+	RootNode   btrfsvol.LogicalAddr
+	Level      uint8
+	Generation Generation
+}
+
+func (fs *FS) lookupTree(treeID ObjID) (*treeRoot, error) {
+	sb, err := fs.Superblock()
+	if err != nil {
+		return nil, err
+	}
+	switch treeID {
+	case ROOT_TREE_OBJECTID:
+		return &treeRoot{
+			TreeID:     treeID,
+			RootNode:   sb.Data.RootTree,
+			Level:      sb.Data.RootLevel,
+			Generation: sb.Data.Generation, // XXX: same generation as LOG_TREE?
+		}, nil
+	case CHUNK_TREE_OBJECTID:
+		return &treeRoot{
+			TreeID:     treeID,
+			RootNode:   sb.Data.ChunkTree,
+			Level:      sb.Data.ChunkLevel,
+			Generation: sb.Data.ChunkRootGeneration,
+		}, nil
+	case TREE_LOG_OBJECTID:
+		return &treeRoot{
+			TreeID:     treeID,
+			RootNode:   sb.Data.LogTree,
+			Level:      sb.Data.LogLevel,
+			Generation: sb.Data.Generation, // XXX: same generation as ROOT_TREE?
+		}, nil
+	case BLOCK_GROUP_TREE_OBJECTID:
+		return &treeRoot{
+			TreeID:     treeID,
+			RootNode:   sb.Data.BlockGroupRoot,
+			Level:      sb.Data.BlockGroupRootLevel,
+			Generation: sb.Data.BlockGroupRootGeneration,
+		}, nil
+	default:
+		rootItem, err := fs.TreeSearch(ROOT_TREE_OBJECTID, func(key Key) int {
+			if key.ObjectID == treeID && key.ItemType == btrfsitem.ROOT_ITEM_KEY {
+				return 0
+			}
+			return Key{
+				ObjectID: treeID,
+				ItemType: btrfsitem.ROOT_ITEM_KEY,
+				Offset:   0,
+			}.Cmp(key)
+		})
+		if err != nil {
+			return nil, err
+		}
+		rootItemBody, ok := rootItem.Body.(btrfsitem.Root)
+		if !ok {
+			return nil, fmt.Errorf("malformed ROOT_ITEM for tree %v", treeID)
+		}
+		return &treeRoot{
+			TreeID:     treeID,
+			RootNode:   rootItemBody.ByteNr,
+			Level:      rootItemBody.Level,
+			Generation: rootItemBody.Generation,
+		}, nil
+	}
+}
+
 type TreeWalkHandler struct {
 	// Callbacks for entire nodes
 	PreNode  func(TreePath) error
@@ -127,12 +192,16 @@ type TreeWalkHandler struct {
 //           else:
 //     004     .Item()
 //     007 .PostNode()
-func (fs *FS) TreeWalk(treeRoot btrfsvol.LogicalAddr, cbs TreeWalkHandler) error {
+func (fs *FS) TreeWalk(treeID ObjID, cbs TreeWalkHandler) error {
+	rootInfo, err := fs.lookupTree(treeID)
+	if err != nil {
+		return err
+	}
 	path := TreePath{
 		TreePathElem{
 			ItemIdx:   -1,
-			NodeAddr:  treeRoot,
-			NodeLevel: math.MaxUint8,
+			NodeAddr:  rootInfo.RootNode,
+			NodeLevel: rootInfo.Level,
 		},
 	}
 	return fs.treeWalk(path, cbs)
@@ -216,12 +285,12 @@ func (fs *FS) treeWalk(path TreePath, cbs TreeWalkHandler) error {
 	return nil
 }
 
-func (fs *FS) treeSearch(treeRoot btrfsvol.LogicalAddr, fn func(Key) int) (TreePath, *util.Ref[btrfsvol.LogicalAddr, Node], error) {
+func (fs *FS) treeSearch(treeRoot treeRoot, fn func(Key) int) (TreePath, *util.Ref[btrfsvol.LogicalAddr, Node], error) {
 	path := TreePath{
 		TreePathElem{
 			ItemIdx:   -1,
-			NodeAddr:  treeRoot,
-			NodeLevel: math.MaxUint8,
+			NodeAddr:  treeRoot.RootNode,
+			NodeLevel: treeRoot.Level,
 		},
 	}
 	for {
@@ -418,16 +487,20 @@ func (fs *FS) next(path TreePath, node *util.Ref[btrfsvol.LogicalAddr, Node]) (T
 	return path, node, nil
 }
 
-func (fs *FS) TreeSearch(treeRoot btrfsvol.LogicalAddr, fn func(Key) int) (Item, error) {
-	path, node, err := fs.treeSearch(treeRoot, fn)
+func (fs *FS) TreeSearch(treeID ObjID, fn func(Key) int) (Item, error) {
+	rootInfo, err := fs.lookupTree(treeID)
+	if err != nil {
+		return Item{}, err
+	}
+	path, node, err := fs.treeSearch(*rootInfo, fn)
 	if err != nil {
 		return Item{}, err
 	}
 	return node.Data.BodyLeaf[path[len(path)-1].ItemIdx], nil
 }
 
-func (fs *FS) TreeLookup(treeRoot btrfsvol.LogicalAddr, key Key) (Item, error) {
-	item, err := fs.TreeSearch(treeRoot, key.Cmp)
+func (fs *FS) TreeLookup(treeID ObjID, key Key) (Item, error) {
+	item, err := fs.TreeSearch(treeID, key.Cmp)
 	if err != nil {
 		err = fmt.Errorf("item with key=%v: %w", key, err)
 	}
@@ -438,8 +511,12 @@ func (fs *FS) TreeLookup(treeRoot btrfsvol.LogicalAddr, key Key) (Item, error) {
 // return *both* a list of items and an error.
 //
 // If no such item is found, an error that is io/fs.ErrNotExist is returned.
-func (fs *FS) TreeSearchAll(treeRoot btrfsvol.LogicalAddr, fn func(Key) int) ([]Item, error) {
-	middlePath, middleNode, err := fs.treeSearch(treeRoot, fn)
+func (fs *FS) TreeSearchAll(treeID ObjID, fn func(Key) int) ([]Item, error) {
+	rootInfo, err := fs.lookupTree(treeID)
+	if err != nil {
+		return nil, err
+	}
+	middlePath, middleNode, err := fs.treeSearch(*rootInfo, fn)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/btrfs/io4_fs.go b/pkg/btrfs/io4_fs.go
index 37d949e..cd8e281 100644
--- a/pkg/btrfs/io4_fs.go
+++ b/pkg/btrfs/io4_fs.go
@@ -66,13 +66,7 @@ type Subvolume struct {
 
 func (sv *Subvolume) init() {
 	sv.rootOnce.Do(func() {
-		sb, err := sv.FS.Superblock()
-		if err != nil {
-			sv.rootErr = err
-			return
-		}
-
-		root, err := sv.FS.TreeLookup(sb.Data.RootTree, Key{
+		root, err := sv.FS.TreeLookup(ROOT_TREE_OBJECTID, Key{
 			ObjectID: sv.TreeID,
 			ItemType: btrfsitem.ROOT_ITEM_KEY,
 			Offset:   0,
@@ -97,22 +91,12 @@ func (sv *Subvolume) GetRootInode() (ObjID, error) {
 	return sv.rootVal.RootDirID, sv.rootErr
 }
 
-func (sv *Subvolume) GetFSTree() (btrfsvol.LogicalAddr, error) {
-	sv.init()
-	return sv.rootVal.ByteNr, sv.rootErr
-}
-
 func (sv *Subvolume) LoadBareInode(inode ObjID) (*BareInode, error) {
 	val := sv.bareInodeCache.GetOrElse(inode, func() (val *BareInode) {
 		val = &BareInode{
 			Inode: inode,
 		}
-		tree, err := sv.GetFSTree()
-		if err != nil {
-			val.Errs = append(val.Errs, err)
-			return
-		}
-		item, err := sv.FS.TreeLookup(tree, Key{
+		item, err := sv.FS.TreeLookup(sv.TreeID, Key{
 			ObjectID: inode,
 			ItemType: btrfsitem.INODE_ITEM_KEY,
 			Offset:   0,
@@ -144,12 +128,7 @@ func (sv *Subvolume) LoadFullInode(inode ObjID) (*FullInode, error) {
 				Inode: inode,
 			},
 		}
-		tree, err := sv.GetFSTree()
-		if err != nil {
-			val.Errs = append(val.Errs, err)
-			return
-		}
-		items, err := sv.FS.TreeSearchAll(tree, func(key Key) int {
+		items, err := sv.FS.TreeSearchAll(sv.TreeID, func(key Key) int {
 			return util.CmpUint(inode, key.ObjectID)
 		})
 		if err != nil {
diff --git a/pkg/btrfs/types_node.go b/pkg/btrfs/types_node.go
index 92f7513..4382a91 100644
--- a/pkg/btrfs/types_node.go
+++ b/pkg/btrfs/types_node.go
@@ -4,7 +4,6 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"math"
 
 	"lukeshu.com/btrfs-tools/pkg/binstruct"
 	"lukeshu.com/btrfs-tools/pkg/btrfs/btrfsitem"
@@ -404,7 +403,7 @@ func (fs *FS) readNodeAtLevel(addr btrfsvol.LogicalAddr, expLevel uint8) (*util.
 	if err != nil {
 		return node, err
 	}
-	if expLevel != math.MaxUint8 && node.Data.Head.Level != expLevel {
+	if node.Data.Head.Level != expLevel {
 		return node, fmt.Errorf("btrfs.FS.ReadNode: node@%v: expected level %v but has level %v",
 			node.Addr, expLevel, node.Data.Head.Level)
 	}
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)
 		}
 	}
 }