rebuildnodes/btrees: Track the tree's root item offset

2 files changed, 12 insertions, 10 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index a4ccd5e..3a6c35b 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -25,9 +25,10 @@ import (
 
 type rebuiltTree struct {
 	// static
-	ID     btrfsprim.ObjID
-	UUID   btrfsprim.UUID
-	Parent *rebuiltTree
+	ID        btrfsprim.ObjID
+	UUID      btrfsprim.UUID
+	Parent    *rebuiltTree
+	ParentGen btrfsprim.Generation // offset of this tree's root item
 
 	// all leafs (lvl=0) that pass .isOwnerOK, even if not in the tree
 	leafToRoots map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
@@ -105,7 +106,7 @@ type RebuiltTrees struct {
 
 	// static callbacks
 	cbAddedItem  func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key)
-	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (item btrfsitem.Root, ok bool)
+	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool)
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
 
 	// mutable
@@ -117,7 +118,7 @@ type RebuiltTrees struct {
 func NewRebuiltTrees(
 	sb btrfstree.Superblock, graph pkggraph.Graph, keyIO keyio.Handle,
 	cbAddedItem func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key),
-	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (item btrfsitem.Root, ok bool),
+	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool),
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool),
 ) *RebuiltTrees {
 	return &RebuiltTrees{
@@ -252,13 +253,14 @@ func (ts *RebuiltTrees) addTree(ctx context.Context, treeID btrfsprim.ObjID, sta
 		if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
 			return false
 		}
-		rootItem, ok := ts.cbLookupRoot(ctx, treeID)
+		rootOff, rootItem, ok := ts.cbLookupRoot(ctx, treeID)
 		if !ok {
 			return false
 		}
 		root = rootItem.ByteNr
 		tree.UUID = rootItem.UUID
 		if rootItem.ParentUUID != (btrfsprim.UUID{}) {
+			tree.ParentGen = rootOff
 			if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
 				return false
 			}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 4e60632..180edab 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -184,11 +184,11 @@ func (o *rebuilder) cbAddedItem(ctx context.Context, tree btrfsprim.ObjID, key b
 	})
 }
 
-func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (item btrfsitem.Root, ok bool) {
+func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool) {
 	key, ok := o._want(ctx, btrfsprim.ROOT_TREE_OBJECTID, tree, btrfsitem.ROOT_ITEM_KEY)
 	if !ok {
 		o.queue = append(o.queue, o.curKey)
-		return btrfsitem.Root{}, false
+		return 0, btrfsitem.Root{}, false
 	}
 	itemBody, ok := o.rebuilt.Load(ctx, btrfsprim.ROOT_TREE_OBJECTID, key)
 	if !ok {
@@ -196,10 +196,10 @@ func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (ite
 	}
 	switch itemBody := itemBody.(type) {
 	case btrfsitem.Root:
-		return itemBody, true
+		return btrfsprim.Generation(key.Offset), itemBody, true
 	case btrfsitem.Error:
 		o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", btrfsprim.ROOT_TREE_OBJECTID, key, itemBody.Err))
-		return btrfsitem.Root{}, false
+		return 0, btrfsitem.Root{}, false
 	default:
 		// This is a panic because the item decoder should not emit ROOT_ITEM items as anything but
 		// btrfsitem.Root or btrfsitem.Error without this code also being updated.