rebuildtrees: Move item-data from btrfsutil.KeyIO to scan

5 files changed, 103 insertions, 95 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
index ca1ce8c..c1391eb 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
@@ -47,8 +47,7 @@ func (o keyAndTree) String() string {
 }
 
 type rebuilder struct {
-	sb    btrfstree.Superblock
-	graph btrfsutil.Graph
+	scan  ScanDevicesResult
 	keyIO *btrfsutil.KeyIO
 
 	rebuilt *btrfsutil.RebuiltForrest
@@ -79,17 +78,18 @@ type Rebuilder interface {
 
 func NewRebuilder(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalAddr) (Rebuilder, error) {
 	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.step", "read-fs-data")
-	sb, nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeList) // ScanDevices does its own logging
+	scanData, err := ScanDevices(ctx, fs, nodeList) // ScanDevices does its own logging
 	if err != nil {
 		return nil, err
 	}
 
+	keyIO := btrfsutil.NewKeyIO(fs, scanData.Superblock, scanData.Graph)
+
 	o := &rebuilder{
-		sb:    sb,
-		graph: nodeGraph,
+		scan:  scanData,
 		keyIO: keyIO,
 	}
-	o.rebuilt = btrfsutil.NewRebuiltForrest(sb, nodeGraph, keyIO, o)
+	o.rebuilt = btrfsutil.NewRebuiltForrest(scanData.Superblock, scanData.Graph, keyIO, o)
 	return o, nil
 }
 
@@ -385,8 +385,8 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, treeID btrfsprim.Ob
 		} else {
 			choices[choice] = ChoiceInfo{
 				Count:      1,
-				Distance:   discardOK(o.rebuilt.Tree(ctx, treeID).COWDistance(o.graph.Nodes[choice].Owner)),
-				Generation: o.graph.Nodes[choice].Generation,
+				Distance:   discardOK(o.rebuilt.Tree(ctx, treeID).COWDistance(o.scan.Graph.Nodes[choice].Owner)),
+				Generation: o.scan.Graph.Nodes[choice].Generation,
 			}
 		}
 	}
