rebuildnodes: Have the graph store keys; avoid I/O when indexing a tree

2 files changed, 15 insertions, 8 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index 1d5ba14..f919b37 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -181,20 +181,20 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 			continue
 		}
 		tree.Leafs.Insert(leaf)
-		for j, item := range ts.keyIO.ReadNode(leaf).Data.BodyLeaf {
+		for j, itemKey := range ts.graph.Nodes[leaf].Items {
 			newPtr := keyio.ItemPtr{
 				Node: leaf,
 				Idx:  j,
 			}
-			if oldPtr, exists := tree.Items.Load(item.Key); !exists {
-				tree.Items.Store(item.Key, newPtr)
+			if oldPtr, exists := tree.Items.Load(itemKey); !exists {
+				tree.Items.Store(itemKey, newPtr)
 				numAdded++
 			} else {
 				oldGen := ts.graph.Nodes[oldPtr.Node].Generation
 				newGen := ts.graph.Nodes[newPtr.Node].Generation
 				switch {
 				case oldGen < newGen:
-					tree.Items.Store(item.Key, newPtr)
+					tree.Items.Store(itemKey, newPtr)
 					numReplaced++
 				case oldGen > newGen:
 					// keep the old one
@@ -203,12 +203,12 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 					// handle this is, and so if this happens I want the program to crash
 					// and force me to figure out how to handle it.
 					panic(fmt.Errorf("dup key=%v in tree=%v: old=%v=%#v ; new=%v=%#v",
-						item.Key, treeID,
+						itemKey, treeID,
 						oldPtr, ts.graph.Nodes[oldPtr.Node],
 						newPtr, ts.graph.Nodes[newPtr.Node]))
 				}
 			}
-			ts.cbAddedItem(ctx, treeID, item.Key)
+			ts.cbAddedItem(ctx, treeID, itemKey)
 			progress(i)
 		}
 	}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
index 46fb4b2..1180b0d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
@@ -6,6 +6,7 @@ package graph
 
 import (
 	"fmt"
+	"reflect"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
@@ -22,10 +23,11 @@ type Node struct {
 	NumItems   uint32
 	MinItem    btrfsprim.Key
 	MaxItem    btrfsprim.Key
+	Items      []btrfsprim.Key
 }
 
 func (n Node) String() string {
-	if n == (Node{}) {
+	if reflect.ValueOf(n).IsZero() {
 		return "{}"
 	}
 	return fmt.Sprintf(`{lvl:%v, gen:%v, tree:%v, cnt:%v, min:(%v,%v,%v), max:(%v,%v,%v)}`,
@@ -130,7 +132,7 @@ func New(sb btrfstree.Superblock) *Graph {
 }
 
 func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
-	g.Nodes[nodeRef.Addr] = Node{
+	nodeData := Node{
 		Level:      nodeRef.Data.Head.Level,
 		Generation: nodeRef.Data.Head.Generation,
 		Owner:      nodeRef.Data.Head.Owner,
@@ -147,7 +149,11 @@ func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.N
 			}
 		}
 		kps := make([]Edge, 0, cnt)
+		keys := make([]btrfsprim.Key, len(nodeRef.Data.BodyLeaf))
+		nodeData.Items = keys
+		g.Nodes[nodeRef.Addr] = nodeData
 		for i, item := range nodeRef.Data.BodyLeaf {
+			keys[i] = item.Key
 			if itemBody, ok := item.Body.(btrfsitem.Root); ok {
 				kps = append(kps, Edge{
 					FromRoot:     nodeRef.Addr,
@@ -161,6 +167,7 @@ func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.N
 			}
 		}
 	} else {
+		g.Nodes[nodeRef.Addr] = nodeData
 		kps := make([]Edge, len(nodeRef.Data.BodyInternal))
 		for i, kp := range nodeRef.Data.BodyInternal {
 			kps[i] = Edge{