@@ -399,8 +399,8 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, treeID btrfsprim.Ob
 			} else {
 				choices[choice] = ChoiceInfo{
 					Count:      1,
-					Distance:   discardOK(o.rebuilt.Tree(ctx, treeID).COWDistance(o.graph.Nodes[choice].Owner)),
-					Generation: o.graph.Nodes[choice].Generation,
+					Distance:   discardOK(o.rebuilt.Tree(ctx, treeID).COWDistance(o.scan.Graph.Nodes[choice].Owner)),
+					Generation: o.scan.Graph.Nodes[choice].Generation,
 				}
 			}
 		}
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
index 704f4ee..78ed8ab 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
@@ -145,7 +145,7 @@ func (o *rebuilder) WantDirIndex(ctx context.Context, reason string, treeID btrf
 			return tgt.Compare(key)
 		},
 		func(_ btrfsprim.Key, ptr btrfsutil.ItemPtr) bool {
-			if itemName, ok := o.keyIO.Names[ptr]; ok && bytes.Equal(itemName, name) {
+			if itemName, ok := o.scan.Names[ptr]; ok && bytes.Equal(itemName, name) {
 				found = true
 			}
 			return !found
@@ -166,7 +166,7 @@ func (o *rebuilder) WantDirIndex(ctx context.Context, reason string, treeID btrf
 			return tgt.Compare(key)
 		},
 		func(_ btrfsprim.Key, ptr btrfsutil.ItemPtr) bool {
-			if itemName, ok := o.keyIO.Names[ptr]; ok && bytes.Equal(itemName, name) {
+			if itemName, ok := o.scan.Names[ptr]; ok && bytes.Equal(itemName, name) {
 				wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(ctx, ptr.Node))
 			}
 			return true
@@ -203,7 +203,7 @@ func (o *rebuilder) _walkRange(
 			}
 		},
 		func(runKey btrfsprim.Key, runPtr btrfsutil.ItemPtr) bool {
-			runSizeAndErr, ok := o.keyIO.Sizes[runPtr]
+			runSizeAndErr, ok := o.scan.Sizes[runPtr]
 			if !ok {
 				panic(fmt.Errorf("should not happen: %v (%v) did not have a size recorded",
 					runPtr, keyAndTree{TreeID: treeID, Key: runKey}))
@@ -372,7 +372,7 @@ func (o *rebuilder) WantCSum(ctx context.Context, reason string, inodeTree, inod
 	if !ok {
 		panic(fmt.Errorf("should not happen: could not load key: %v", inodeWant))
 	}
-	inodeFlags, ok := o.keyIO.Flags[inodePtr]
+	inodeFlags, ok := o.scan.Flags[inodePtr]
 	if !ok {
 		panic(fmt.Errorf("should not happen: INODE_ITEM did not have flags recorded"))
 	}
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
index 4bc09ad..961ff8d 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
@@ -6,11 +6,14 @@ package rebuildtrees
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/datawire/dlib/dlog"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"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"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsutil"
@@ -18,11 +21,31 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
-func ScanDevices(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalAddr) (btrfstree.Superblock, btrfsutil.Graph, *btrfsutil.KeyIO, error) {
+type SizeAndErr struct {
+	Size uint64
+	Err  error
+}
+
+type FlagsAndErr struct {
+	NoDataSum bool
+	Err       error
+}
+
+type ScanDevicesResult struct {
+	Superblock btrfstree.Superblock
+
+	Graph btrfsutil.Graph
+
+	Flags map[btrfsutil.ItemPtr]FlagsAndErr // INODE_ITEM
+	Names map[btrfsutil.ItemPtr][]byte      // DIR_INDEX
+	Sizes map[btrfsutil.ItemPtr]SizeAndErr  // EXTENT_CSUM and EXTENT_DATA
+}
+
+func ScanDevices(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalAddr) (ScanDevicesResult, error) {
 	dlog.Info(ctx, "Reading superblock...")
 	sb, err := fs.Superblock()
 	if err != nil {
-		return btrfstree.Superblock{}, btrfsutil.Graph{}, nil, err
+		return ScanDevicesResult{}, err
 	}
 
 	dlog.Infof(ctx, "Reading node data from FS...")
@@ -33,24 +56,30 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalA
 		dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.read.substep", "read-nodes"),
 		dlog.LogLevelInfo, textui.Tunable(1*time.Second))
 
-	nodeGraph := btrfsutil.NewGraph(*sb)
-	keyIO := btrfsutil.NewKeyIO(fs, *sb)
+	ret := ScanDevicesResult{
+		Superblock: *sb,
+
+		Graph: btrfsutil.NewGraph(*sb),
+
+		Flags: make(map[btrfsutil.ItemPtr]FlagsAndErr),
+		Names: make(map[btrfsutil.ItemPtr][]byte),
+		Sizes: make(map[btrfsutil.ItemPtr]SizeAndErr),
+	}
 
 	progressWriter.Set(stats)
 	for _, laddr := range nodeList {
 		if err := ctx.Err(); err != nil {
-			return btrfstree.Superblock{}, btrfsutil.Graph{}, nil, err
+			return ScanDevicesResult{}, err
 		}
 		node, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, laddr, btrfstree.NodeExpectations{
 			LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 		})
 		if err != nil {
 			node.Free()
-			return btrfstree.Superblock{}, btrfsutil.Graph{}, nil, err
+			return ScanDevicesResult{}, err
 		}
 
-		nodeGraph.InsertNode(node)
-		keyIO.InsertNode(node)
+		ret.insertNode(node)
 
 		node.Free()
 
@@ -64,10 +93,54 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalA
 	dlog.Info(ctx, "... done reading node data")
 
 	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.read.substep", "check")
-	if err := nodeGraph.FinalCheck(ctx, fs, *sb); err != nil {
-		return btrfstree.Superblock{}, btrfsutil.Graph{}, nil, err
+	if err := ret.Graph.FinalCheck(ctx, fs, *sb); err != nil {
+		return ScanDevicesResult{}, err
 	}
-	keyIO.SetGraph(*nodeGraph)
 
-	return *sb, *nodeGraph, keyIO, nil
+	return ret, nil
+}
+
+func (o *ScanDevicesResult) insertNode(node *btrfstree.Node) {
+	o.Graph.InsertNode(node)
+	for i, item := range node.BodyLeaf {
+		ptr := btrfsutil.ItemPtr{
+			Node: node.Head.Addr,
+			Slot: i,
+		}
+		switch itemBody := item.Body.(type) {
+		case *btrfsitem.Inode:
+			o.Flags[ptr] = FlagsAndErr{
+				NoDataSum: itemBody.Flags.Has(btrfsitem.INODE_NODATASUM),
+				Err:       nil,
+			}
+		case *btrfsitem.DirEntry:
+			if item.Key.ItemType == btrfsprim.DIR_INDEX_KEY {
+				o.Names[ptr] = append([]byte(nil), itemBody.Name...)
+			}
+		case *btrfsitem.ExtentCSum:
+			o.Sizes[ptr] = SizeAndErr{
+				Size: uint64(itemBody.Size()),
+				Err:  nil,
+			}
+		case *btrfsitem.FileExtent:
+			size, err := itemBody.Size()
+			o.Sizes[ptr] = SizeAndErr{
+				Size: uint64(size),
+				Err:  err,
+			}
+		case *btrfsitem.Error:
+			switch item.Key.ItemType {
+			case btrfsprim.INODE_ITEM_KEY:
+				o.Flags[ptr] = FlagsAndErr{
+					Err: fmt.Errorf("error decoding item: ptr=%v (tree=%v key=%v): %w",
+						ptr, node.Head.Owner, item.Key, itemBody.Err),
+				}
+			case btrfsprim.EXTENT_CSUM_KEY, btrfsprim.EXTENT_DATA_KEY:
+				o.Sizes[ptr] = SizeAndErr{
+					Err: fmt.Errorf("error decoding item: ptr=%v (tree=%v key=%v): %w",
+						ptr, node.Head.Owner, item.Key, itemBody.Err),
+				}
+			}
+		}
+	}
 }
diff --git a/lib/btrfsutil/graph.go b/lib/btrfsutil/graph.go
index a7c299a..45e9878 100644
--- a/lib/btrfsutil/graph.go
+++ b/lib/btrfsutil/graph.go
@@ -131,8 +131,8 @@ func (g Graph) insertTreeRoot(sb btrfstree.Superblock, treeID btrfsprim.ObjID) {
 	})
 }
 
-func NewGraph(sb btrfstree.Superblock) *Graph {
-	g := &Graph{
+func NewGraph(sb btrfstree.Superblock) Graph {
+	g := Graph{
 		Nodes:     make(map[btrfsvol.LogicalAddr]GraphNode),
 		BadNodes:  make(map[btrfsvol.LogicalAddr]error),
 		EdgesFrom: make(map[btrfsvol.LogicalAddr][]*GraphEdge),
diff --git a/lib/btrfsutil/rebuilt_readitem.go b/lib/btrfsutil/rebuilt_readitem.go
index b1a0656..7e9be09 100644
--- a/lib/btrfsutil/rebuilt_readitem.go
+++ b/lib/btrfsutil/rebuilt_readitem.go
@@ -29,37 +29,20 @@ func (ptr ItemPtr) String() string {
 	return fmt.Sprintf("node@%v[%v]", ptr.Node, ptr.Slot)
 }
 
-type SizeAndErr struct {
-	Size uint64
-	Err  error
-}
-
-type FlagsAndErr struct {
-	NoDataSum bool
-	Err       error
-}
-
 type KeyIO struct {
 	rawFile diskio.File[btrfsvol.LogicalAddr]
 	sb      btrfstree.Superblock
 	graph   Graph
 
-	Flags map[ItemPtr]FlagsAndErr // INODE_ITEM
-	Names map[ItemPtr][]byte      // DIR_INDEX
-	Sizes map[ItemPtr]SizeAndErr  // EXTENT_CSUM and EXTENT_DATA
-
 	mu    sync.Mutex
 	cache containers.ARCache[btrfsvol.LogicalAddr, *btrfstree.Node]
 }
 
-func NewKeyIO(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock) *KeyIO {
+func NewKeyIO(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph Graph) *KeyIO {
 	return &KeyIO{
 		rawFile: file,
 		sb:      sb,
-
-		Flags: make(map[ItemPtr]FlagsAndErr),
-		Names: make(map[ItemPtr][]byte),
-		Sizes: make(map[ItemPtr]SizeAndErr),
+		graph:   graph,
 
 		cache: containers.ARCache[btrfsvol.LogicalAddr, *btrfstree.Node]{
 			MaxLen: textui.Tunable(8),
@@ -70,54 +53,6 @@ func NewKeyIO(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock) *
 	}
 }
 
-func (o *KeyIO) InsertNode(node *btrfstree.Node) {
-	for i, item := range node.BodyLeaf {
-		ptr := ItemPtr{
-			Node: node.Head.Addr,
-			Slot: i,
-		}
-		switch itemBody := item.Body.(type) {
-		case *btrfsitem.Inode:
-			o.Flags[ptr] = FlagsAndErr{
-				NoDataSum: itemBody.Flags.Has(btrfsitem.INODE_NODATASUM),
-				Err:       nil,
-			}
-		case *btrfsitem.DirEntry:
-			if item.Key.ItemType == btrfsprim.DIR_INDEX_KEY {
-				o.Names[ptr] = append([]byte(nil), itemBody.Name...)
-			}
-		case *btrfsitem.ExtentCSum:
-			o.Sizes[ptr] = SizeAndErr{
-				Size: uint64(itemBody.Size()),
-				Err:  nil,
-			}
-		case *btrfsitem.FileExtent:
-			size, err := itemBody.Size()
-			o.Sizes[ptr] = SizeAndErr{
-				Size: uint64(size),
-				Err:  err,
-			}
-		case *btrfsitem.Error:
-			switch item.Key.ItemType {
-			case btrfsprim.INODE_ITEM_KEY:
-				o.Flags[ptr] = FlagsAndErr{
-					Err: fmt.Errorf("error decoding item: ptr=%v (tree=%v key=%v): %w",
-						ptr, node.Head.Owner, item.Key, itemBody.Err),
-				}
-			case btrfsprim.EXTENT_CSUM_KEY, btrfsprim.EXTENT_DATA_KEY:
-				o.Sizes[ptr] = SizeAndErr{
-					Err: fmt.Errorf("error decoding item: ptr=%v (tree=%v key=%v): %w",
-						ptr, node.Head.Owner, item.Key, itemBody.Err),
-				}
-			}
-		}
-	}
-}
-
-func (o *KeyIO) SetGraph(graph Graph) {
-	o.graph = graph
-}
-
 func (o *KeyIO) readNode(ctx context.Context, laddr btrfsvol.LogicalAddr) *btrfstree.Node {
 	if cached, ok := o.cache.Load(laddr); ok {
 		dlog.Tracef(ctx, "cache-hit node@%v", laddr)