From 7c0e009092f593ca0990a50de8db3f81c112e58e Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sun, 18 Dec 2022 00:34:06 -0700
Subject: rebuildnodes: Reduce scope of the orphans set

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go | 7 +++----
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 4 +---
 2 files changed, 4 insertions(+), 7 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
index 51313a9..0ecf852 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
@@ -83,13 +83,12 @@ func (s readStats) String() string {
 }
 
 func indexOrphans(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph graph.Graph) (
-	orphans containers.Set[btrfsvol.LogicalAddr],
 	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr],
 	key2leaf *containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr],
 	err error,
 ) {
 	dlog.Info(ctx, "... counting orphans")
-	orphans = make(containers.Set[btrfsvol.LogicalAddr])
+	orphans := make(containers.Set[btrfsvol.LogicalAddr])
 	for node := range graph.Nodes {
 		if len(graph.EdgesTo[node]) == 0 {
 			orphans.Insert(node)
@@ -145,7 +144,7 @@ func indexOrphans(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb
 			Level: containers.Optional[uint8]{OK: true, Val: 0},
 		})
 		if err != nil {
-			return nil, nil, nil, err
+			return nil, nil, err
 		}
 
 		for _, item := range nodeRef.Data.BodyLeaf {
@@ -162,5 +161,5 @@ func indexOrphans(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb
 	}
 	readProgressWriter.Done()
 
-	return orphans, leaf2orphans, key2leaf, nil
+	return leaf2orphans, key2leaf, nil
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index ef50653..d197bf8 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -37,7 +37,6 @@ type Rebuilder struct {
 	graph        graph.Graph
 	uuidMap      uuidmap.UUIDMap
 	csums        containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
-	orphans      containers.Set[btrfsvol.LogicalAddr]
 	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
 	key2leaf     containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr]
 
@@ -65,7 +64,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Indexing orphans...")
-	orphans, leaf2orphans, key2leaf, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
+	leaf2orphans, key2leaf, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
 	if err != nil {
 		return nil, err
 	}
@@ -79,7 +78,6 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		graph:        *scanData.nodeGraph,
 		uuidMap:      *scanData.uuidMap,
 		csums:        *csums,
-		orphans:      orphans,
 		leaf2orphans: leaf2orphans,
 		key2leaf:     *key2leaf,
 
-- 
cgit v1.2.3-54-g00ecf


From 01a9a607b9412f482e0a9784ab9013a7213077a0 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sun, 18 Dec 2022 00:36:05 -0700
Subject: rebuildnodes: Optimize indexing of orphans

---
 .../btrfsinspect/rebuildnodes/orphans.go           | 83 +++++++---------------
 1 file changed, 24 insertions(+), 59 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
index 0ecf852..5d3f031 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
@@ -33,18 +33,6 @@ func listRoots(graph graph.Graph, leaf btrfsvol.LogicalAddr) containers.Set[btrf
 	return ret
 }
 
-func walk(graph graph.Graph, root btrfsvol.LogicalAddr, fn func(btrfsvol.LogicalAddr) bool) {
-	if _, ok := graph.Nodes[root]; !ok {
-		return
-	}
-	if !fn(root) {
-		return
-	}
-	for _, kp := range graph.EdgesFrom[root] {
-		walk(graph, kp.ToNode, fn)
-	}
-}
-
 type keyAndTree struct {
 	btrfsprim.Key
 	TreeID btrfsprim.ObjID
@@ -58,17 +46,15 @@ func (a keyAndTree) Cmp(b keyAndTree) int {
 }
 
 type crawlStats struct {
-	DoneOrphans  int
-	TotalOrphans int
-
-	VisitedNodes int
-	FoundLeafs   int
+	DoneNodes  int
+	TotalNodes int
+	FoundLeafs int
 }
 
 func (s crawlStats) String() string {
-	return fmt.Sprintf("... crawling orphans %v%% (%v/%v); visited %d nodes, found %d leaf nodes",
-		int(100*float64(s.DoneOrphans)/float64(s.TotalOrphans)),
-		s.DoneOrphans, s.TotalOrphans, s.VisitedNodes, s.FoundLeafs)
+	return fmt.Sprintf("... indexing orphaned leafs %v%% (%v/%v); found %d leaf nodes",
+		int(100*float64(s.DoneNodes)/float64(s.TotalNodes)),
+		s.DoneNodes, s.TotalNodes, s.FoundLeafs)
 }
 
 type readStats struct {
@@ -87,58 +73,38 @@ func indexOrphans(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb
 	key2leaf *containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr],
 	err error,
 ) {
-	dlog.Info(ctx, "... counting orphans")
-	orphans := make(containers.Set[btrfsvol.LogicalAddr])
-	for node := range graph.Nodes {
-		if len(graph.EdgesTo[node]) == 0 {
-			orphans.Insert(node)
-		}
-	}
 
-	leaf2orphans = make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
-	visited := make(containers.Set[btrfsvol.LogicalAddr])
-
-	done := 0
 	crawlProgressWriter := textui.NewProgress[crawlStats](ctx, dlog.LogLevelInfo, 1*time.Second)
-	progress := func() {
+	progress := func(done int) {
 		crawlProgressWriter.Set(crawlStats{
-			DoneOrphans:  done,
-			TotalOrphans: len(orphans),
-			VisitedNodes: len(visited),
-			FoundLeafs:   len(leaf2orphans),
+			DoneNodes:  done,
+			TotalNodes: len(graph.Nodes),
+			FoundLeafs: len(leaf2orphans),
 		})
 	}
-	progress()
-	for _, orphan := range maps.SortedKeys(orphans) {
-		walk(graph, orphan, func(node btrfsvol.LogicalAddr) bool {
-			if visited.Has(node) {
-				return false
-			}
-			visited.Insert(node)
-			if graph.Nodes[node].Level == 0 {
-				if roots := listRoots(graph, node); !roots.Has(0) {
-					leaf2orphans[node] = roots
-				}
-			}
-			progress()
-			return true
-		})
-		done++
-		progress()
+	leaf2orphans = make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+	for i, node := range maps.SortedKeys(graph.Nodes) {
+		progress(i)
+		if graph.Nodes[node].Level != 0 {
+			continue
+		}
+		if roots := listRoots(graph, node); !roots.Has(0) {
+			leaf2orphans[node] = roots
+		}
 	}
+	progress(len(graph.Nodes))
 	crawlProgressWriter.Done()
 
 	key2leaf = new(containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr])
-	done = 0
 	readProgressWriter := textui.NewProgress[readStats](ctx, dlog.LogLevelInfo, 1*time.Second)
-	progress = func() {
+	progress = func(done int) {
 		readProgressWriter.Set(readStats{
 			DoneLeafNodes:  done,
 			TotalLeafNodes: len(leaf2orphans),
 		})
 	}
-	progress()
-	for _, laddr := range maps.SortedKeys(leaf2orphans) {
+	for i, laddr := range maps.SortedKeys(leaf2orphans) {
+		progress(i)
 		nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, sb, laddr, btrfstree.NodeExpectations{
 			LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 			Level: containers.Optional[uint8]{OK: true, Val: 0},
@@ -156,9 +122,8 @@ func indexOrphans(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb
 				key2leaf.Store(k, laddr)
 			}
 		}
-		done++
-		progress()
 	}
+	progress(len(leaf2orphans))
 	readProgressWriter.Done()
 
 	return leaf2orphans, key2leaf, nil
-- 
cgit v1.2.3-54-g00ecf


From aebdf099d1b02394debd71a40b18050370ba7416 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sun, 18 Dec 2022 00:58:08 -0700
Subject: rebuildnodes: Optimize memory access when indexing orphans

---
 .../btrfsinspect/rebuildnodes/graph/graph.go       | 69 +++++++++++++---------
 .../btrfsinspect/rebuildnodes/orphans.go           | 12 ++--
 2 files changed, 49 insertions(+), 32 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
index 7ae3b89..0a24588 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
@@ -51,13 +51,12 @@ type Graph struct {
 	EdgesTo   map[btrfsvol.LogicalAddr][]*Edge
 }
 
-func (g Graph) insertEdge(kp Edge) {
-	ptr := &kp
-	if kp.ToNode == 0 {
+func (g Graph) insertEdge(ptr *Edge) {
+	if ptr.ToNode == 0 {
 		panic("kp.ToNode should not be zero")
 	}
-	g.EdgesFrom[kp.FromNode] = append(g.EdgesFrom[kp.FromNode], ptr)
-	g.EdgesTo[kp.ToNode] = append(g.EdgesTo[kp.ToNode], ptr)
+	g.EdgesFrom[ptr.FromNode] = append(g.EdgesFrom[ptr.FromNode], ptr)
+	g.EdgesTo[ptr.ToNode] = append(g.EdgesTo[ptr.ToNode], ptr)
 }
 
 func (g Graph) insertTreeRoot(sb btrfstree.Superblock, treeID btrfsprim.ObjID) {
@@ -72,7 +71,7 @@ func (g Graph) insertTreeRoot(sb btrfstree.Superblock, treeID btrfsprim.ObjID) {
 	if treeInfo.RootNode == 0 {
 		return
 	}
-	g.insertEdge(Edge{
+	g.insertEdge(&Edge{
 		FromTree:     treeID,
 		ToNode:       treeInfo.RootNode,
 		ToLevel:      treeInfo.Level,
@@ -99,18 +98,6 @@ func New(sb btrfstree.Superblock) *Graph {
 }
 
 func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
-	for _, item := range nodeRef.Data.BodyLeaf {
-		switch itemBody := item.Body.(type) {
-		case btrfsitem.Root:
-			g.insertEdge(Edge{
-				FromTree:     item.Key.ObjectID,
-				ToNode:       itemBody.ByteNr,
-				ToLevel:      itemBody.Level,
-				ToGeneration: itemBody.Generation,
-			})
-		}
-	}
-
 	g.Nodes[nodeRef.Addr] = Node{
 		Level:      nodeRef.Data.Head.Level,
 		Generation: nodeRef.Data.Head.Generation,
@@ -119,16 +106,42 @@ func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.N
 		MinItem:    discardOK(nodeRef.Data.MinItem()),
 		MaxItem:    discardOK(nodeRef.Data.MaxItem()),
 	}
-	for i, kp := range nodeRef.Data.BodyInternal {
-		g.insertEdge(Edge{
-			FromTree:     nodeRef.Data.Head.Owner,
-			FromNode:     nodeRef.Addr,
-			FromItem:     i,
-			ToNode:       kp.BlockPtr,
-			ToLevel:      nodeRef.Data.Head.Level - 1,
-			ToKey:        kp.Key,
-			ToGeneration: kp.Generation,
-		})
+
+	if nodeRef.Data.Head.Level == 0 {
+		cnt := 0
+		for _, item := range nodeRef.Data.BodyLeaf {
+			switch item.Body.(type) {
+			case btrfsitem.Root:
+				cnt++
+			}
+		}
+		kps := make([]Edge, 0, cnt)
+		for _, item := range nodeRef.Data.BodyLeaf {
+			switch itemBody := item.Body.(type) {
+			case btrfsitem.Root:
+				kps = append(kps, Edge{
+					FromTree:     item.Key.ObjectID,
+					ToNode:       itemBody.ByteNr,
+					ToLevel:      itemBody.Level,
+					ToGeneration: itemBody.Generation,
+				})
+				g.insertEdge(&kps[len(kps)-1])
+			}
+		}
+	} else {
+		kps := make([]Edge, len(nodeRef.Data.BodyInternal))
+		for i, kp := range nodeRef.Data.BodyInternal {
+			kps[i] = Edge{
+				FromTree:     nodeRef.Data.Head.Owner,
+				FromNode:     nodeRef.Addr,
+				FromItem:     i,
+				ToNode:       kp.BlockPtr,
+				ToLevel:      nodeRef.Data.Head.Level - 1,
+				ToKey:        kp.Key,
+				ToGeneration: kp.Generation,
+			}
+			g.insertEdge(&kps[i])
+		}
 	}
 }
 
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
index 5d3f031..89e9ad4 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
@@ -22,15 +22,19 @@ import (
 )
 
 func listRoots(graph graph.Graph, leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
+	ret := make(containers.Set[btrfsvol.LogicalAddr])
+	_listRoots(ret, graph, leaf)
+	return ret
+}
+
+func _listRoots(ret containers.Set[btrfsvol.LogicalAddr], graph graph.Graph, leaf btrfsvol.LogicalAddr) {
 	kps := graph.EdgesTo[leaf]
 	if len(kps) == 0 {
-		return containers.NewSet(leaf)
+		ret.Insert(leaf)
 	}
-	ret := make(containers.Set[btrfsvol.LogicalAddr])
 	for _, kp := range kps {
-		ret.InsertFrom(listRoots(graph, kp.FromNode))
+		_listRoots(ret, graph, kp.FromNode)
 	}
-	return ret
 }
 
 type keyAndTree struct {
-- 
cgit v1.2.3-54-g00ecf


From 15f0adbfccdf8e4be9f6f493ef5b1fd6bce1488f Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 20 Dec 2022 03:09:23 -0700
Subject: rebuildnodes/graph: Track root items, improve string representations

---
 .../btrfsinspect/rebuildnodes/graph/graph.go       | 52 +++++++++++++++++-----
 1 file changed, 42 insertions(+), 10 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
index 0a24588..46fb4b2 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
@@ -24,10 +24,25 @@ type Node struct {
 	MaxItem    btrfsprim.Key
 }
 
+func (n Node) String() string {
+	if n == (Node{}) {
+		return "{}"
+	}
+	return fmt.Sprintf(`{lvl:%v, gen:%v, tree:%v, cnt:%v, min:(%v,%v,%v), max:(%v,%v,%v)}`,
+		n.Level, n.Generation, n.Owner, n.NumItems,
+		n.MinItem.ObjectID, n.MinItem.ItemType, n.MinItem.Offset,
+		n.MaxItem.ObjectID, n.MaxItem.ItemType, n.MaxItem.Offset)
+}
+
 type Edge struct {
-	FromTree btrfsprim.ObjID
+	// It is invalid both 'FromRoot' and 'FromNode' to be
+	// non-zero.  If both are zero, then the Edge is from the
+	// superblock.
+	FromRoot btrfsvol.LogicalAddr
 	FromNode btrfsvol.LogicalAddr
-	FromItem int
+	FromItem int // only valid if one of FromRoot or FromNode is non-zero
+
+	FromTree btrfsprim.ObjID
 
 	ToNode       btrfsvol.LogicalAddr
 	ToLevel      uint8
@@ -36,8 +51,19 @@ type Edge struct {
 }
 
 func (kp Edge) String() string {
-	return fmt.Sprintf(`{t:%v,n:%v}[%d]->{n:%v,l:%v,g:%v,k:(%d,%v,%d)}`,
-		kp.FromTree, kp.FromNode, kp.FromItem,
+	var from string
+	switch {
+	case kp.FromRoot != 0:
+		from = fmt.Sprintf("root@%v[%d]:%v",
+			kp.FromRoot, kp.FromItem, kp.FromTree)
+	case kp.FromNode != 0:
+		from = fmt.Sprintf("{node:%v, tree:%v}[%d]",
+			kp.FromNode, kp.FromTree, kp.FromItem)
+	default:
+		from = fmt.Sprintf("superblock:%v", kp.FromTree)
+	}
+	return fmt.Sprintf(`%s -> {n:%v,l:%v,g:%v,k:(%v,%v,%v)}`,
+		from,
 		kp.ToNode, kp.ToLevel, kp.ToGeneration,
 		kp.ToKey.ObjectID,
 		kp.ToKey.ItemType,
@@ -55,6 +81,12 @@ func (g Graph) insertEdge(ptr *Edge) {
 	if ptr.ToNode == 0 {
 		panic("kp.ToNode should not be zero")
 	}
+	if ptr.FromRoot != 0 && ptr.FromNode != 0 {
+		panic("kp.FromRoot and kp.FromNode should not both be set")
+	}
+	if (ptr.FromRoot == 0 && ptr.FromNode == 0) && ptr.FromItem != 0 {
+		panic("kp.FromItem should only be set if either kp.FromRoot or kp.FromItem is set")
+	}
 	g.EdgesFrom[ptr.FromNode] = append(g.EdgesFrom[ptr.FromNode], ptr)
 	g.EdgesTo[ptr.ToNode] = append(g.EdgesTo[ptr.ToNode], ptr)
 }
@@ -110,16 +142,16 @@ func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.N
 	if nodeRef.Data.Head.Level == 0 {
 		cnt := 0
 		for _, item := range nodeRef.Data.BodyLeaf {
-			switch item.Body.(type) {
-			case btrfsitem.Root:
+			if _, ok := item.Body.(btrfsitem.Root); ok {
 				cnt++
 			}
 		}
 		kps := make([]Edge, 0, cnt)
-		for _, item := range nodeRef.Data.BodyLeaf {
-			switch itemBody := item.Body.(type) {
-			case btrfsitem.Root:
+		for i, item := range nodeRef.Data.BodyLeaf {
+			if itemBody, ok := item.Body.(btrfsitem.Root); ok {
 				kps = append(kps, Edge{
+					FromRoot:     nodeRef.Addr,
+					FromItem:     i,
 					FromTree:     item.Key.ObjectID,
 					ToNode:       itemBody.ByteNr,
 					ToLevel:      itemBody.Level,
@@ -132,9 +164,9 @@ func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.N
 		kps := make([]Edge, len(nodeRef.Data.BodyInternal))
 		for i, kp := range nodeRef.Data.BodyInternal {
 			kps[i] = Edge{
-				FromTree:     nodeRef.Data.Head.Owner,
 				FromNode:     nodeRef.Addr,
 				FromItem:     i,
+				FromTree:     nodeRef.Data.Head.Owner,
 				ToNode:       kp.BlockPtr,
 				ToLevel:      nodeRef.Data.Head.Level - 1,
 				ToKey:        kp.Key,
-- 
cgit v1.2.3-54-g00ecf


From abd7485b9fb78db233839effcd06e69d44da3859 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 20 Dec 2022 18:42:24 -0700
Subject: rebuildnodes: Don't export `Rebuilder`

---
 .../btrfsinspect/rebuildnodes/rebuild.go           | 24 +++++++++++-----------
 1 file changed, 12 insertions(+), 12 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index d197bf8..e441abd 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -26,7 +26,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
-type Rebuilder struct {
+type rebuilder struct {
 	raw   *btrfs.FS
 	inner interface {
 		btrfstree.TreeOperator
@@ -70,7 +70,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Rebuilding node tree...")
-	o := &Rebuilder{
+	o := &rebuilder{
 		raw:   fs,
 		inner: btrfsutil.NewBrokenTrees(ctx, fs),
 		sb:    *sb,
@@ -90,13 +90,13 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	return o.augments, nil
 }
 
-func (o *Rebuilder) ioErr(ctx context.Context, err error) {
+func (o *rebuilder) ioErr(ctx context.Context, err error) {
 	err = fmt.Errorf("should not happen: i/o error: %w", err)
 	dlog.Error(ctx, err)
 	panic(err)
 }
 
-func (o *Rebuilder) rebuild(ctx context.Context) error {
+func (o *rebuilder) rebuild(ctx context.Context) error {
 	passNum := 0
 	dlog.Infof(ctx, "... pass %d: scanning for implied items", passNum)
 	o.pendingAugments = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
@@ -151,7 +151,7 @@ func (o *Rebuilder) rebuild(ctx context.Context) error {
 	return nil
 }
 
-func (o *Rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances []map[btrfsvol.LogicalAddr]int) containers.Set[btrfsvol.LogicalAddr] {
+func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances []map[btrfsvol.LogicalAddr]int) containers.Set[btrfsvol.LogicalAddr] {
 	distances := make(map[btrfsvol.LogicalAddr]int)
 	generations := make(map[btrfsvol.LogicalAddr]btrfsprim.Generation)
 	lists := make([]containers.Set[btrfsvol.LogicalAddr], len(listsWithDistances))
@@ -294,7 +294,7 @@ func treeDistance(fs _NodeFile, tree, leaf btrfsprim.ObjID) (int, bool) {
 	}
 }
 
-func (o *Rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, choices containers.Set[btrfsvol.LogicalAddr]) {
+func (o *rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, choices containers.Set[btrfsvol.LogicalAddr]) {
 	choicesWithDist := make(map[btrfsvol.LogicalAddr]int)
 	for choice := range choices {
 		if dist, ok := treeDistance(o.uuidMap, treeID, o.graph.Nodes[choice].Owner); ok {
@@ -312,12 +312,12 @@ func (o *Rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, cho
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 // fsErr implements rebuildCallbacks.
-func (o *Rebuilder) fsErr(ctx context.Context, e error) {
+func (o *rebuilder) fsErr(ctx context.Context, e error) {
 	dlog.Errorf(ctx, "filesystem error: %v", e)
 }
 
 // want implements rebuildCallbacks.
-func (o *Rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) {
+func (o *rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) {
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
@@ -345,7 +345,7 @@ func (o *Rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrf
 }
 
 // wantOff implements rebuildCallbacks.
-func (o *Rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) {
+func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) {
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
@@ -371,7 +371,7 @@ func (o *Rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID b
 }
 
 // wantFunc implements rebuildCallbacks.
-func (o *Rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(btrfsitem.Item) bool) {
+func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(btrfsitem.Item) bool) {
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
@@ -415,7 +415,7 @@ func (o *Rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 // func implements rebuildCallbacks.
 //
 // interval is [beg, end)
-func (o *Rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
+func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
 	for beg < end {
 		// check if we already have it
 		if run, err := btrfs.LookupCSum(o.inner, o.sb.ChecksumType, beg); err == nil {
@@ -460,7 +460,7 @@ func (o *Rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 }
 
 // wantFileExt implements rebuildCallbacks.
-func (o *Rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
+func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
 	min := btrfsprim.Key{
 		ObjectID: ino,
 		ItemType: btrfsitem.EXTENT_DATA_KEY,
-- 
cgit v1.2.3-54-g00ecf


From f7701686ccdb294e3a67f24a3b333342925e89bc Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 20 Dec 2022 21:17:21 -0700
Subject: rebuildnodes: rebuild_graph.go: Have Root items want their UUIDs

---
 lib/btrfs/btrfsitem/item_uuid.go                         | 10 +++++++++-
 .../btrfsinspect/rebuildnodes/rebuild_graph.go           | 16 ++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

(limited to 'lib')

diff --git a/lib/btrfs/btrfsitem/item_uuid.go b/lib/btrfs/btrfsitem/item_uuid.go
index 451ccae..e1a6cf9 100644
--- a/lib/btrfs/btrfsitem/item_uuid.go
+++ b/lib/btrfs/btrfsitem/item_uuid.go
@@ -24,6 +24,14 @@ type UUIDMap struct { // UUID_SUBVOL=251 UUID_RECEIVED_SUBVOL=252
 func KeyToUUID(key btrfsprim.Key) btrfsprim.UUID {
 	var uuid btrfsprim.UUID
 	binary.LittleEndian.PutUint64(uuid[:8], uint64(key.ObjectID))
-	binary.LittleEndian.PutUint64(uuid[8:], uint64(key.Offset))
+	binary.LittleEndian.PutUint64(uuid[8:], key.Offset)
 	return uuid
 }
+
+func UUIDToKey(uuid btrfsprim.UUID) btrfsprim.Key {
+	return btrfsprim.Key{
+		ObjectID: btrfsprim.ObjID(binary.LittleEndian.Uint64(uuid[:8])),
+		ItemType: UUID_SUBVOL_KEY,
+		Offset:   binary.LittleEndian.Uint64(uuid[8:]),
+	}
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_graph.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_graph.go
index 976716d..db7a7a5 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_graph.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_graph.go
@@ -269,6 +269,22 @@ func handleItem(o rebuildCallbacks, ctx context.Context, treeID btrfsprim.ObjID,
 				btrfsitem.INODE_ITEM_KEY,
 				0)
 		}
+		if body.UUID != (btrfsprim.UUID{}) {
+			key := btrfsitem.UUIDToKey(body.UUID)
+			o.wantOff(dlog.WithField(ctx, "wants", "uuid"),
+				btrfsprim.UUID_TREE_OBJECTID,
+				key.ObjectID,
+				key.ItemType,
+				key.Offset)
+		}
+		if body.ParentUUID != (btrfsprim.UUID{}) {
+			key := btrfsitem.UUIDToKey(body.ParentUUID)
+			o.wantOff(dlog.WithField(ctx, "wants", "parent uuid"),
+				btrfsprim.UUID_TREE_OBJECTID,
+				key.ObjectID,
+				key.ItemType,
+				key.Offset)
+		}
 	case btrfsitem.RootRef:
 		var otherType btrfsprim.ItemType
 		var parent, child btrfsprim.ObjID
-- 
cgit v1.2.3-54-g00ecf


From cb93353c360258569e4b91e164136e321707470a Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 13:53:04 -0700
Subject: rebuildnodes: Fix a missing word in a comment

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index e441abd..e0ca7a0 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -590,8 +590,8 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 							continue
 						}
 						itemEnd := itemBeg + itemSize
-						// We're being and "wanting" any extent that has any overlap with the
-						// gap.  But maybe instead we sould only want extents that are
+						// We're being greedy and "wanting" any extent that has any overlap with
+						// the gap.  But maybe instead we sould only want extents that are
 						// *entirely* within the gap.  I'll have to run it on real filesystems
 						// to see what works better.
 						//
-- 
cgit v1.2.3-54-g00ecf


From 3cd6f42caa633b6726b7b4aebeef92c86b2cf213 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 20 Dec 2022 21:35:35 -0700
Subject: btrfsutil: broken_btree.go: Fix an outdated comment

---
 lib/btrfsprogs/btrfsutil/broken_btree.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsutil/broken_btree.go b/lib/btrfsprogs/btrfsutil/broken_btree.go
index e44c9e1..28e8b1d 100644
--- a/lib/btrfsprogs/btrfsutil/broken_btree.go
+++ b/lib/btrfsprogs/btrfsutil/broken_btree.go
@@ -75,7 +75,7 @@ var _ btrfstree.TreeOperator = (*brokenTrees)(nil)
 // NewBrokenTrees wraps a *btrfs.FS to support looking up information
 // from broken trees.
 //
-// Of the btrfs.FS.Tree{Verb} methods:
+// Of the btrfstree.TreeOperator methods:
 //
 //   - TreeWalk works on broken trees
 //   - TreeLookup relies on the tree being properly ordered (which a
-- 
cgit v1.2.3-54-g00ecf


From aa923f4a8bdd679eb2e856780c795f93c4dc17b2 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 20 Dec 2022 21:33:28 -0700
Subject: containers: sortedmap: Add new .Search() and .SearchAll() methods

---
 lib/containers/sortedmap.go | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

(limited to 'lib')

diff --git a/lib/containers/sortedmap.go b/lib/containers/sortedmap.go
index b759fa3..d6d2f9e 100644
--- a/lib/containers/sortedmap.go
+++ b/lib/containers/sortedmap.go
@@ -8,13 +8,13 @@ import (
 	"errors"
 )
 
-type orderedKV[K Ordered[K], V any] struct {
+type OrderedKV[K Ordered[K], V any] struct {
 	K K
 	V V
 }
 
 type SortedMap[K Ordered[K], V any] struct {
-	inner RBTree[K, orderedKV[K, V]]
+	inner RBTree[K, OrderedKV[K, V]]
 }
 
 func (m *SortedMap[K, V]) init() {
@@ -23,7 +23,7 @@ func (m *SortedMap[K, V]) init() {
 	}
 }
 
-func (m *SortedMap[K, V]) keyFn(kv orderedKV[K, V]) K {
+func (m *SortedMap[K, V]) keyFn(kv OrderedKV[K, V]) K {
 	return kv.K
 }
 
@@ -46,7 +46,7 @@ var errStop = errors.New("stop")
 
 func (m *SortedMap[K, V]) Range(f func(key K, value V) bool) {
 	m.init()
-	_ = m.inner.Walk(func(node *RBNode[orderedKV[K, V]]) error {
+	_ = m.inner.Walk(func(node *RBNode[OrderedKV[K, V]]) error {
 		if f(node.Value.K, node.Value.V) {
 			return nil
 		} else {
@@ -57,7 +57,7 @@ func (m *SortedMap[K, V]) Range(f func(key K, value V) bool) {
 
 func (m *SortedMap[K, V]) Subrange(rangeFn func(K, V) int, handleFn func(K, V) bool) {
 	m.init()
-	kvs := m.inner.SearchRange(func(kv orderedKV[K, V]) int {
+	kvs := m.inner.SearchRange(func(kv OrderedKV[K, V]) int {
 		return rangeFn(kv.K, kv.V)
 	})
 	for _, kv := range kvs {
@@ -69,8 +69,26 @@ func (m *SortedMap[K, V]) Subrange(rangeFn func(K, V) int, handleFn func(K, V) b
 
 func (m *SortedMap[K, V]) Store(key K, value V) {
 	m.init()
-	m.inner.Insert(orderedKV[K, V]{
+	m.inner.Insert(OrderedKV[K, V]{
 		K: key,
 		V: value,
 	})
 }
+
+func (m *SortedMap[K, V]) Search(fn func(K, V) int) (K, V, bool) {
+	node := m.inner.Search(func(kv OrderedKV[K, V]) int {
+		return fn(kv.K, kv.V)
+	})
+	if node == nil {
+		var zeroK K
+		var zeroV V
+		return zeroK, zeroV, false
+	}
+	return node.Value.K, node.Value.V, true
+}
+
+func (m *SortedMap[K, V]) SearchAll(fn func(K, V) int) []OrderedKV[K, V] {
+	return m.inner.SearchRange(func(kv OrderedKV[K, V]) int {
+		return fn(kv.K, kv.V)
+	})
+}
-- 
cgit v1.2.3-54-g00ecf


From af795797565e642a9e7fd2a962add8e8b4d5756b Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 01:11:05 -0700
Subject: rebuildnodes/btrees: Implement a new rebuilt-trees system

This will replace its use of btrfsutil.NewBrokenTrees.
---
 .../rebuildnodes/btrees/rebuilt_btrees.go          | 343 +++++++++++++++++++++
 1 file changed, 343 insertions(+)
 create mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
new file mode 100644
index 0000000..28d93b9
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -0,0 +1,343 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrees
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/datawire/dlib/dlog"
+
+	"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/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
+)
+
+type itemPtr struct {
+	Node btrfsvol.LogicalAddr
+	Idx  int
+}
+
+type rebuiltTree struct {
+	// static
+	ID     btrfsprim.ObjID
+	UUID   btrfsprim.UUID
+	Parent *rebuiltTree
+
+	// mutable
+	Roots containers.Set[btrfsvol.LogicalAddr]
+	Items containers.SortedMap[btrfsprim.Key, itemPtr]
+}
+
+// isOwnerOK returns whether it is permissible for a node with
+// .Head.Owner=owner to be in this tree.
+func (t *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID) bool {
+	for {
+		if t == nil {
+			return false
+		}
+		if owner == t.ID {
+			return true
+		}
+		t = t.Parent
+	}
+}
+
+// RebuiltTrees is an abstraction for rebuilding and accessing
+// potentially broken btrees.
+//
+// It is conceptually a btrfstree.TreeOperator, and adds similar
+// broken-tree handling to btrfsutil.BrokenTrees.  However, the API is
+// different thant btrfstree.TreeOperator, and is much more efficient
+// than btrfsutil.BrokenTrees.
+//
+// The efficiency improvements are possible because of the API
+// differences, which are necessary for how it is used in
+// rebuildnodes:
+//
+//   - it consumes an already-read graph.Graph instead of reading the
+//     graph itself
+//
+//   - it does not use `btrfstree.TreePath`
+//
+//   - it does not keep track of errors encountered in a tree
+//
+// A zero RebuiltTrees is invalid; it must be initialized with
+// NewRebuiltTrees().
+type RebuiltTrees struct {
+	// static
+	rawFile diskio.File[btrfsvol.LogicalAddr]
+	sb      btrfstree.Superblock
+	graph   graph.Graph
+
+	// 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)
+	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
+
+	// mutable
+	trees map[btrfsprim.ObjID]*rebuiltTree
+}
+
+// NewRebuiltTrees returns a new RebuiltTrees instance.  All of the
+// callbacks must be non-nil.
+func NewRebuiltTrees(
+	file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph graph.Graph,
+	cbAddedItem func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key),
+	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (item btrfsitem.Root, ok bool),
+	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool),
+) *RebuiltTrees {
+	return &RebuiltTrees{
+		rawFile: file,
+		sb:      sb,
+		graph:   graph,
+
+		cbAddedItem:  cbAddedItem,
+		cbLookupRoot: cbLookupRoot,
+		cbLookupUUID: cbLookupUUID,
+
+		trees: make(map[btrfsprim.ObjID]*rebuiltTree),
+	}
+}
+
+func (ts *RebuiltTrees) readNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node] {
+	graphInfo, ok := ts.graph.Nodes[laddr]
+	if !ok {
+		panic(fmt.Errorf("should not happen: node@%v is not mentioned in the in-memory graph", laddr))
+	}
+	ref, err := btrfstree.ReadNode(ts.rawFile, ts.sb, laddr, btrfstree.NodeExpectations{
+		LAddr:      containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
+		Level:      containers.Optional[uint8]{OK: true, Val: graphInfo.Level},
+		Generation: containers.Optional[btrfsprim.Generation]{OK: true, Val: graphInfo.Generation},
+		Owner: func(treeID btrfsprim.ObjID) error {
+			if treeID != graphInfo.Owner {
+				return fmt.Errorf("expected owner=%v but claims to have owner=%v",
+					graphInfo.Owner, treeID)
+			}
+			return nil
+		},
+		MinItem: containers.Optional[btrfsprim.Key]{OK: true, Val: graphInfo.MinItem},
+		MaxItem: containers.Optional[btrfsprim.Key]{OK: true, Val: graphInfo.MaxItem},
+	})
+	if err != nil {
+		panic(fmt.Errorf("should not happen: i/o error: %w", err))
+	}
+	return ref
+}
+
+func walk(graph graph.Graph, root btrfsvol.LogicalAddr, visited containers.Set[btrfsvol.LogicalAddr], fn func(btrfsvol.LogicalAddr) bool) {
+	if _, ok := graph.Nodes[root]; !ok {
+		return
+	}
+	if visited.Has(root) {
+		return
+	}
+	defer visited.Insert(root)
+	if !fn(root) {
+		return
+	}
+	for _, kp := range graph.EdgesFrom[root] {
+		walk(graph, kp.ToNode, visited, fn)
+	}
+}
+
+type rootStats struct {
+	TreeID   btrfsprim.ObjID
+	RootNode btrfsvol.LogicalAddr
+
+	VisitedNodes int
+	TotalNodes   int
+	AddedItems   int
+}
+
+func (s rootStats) String() string {
+	return fmt.Sprintf("tree %v: adding root node@%v: %v%% (%v/%v) (added %v items)",
+		s.TreeID, s.RootNode,
+		int(100*float64(s.VisitedNodes)/float64(s.TotalNodes)),
+		s.VisitedNodes, s.TotalNodes,
+		s.AddedItems)
+}
+
+// AddRoot adds an additional root node to an existing tree.  It is
+// useful to call .AddRoot() to re-attach part of the tree that has
+// been broken off.
+//
+// It is invalid (panic) to call AddRoot for a tree without having
+// called AddTree first.
+func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, rootNode btrfsvol.LogicalAddr) {
+	tree := ts.trees[treeID]
+	tree.Roots.Insert(rootNode)
+
+	visited := make(containers.Set[btrfsvol.LogicalAddr])
+	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, 1*time.Second)
+	numAdded := 0
+	progress := func() {
+		progressWriter.Set(rootStats{
+			TreeID:       treeID,
+			RootNode:     rootNode,
+			VisitedNodes: len(visited),
+			TotalNodes:   len(ts.graph.Nodes),
+			AddedItems:   numAdded,
+		})
+	}
+	progress()
+	walk(ts.graph, rootNode, visited, func(node btrfsvol.LogicalAddr) bool {
+		progress()
+		if !tree.isOwnerOK(ts.graph.Nodes[node].Owner) {
+			return false
+		}
+		if ts.graph.Nodes[node].Level == 0 {
+			for i, item := range ts.readNode(node).Data.BodyLeaf {
+				if _, exists := tree.Items.Load(item.Key); exists {
+					// This is a panic because I'm not really sure what the best way to
+					// 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", item.Key, treeID))
+				}
+				tree.Items.Store(item.Key, itemPtr{
+					Node: node,
+					Idx:  i,
+				})
+				numAdded++
+				ts.cbAddedItem(ctx, treeID, item.Key)
+			}
+		}
+		return true
+	})
+	progress()
+	progressWriter.Done()
+}
+
+// AddTree initializes the given tree, returning true if it was able
+// to do so, or false if there was a problem and nothing was done.
+// The tree is initialized with the normal root node of the tree.
+//
+// Subsequent calls to AddTree for the same tree are no-ops.
+func (ts *RebuiltTrees) AddTree(ctx context.Context, treeID btrfsprim.ObjID) (ok bool) {
+	return ts.addTree(ctx, treeID, nil)
+}
+
+func (ts *RebuiltTrees) addTree(ctx context.Context, treeID btrfsprim.ObjID, stack []btrfsprim.ObjID) (ok bool) {
+	if _, ok := ts.trees[treeID]; ok {
+		return true
+	}
+	if slices.Contains(treeID, stack) {
+		return false
+	}
+	tree := &rebuiltTree{
+		ID:    treeID,
+		Roots: make(containers.Set[btrfsvol.LogicalAddr]),
+	}
+	var root btrfsvol.LogicalAddr
+	switch treeID {
+	case btrfsprim.ROOT_TREE_OBJECTID:
+		root = ts.sb.RootTree
+	case btrfsprim.CHUNK_TREE_OBJECTID:
+		root = ts.sb.ChunkTree
+	case btrfsprim.TREE_LOG_OBJECTID:
+		root = ts.sb.LogTree
+	case btrfsprim.BLOCK_GROUP_TREE_OBJECTID:
+		root = ts.sb.BlockGroupRoot
+	default:
+		stack := append(stack, treeID)
+		if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
+			return false
+		}
+		rootItem, ok := ts.cbLookupRoot(ctx, treeID)
+		if !ok {
+			return false
+		}
+		root = rootItem.ByteNr
+		tree.UUID = rootItem.UUID
+		if rootItem.ParentUUID != (btrfsprim.UUID{}) {
+			if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
+				return false
+			}
+			parentID, ok := ts.cbLookupUUID(ctx, rootItem.ParentUUID)
+			if !ok {
+				return false
+			}
+			if !ts.addTree(ctx, parentID, append(stack, treeID)) {
+				return false
+			}
+			tree.Parent = ts.trees[parentID]
+		}
+	}
+	ts.trees[treeID] = tree
+	if root != 0 {
+		ts.AddRoot(ctx, treeID, root)
+	}
+	return true
+}
+
+// Load reads an item from a tree.
+//
+// It is not nescessary to call AddTree for that tree first; Load will
+// call it for you.
+func (ts *RebuiltTrees) Load(ctx context.Context, treeID btrfsprim.ObjID, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
+	if !ts.AddTree(ctx, treeID) {
+		return nil, false
+	}
+	ptr, ok := ts.trees[treeID].Items.Load(key)
+	if !ok {
+		return nil, false
+	}
+	return ts.readNode(ptr.Node).Data.BodyLeaf[ptr.Idx].Body, true
+}
+
+// Search searches for an item from a tree.
+//
+// It is not nescessary to call AddTree for that tree first; Search
+// will call it for you.
+func (ts *RebuiltTrees) Search(ctx context.Context, treeID btrfsprim.ObjID, fn func(btrfsprim.Key) int) (key btrfsprim.Key, ok bool) {
+	if !ts.AddTree(ctx, treeID) {
+		return btrfsprim.Key{}, false
+	}
+	k, _, ok := ts.trees[treeID].Items.Search(func(k btrfsprim.Key, _ itemPtr) int {
+		return fn(k)
+	})
+	return k, ok
+}
+
+// Search searches for a range of items from a tree.
+//
+// It is not nescessary to call AddTree for that tree first; SearchAll
+// will call it for you.
+func (ts *RebuiltTrees) SearchAll(ctx context.Context, treeID btrfsprim.ObjID, fn func(btrfsprim.Key) int) []btrfsprim.Key {
+	if !ts.AddTree(ctx, treeID) {
+		return nil
+	}
+	kvs := ts.trees[treeID].Items.SearchAll(func(k btrfsprim.Key, _ itemPtr) int {
+		return fn(k)
+	})
+	if len(kvs) == 0 {
+		return nil
+	}
+	ret := make([]btrfsprim.Key, len(kvs))
+	for i := range kvs {
+		ret[i] = kvs[i].K
+	}
+	return ret
+}
+
+// ListRoots returns a listing of all initialized trees and their root
+// nodes.
+//
+// Do not mutate the set of roots for a tree; it is a pointer to the
+// RebuiltTrees' internal set!
+func (ts *RebuiltTrees) ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] {
+	ret := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], len(ts.trees))
+	for treeID := range ts.trees {
+		ret[treeID] = ts.trees[treeID].Roots
+	}
+	return ret
+}
-- 
cgit v1.2.3-54-g00ecf


From 3f88bf0926046b245844a72a2a0c97654a410f0a Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 02:41:45 -0700
Subject: rebuildnodes/btrees: Cache nodes

---
 .../btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go     | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index 28d93b9..dea7ec3 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -85,7 +85,8 @@ type RebuiltTrees struct {
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
 
 	// mutable
-	trees map[btrfsprim.ObjID]*rebuiltTree
+	trees     map[btrfsprim.ObjID]*rebuiltTree
+	nodeCache *containers.LRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]]
 }
 
 // NewRebuiltTrees returns a new RebuiltTrees instance.  All of the
@@ -105,15 +106,21 @@ func NewRebuiltTrees(
 		cbLookupRoot: cbLookupRoot,
 		cbLookupUUID: cbLookupUUID,
 
-		trees: make(map[btrfsprim.ObjID]*rebuiltTree),
+		trees:     make(map[btrfsprim.ObjID]*rebuiltTree),
+		nodeCache: containers.NewLRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]](8),
 	}
 }
 
 func (ts *RebuiltTrees) readNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node] {
+	if cached, ok := ts.nodeCache.Get(laddr); ok {
+		return cached
+	}
+
 	graphInfo, ok := ts.graph.Nodes[laddr]
 	if !ok {
 		panic(fmt.Errorf("should not happen: node@%v is not mentioned in the in-memory graph", laddr))
 	}
+
 	ref, err := btrfstree.ReadNode(ts.rawFile, ts.sb, laddr, btrfstree.NodeExpectations{
 		LAddr:      containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 		Level:      containers.Optional[uint8]{OK: true, Val: graphInfo.Level},
@@ -131,6 +138,9 @@ func (ts *RebuiltTrees) readNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvo
 	if err != nil {
 		panic(fmt.Errorf("should not happen: i/o error: %w", err))
 	}
+
+	ts.nodeCache.Add(laddr, ref)
+
 	return ref
 }
 
-- 
cgit v1.2.3-54-g00ecf


From 083444569ac882c470da8adee8059989ba176553 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 03:59:49 -0700
Subject: rebuildnodes/btrees: Index all (orphaned) leafs

---
 .../rebuildnodes/btrees/rebuilt_btrees.go          | 195 +++++++++++++++------
 1 file changed, 138 insertions(+), 57 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index dea7ec3..6e68a84 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -15,9 +15,10 @@ import (
 	"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/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	pkggraph "git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
@@ -33,6 +34,9 @@ type rebuiltTree struct {
 	UUID   btrfsprim.UUID
 	Parent *rebuiltTree
 
+	// all leafs (lvl=0) that pass .isOwnerOK, even if not in the tree
+	leafToRoots map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
+
 	// mutable
 	Roots containers.Set[btrfsvol.LogicalAddr]
 	Items containers.SortedMap[btrfsprim.Key, itemPtr]
@@ -40,15 +44,15 @@ type rebuiltTree struct {
 
 // isOwnerOK returns whether it is permissible for a node with
 // .Head.Owner=owner to be in this tree.
-func (t *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID) bool {
+func (tree *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID) bool {
 	for {
-		if t == nil {
+		if tree == nil {
 			return false
 		}
-		if owner == t.ID {
+		if owner == tree.ID {
 			return true
 		}
-		t = t.Parent
+		tree = tree.Parent
 	}
 }
 
@@ -71,13 +75,19 @@ func (t *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID) bool {
 //
 //   - it does not keep track of errors encountered in a tree
 //
+// Additionally, it provides a piece of functionality that
+// btrfsutil.BrokenTrees does not:
+//
+//   - it provides a .LeafToRoots() method to advise on what
+//     additional roots should be added
+//
 // A zero RebuiltTrees is invalid; it must be initialized with
 // NewRebuiltTrees().
 type RebuiltTrees struct {
 	// static
 	rawFile diskio.File[btrfsvol.LogicalAddr]
 	sb      btrfstree.Superblock
-	graph   graph.Graph
+	graph   pkggraph.Graph
 
 	// static callbacks
 	cbAddedItem  func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key)
@@ -92,7 +102,7 @@ type RebuiltTrees struct {
 // NewRebuiltTrees returns a new RebuiltTrees instance.  All of the
 // callbacks must be non-nil.
 func NewRebuiltTrees(
-	file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph graph.Graph,
+	file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph pkggraph.Graph,
 	cbAddedItem func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key),
 	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (item btrfsitem.Root, ok bool),
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool),
@@ -144,36 +154,20 @@ func (ts *RebuiltTrees) readNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvo
 	return ref
 }
 
-func walk(graph graph.Graph, root btrfsvol.LogicalAddr, visited containers.Set[btrfsvol.LogicalAddr], fn func(btrfsvol.LogicalAddr) bool) {
-	if _, ok := graph.Nodes[root]; !ok {
-		return
-	}
-	if visited.Has(root) {
-		return
-	}
-	defer visited.Insert(root)
-	if !fn(root) {
-		return
-	}
-	for _, kp := range graph.EdgesFrom[root] {
-		walk(graph, kp.ToNode, visited, fn)
-	}
-}
-
 type rootStats struct {
 	TreeID   btrfsprim.ObjID
 	RootNode btrfsvol.LogicalAddr
 
-	VisitedNodes int
-	TotalNodes   int
-	AddedItems   int
+	DoneLeafs  int
+	TotalLeafs int
+	AddedItems int
 }
 
 func (s rootStats) String() string {
 	return fmt.Sprintf("tree %v: adding root node@%v: %v%% (%v/%v) (added %v items)",
 		s.TreeID, s.RootNode,
-		int(100*float64(s.VisitedNodes)/float64(s.TotalNodes)),
-		s.VisitedNodes, s.TotalNodes,
+		int(100*float64(s.DoneLeafs)/float64(s.TotalLeafs)),
+		s.DoneLeafs, s.TotalLeafs,
 		s.AddedItems)
 }
 
@@ -187,43 +181,40 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 	tree := ts.trees[treeID]
 	tree.Roots.Insert(rootNode)
 
-	visited := make(containers.Set[btrfsvol.LogicalAddr])
 	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, 1*time.Second)
 	numAdded := 0
-	progress := func() {
+	progress := func(done int) {
 		progressWriter.Set(rootStats{
-			TreeID:       treeID,
-			RootNode:     rootNode,
-			VisitedNodes: len(visited),
-			TotalNodes:   len(ts.graph.Nodes),
-			AddedItems:   numAdded,
+			TreeID:     treeID,
+			RootNode:   rootNode,
+			DoneLeafs:  done,
+			TotalLeafs: len(tree.leafToRoots),
+			AddedItems: numAdded,
 		})
 	}
-	progress()
-	walk(ts.graph, rootNode, visited, func(node btrfsvol.LogicalAddr) bool {
-		progress()
-		if !tree.isOwnerOK(ts.graph.Nodes[node].Owner) {
-			return false
+	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
+		progress(i)
+		roots := tree.leafToRoots[leaf]
+		if !roots.Has(rootNode) {
+			continue
 		}
-		if ts.graph.Nodes[node].Level == 0 {
-			for i, item := range ts.readNode(node).Data.BodyLeaf {
-				if _, exists := tree.Items.Load(item.Key); exists {
-					// This is a panic because I'm not really sure what the best way to
-					// 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", item.Key, treeID))
-				}
-				tree.Items.Store(item.Key, itemPtr{
-					Node: node,
-					Idx:  i,
-				})
-				numAdded++
-				ts.cbAddedItem(ctx, treeID, item.Key)
+		for j, item := range ts.readNode(leaf).Data.BodyLeaf {
+			if _, exists := tree.Items.Load(item.Key); exists {
+				// This is a panic because I'm not really sure what the best way to
+				// 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", item.Key, treeID))
 			}
+			tree.Items.Store(item.Key, itemPtr{
+				Node: leaf,
+				Idx:  j,
+			})
+			numAdded++
+			ts.cbAddedItem(ctx, treeID, item.Key)
+			progress(i)
 		}
-		return true
-	})
-	progress()
+	}
+	progress(len(tree.leafToRoots))
 	progressWriter.Done()
 }
 
@@ -243,6 +234,7 @@ func (ts *RebuiltTrees) addTree(ctx context.Context, treeID btrfsprim.ObjID, sta
 	if slices.Contains(treeID, stack) {
 		return false
 	}
+
 	tree := &rebuiltTree{
 		ID:    treeID,
 		Roots: make(containers.Set[btrfsvol.LogicalAddr]),
@@ -282,13 +274,81 @@ func (ts *RebuiltTrees) addTree(ctx context.Context, treeID btrfsprim.ObjID, sta
 			tree.Parent = ts.trees[parentID]
 		}
 	}
+	tree.indexLeafs(ctx, ts.graph)
+
 	ts.trees[treeID] = tree
 	if root != 0 {
 		ts.AddRoot(ctx, treeID, root)
 	}
+
 	return true
 }
 
+type indexStats struct {
+	TreeID     btrfsprim.ObjID
+	DoneNodes  int
+	TotalNodes int
+}
+
+func (s indexStats) String() string {
+	return fmt.Sprintf("tree %v: indexing leaf nodes: %v%% (%v/%v)",
+		s.TreeID,
+		int(100*float64(s.DoneNodes)/float64(s.TotalNodes)),
+		s.DoneNodes, s.TotalNodes)
+}
+
+func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) {
+	nodeToRoots := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+	progressWriter := textui.NewProgress[indexStats](ctx, dlog.LogLevelInfo, 1*time.Second)
+	progress := func() {
+		progressWriter.Set(indexStats{
+			TreeID:     tree.ID,
+			DoneNodes:  len(nodeToRoots),
+			TotalNodes: len(graph.Nodes),
+		})
+	}
+	progress()
+	for _, node := range maps.SortedKeys(graph.Nodes) {
+		tree.indexNode(graph, node, nodeToRoots, progress, nil)
+	}
+	progressWriter.Done()
+
+	tree.leafToRoots = make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+	for node, roots := range nodeToRoots {
+		if graph.Nodes[node].Level == 0 && len(roots) > 0 {
+			tree.leafToRoots[node] = roots
+		}
+	}
+}
+
+func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) {
+	defer progress()
+	if _, done := index[node]; done {
+		return
+	}
+	if slices.Contains(node, stack) {
+		return
+	}
+	if !tree.isOwnerOK(graph.Nodes[node].Owner) {
+		index[node] = nil
+		return
+	}
+	kps := slices.RemoveAllFunc(graph.EdgesTo[node], func(kp *pkggraph.Edge) bool {
+		return !tree.isOwnerOK(graph.Nodes[kp.FromNode].Owner)
+	})
+	if len(kps) == 0 {
+		index[node] = containers.NewSet[btrfsvol.LogicalAddr](node)
+		return
+	}
+	stack = append(stack, node)
+	roots := make(containers.Set[btrfsvol.LogicalAddr])
+	for _, kp := range kps {
+		tree.indexNode(graph, kp.FromNode, index, progress, stack)
+		roots.InsertFrom(index[kp.FromNode])
+	}
+	index[node] = roots
+}
+
 // Load reads an item from a tree.
 //
 // It is not nescessary to call AddTree for that tree first; Load will
@@ -339,6 +399,27 @@ func (ts *RebuiltTrees) SearchAll(ctx context.Context, treeID btrfsprim.ObjID, f
 	return ret
 }
 
+// LeafToRoots returns the list of potential roots (to pass to
+// .AddRoot) that include a given leaf-node.
+//
+// It is not nescessary to call AddTree for the tree first;
+// LeafToRoots will call it for you.
+func (ts *RebuiltTrees) LeafToRoots(ctx context.Context, treeID btrfsprim.ObjID, leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
+	if !ts.AddTree(ctx, treeID) {
+		return nil
+	}
+	ret := make(containers.Set[btrfsvol.LogicalAddr])
+	for root := range ts.trees[treeID].leafToRoots[leaf] {
+		if !ts.trees[treeID].Roots.Has(root) {
+			ret.Insert(root)
+		}
+	}
+	if len(ret) == 0 {
+		return nil
+	}
+	return ret
+}
+
 // ListRoots returns a listing of all initialized trees and their root
 // nodes.
 //
-- 
cgit v1.2.3-54-g00ecf


From 30be1eacbe4ff253c82c6e6163234dc20b4de550 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 05:30:25 -0700
Subject: rebuildnodes: Refactor existing key-io code in to a sub-package

---
 .../rebuildnodes/btrees/rebuilt_btrees.go          |  72 +++----------
 .../btrfsinspect/rebuildnodes/keyio/keyio.go       | 117 ++++++++++++++++++++
 .../btrfsinspect/rebuildnodes/orphans.go           |  58 +---------
 .../btrfsinspect/rebuildnodes/rebuild.go           | 118 +++++++++------------
 lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go   |   4 +
 5 files changed, 191 insertions(+), 178 deletions(-)
 create mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index 6e68a84..613f3ca 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -16,18 +16,13 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	pkggraph "git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
-type itemPtr struct {
-	Node btrfsvol.LogicalAddr
-	Idx  int
-}
-
 type rebuiltTree struct {
 	// static
 	ID     btrfsprim.ObjID
@@ -39,7 +34,7 @@ type rebuiltTree struct {
 
 	// mutable
 	Roots containers.Set[btrfsvol.LogicalAddr]
-	Items containers.SortedMap[btrfsprim.Key, itemPtr]
+	Items containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
 }
 
 // isOwnerOK returns whether it is permissible for a node with
@@ -85,9 +80,9 @@ func (tree *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID) bool {
 // NewRebuiltTrees().
 type RebuiltTrees struct {
 	// static
-	rawFile diskio.File[btrfsvol.LogicalAddr]
-	sb      btrfstree.Superblock
-	graph   pkggraph.Graph
+	sb    btrfstree.Superblock
+	graph pkggraph.Graph
+	keyIO keyio.Handle
 
 	// static callbacks
 	cbAddedItem  func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key)
@@ -95,63 +90,28 @@ type RebuiltTrees struct {
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
 
 	// mutable
-	trees     map[btrfsprim.ObjID]*rebuiltTree
-	nodeCache *containers.LRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]]
+	trees map[btrfsprim.ObjID]*rebuiltTree
 }
 
 // NewRebuiltTrees returns a new RebuiltTrees instance.  All of the
 // callbacks must be non-nil.
 func NewRebuiltTrees(
-	file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph pkggraph.Graph,
+	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),
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool),
 ) *RebuiltTrees {
 	return &RebuiltTrees{
-		rawFile: file,
-		sb:      sb,
-		graph:   graph,
+		sb:    sb,
+		graph: graph,
+		keyIO: keyIO,
 
 		cbAddedItem:  cbAddedItem,
 		cbLookupRoot: cbLookupRoot,
 		cbLookupUUID: cbLookupUUID,
 
-		trees:     make(map[btrfsprim.ObjID]*rebuiltTree),
-		nodeCache: containers.NewLRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]](8),
-	}
-}
-
-func (ts *RebuiltTrees) readNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node] {
-	if cached, ok := ts.nodeCache.Get(laddr); ok {
-		return cached
-	}
-
-	graphInfo, ok := ts.graph.Nodes[laddr]
-	if !ok {
-		panic(fmt.Errorf("should not happen: node@%v is not mentioned in the in-memory graph", laddr))
+		trees: make(map[btrfsprim.ObjID]*rebuiltTree),
 	}
-
-	ref, err := btrfstree.ReadNode(ts.rawFile, ts.sb, laddr, btrfstree.NodeExpectations{
-		LAddr:      containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
-		Level:      containers.Optional[uint8]{OK: true, Val: graphInfo.Level},
-		Generation: containers.Optional[btrfsprim.Generation]{OK: true, Val: graphInfo.Generation},
-		Owner: func(treeID btrfsprim.ObjID) error {
-			if treeID != graphInfo.Owner {
-				return fmt.Errorf("expected owner=%v but claims to have owner=%v",
-					graphInfo.Owner, treeID)
-			}
-			return nil
-		},
-		MinItem: containers.Optional[btrfsprim.Key]{OK: true, Val: graphInfo.MinItem},
-		MaxItem: containers.Optional[btrfsprim.Key]{OK: true, Val: graphInfo.MaxItem},
-	})
-	if err != nil {
-		panic(fmt.Errorf("should not happen: i/o error: %w", err))
-	}
-
-	ts.nodeCache.Add(laddr, ref)
-
-	return ref
 }
 
 type rootStats struct {
@@ -198,14 +158,14 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 		if !roots.Has(rootNode) {
 			continue
 		}
-		for j, item := range ts.readNode(leaf).Data.BodyLeaf {
+		for j, item := range ts.keyIO.ReadNode(leaf).Data.BodyLeaf {
 			if _, exists := tree.Items.Load(item.Key); exists {
 				// This is a panic because I'm not really sure what the best way to
 				// 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", item.Key, treeID))
 			}
-			tree.Items.Store(item.Key, itemPtr{
+			tree.Items.Store(item.Key, keyio.ItemPtr{
 				Node: leaf,
 				Idx:  j,
 			})
@@ -361,7 +321,7 @@ func (ts *RebuiltTrees) Load(ctx context.Context, treeID btrfsprim.ObjID, key bt
 	if !ok {
 		return nil, false
 	}
-	return ts.readNode(ptr.Node).Data.BodyLeaf[ptr.Idx].Body, true
+	return ts.keyIO.ReadItem(ptr)
 }
 
 // Search searches for an item from a tree.
@@ -372,7 +332,7 @@ func (ts *RebuiltTrees) Search(ctx context.Context, treeID btrfsprim.ObjID, fn f
 	if !ts.AddTree(ctx, treeID) {
 		return btrfsprim.Key{}, false
 	}
-	k, _, ok := ts.trees[treeID].Items.Search(func(k btrfsprim.Key, _ itemPtr) int {
+	k, _, ok := ts.trees[treeID].Items.Search(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
 		return fn(k)
 	})
 	return k, ok
@@ -386,7 +346,7 @@ func (ts *RebuiltTrees) SearchAll(ctx context.Context, treeID btrfsprim.ObjID, f
 	if !ts.AddTree(ctx, treeID) {
 		return nil
 	}
-	kvs := ts.trees[treeID].Items.SearchAll(func(k btrfsprim.Key, _ itemPtr) int {
+	kvs := ts.trees[treeID].Items.SearchAll(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
 		return fn(k)
 	})
 	if len(kvs) == 0 {
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
new file mode 100644
index 0000000..875b616
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
@@ -0,0 +1,117 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package keyio
+
+import (
+	"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"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
+)
+
+type KeyAndTree struct {
+	btrfsprim.Key
+	TreeID btrfsprim.ObjID
+}
+
+func (a KeyAndTree) Cmp(b KeyAndTree) int {
+	if d := a.Key.Cmp(b.Key); d != 0 {
+		return d
+	}
+	return containers.NativeCmp(a.TreeID, b.TreeID)
+}
+
+type ItemPtr struct {
+	Node btrfsvol.LogicalAddr
+	Idx  int
+}
+
+func (ptr ItemPtr) String() string {
+	return fmt.Sprintf("node@%v[%v]", ptr.Node, ptr.Idx)
+}
+
+type Handle struct {
+	Keys containers.SortedMap[KeyAndTree, ItemPtr]
+
+	rawFile diskio.File[btrfsvol.LogicalAddr]
+	sb      btrfstree.Superblock
+	graph   *graph.Graph
+
+	cache *containers.LRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]]
+}
+
+func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph *graph.Graph) *Handle {
+	return &Handle{
+		rawFile: file,
+		sb:      sb,
+		graph:   graph,
+
+		cache: containers.NewLRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]](8),
+	}
+}
+
+func (o *Handle) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
+	for i, item := range nodeRef.Data.BodyLeaf {
+		k := KeyAndTree{
+			Key:    item.Key,
+			TreeID: nodeRef.Data.Head.Owner,
+		}
+		if cur, ok := o.Keys.Load(k); !ok || o.graph.Nodes[cur.Node].Generation < nodeRef.Data.Head.Generation {
+			o.Keys.Store(k, ItemPtr{
+				Node: nodeRef.Addr,
+				Idx:  i,
+			})
+		}
+	}
+}
+
+func (o *Handle) ReadNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node] {
+	if cached, ok := o.cache.Get(laddr); ok {
+		return cached
+	}
+
+	graphInfo, ok := o.graph.Nodes[laddr]
+	if !ok {
+		panic(fmt.Errorf("should not happen: node@%v is not mentioned in the in-memory graph", laddr))
+	}
+
+	ref, err := btrfstree.ReadNode(o.rawFile, o.sb, laddr, btrfstree.NodeExpectations{
+		LAddr:      containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
+		Level:      containers.Optional[uint8]{OK: true, Val: graphInfo.Level},
+		Generation: containers.Optional[btrfsprim.Generation]{OK: true, Val: graphInfo.Generation},
+		Owner: func(treeID btrfsprim.ObjID) error {
+			if treeID != graphInfo.Owner {
+				return fmt.Errorf("expected owner=%v but claims to have owner=%v",
+					graphInfo.Owner, treeID)
+			}
+			return nil
+		},
+		MinItem: containers.Optional[btrfsprim.Key]{OK: true, Val: graphInfo.MinItem},
+		MaxItem: containers.Optional[btrfsprim.Key]{OK: true, Val: graphInfo.MaxItem},
+	})
+	if err != nil {
+		panic(fmt.Errorf("should not happen: i/o error: %w", err))
+	}
+
+	o.cache.Add(laddr, ref)
+
+	return ref
+}
+
+func (o *Handle) ReadItem(ptr ItemPtr) (item btrfsitem.Item, ok bool) {
+	if o.graph.Nodes[ptr.Node].Level != 0 || ptr.Idx < 0 {
+		return nil, false
+	}
+	items := o.ReadNode(ptr.Node).Data.BodyLeaf
+	if ptr.Idx >= len(items) {
+		return nil, false
+	}
+	return items[ptr.Idx].Body, true
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
index 89e9ad4..517070d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/datawire/dlib/dlog"
 
-	"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/btrfsprogs/btrfsinspect/rebuildnodes/graph"
@@ -37,18 +36,6 @@ func _listRoots(ret containers.Set[btrfsvol.LogicalAddr], graph graph.Graph, lea
 	}
 }
 
-type keyAndTree struct {
-	btrfsprim.Key
-	TreeID btrfsprim.ObjID
-}
-
-func (a keyAndTree) Cmp(b keyAndTree) int {
-	if d := a.Key.Cmp(b.Key); d != 0 {
-		return d
-	}
-	return containers.NativeCmp(a.TreeID, b.TreeID)
-}
-
 type crawlStats struct {
 	DoneNodes  int
 	TotalNodes int
@@ -61,20 +48,8 @@ func (s crawlStats) String() string {
 		s.DoneNodes, s.TotalNodes, s.FoundLeafs)
 }
 
-type readStats struct {
-	DoneLeafNodes  int
-	TotalLeafNodes int
-}
-
-func (s readStats) String() string {
-	return fmt.Sprintf("... reading leafs %v%% (%v/%v)",
-		int(100*float64(s.DoneLeafNodes)/float64(s.TotalLeafNodes)),
-		s.DoneLeafNodes, s.TotalLeafNodes)
-}
-
 func indexOrphans(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph graph.Graph) (
 	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr],
-	key2leaf *containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr],
 	err error,
 ) {
 
@@ -99,36 +74,5 @@ func indexOrphans(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb
 	progress(len(graph.Nodes))
 	crawlProgressWriter.Done()
 
-	key2leaf = new(containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr])
-	readProgressWriter := textui.NewProgress[readStats](ctx, dlog.LogLevelInfo, 1*time.Second)
-	progress = func(done int) {
-		readProgressWriter.Set(readStats{
-			DoneLeafNodes:  done,
-			TotalLeafNodes: len(leaf2orphans),
-		})
-	}
-	for i, laddr := range maps.SortedKeys(leaf2orphans) {
-		progress(i)
-		nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, sb, laddr, btrfstree.NodeExpectations{
-			LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
-			Level: containers.Optional[uint8]{OK: true, Val: 0},
-		})
-		if err != nil {
-			return nil, nil, err
-		}
-
-		for _, item := range nodeRef.Data.BodyLeaf {
-			k := keyAndTree{
-				Key:    item.Key,
-				TreeID: nodeRef.Data.Head.Owner,
-			}
-			if cur, ok := key2leaf.Load(k); !ok || graph.Nodes[cur].Generation < nodeRef.Data.Head.Generation {
-				key2leaf.Store(k, laddr)
-			}
-		}
-	}
-	progress(len(leaf2orphans))
-	readProgressWriter.Done()
-
-	return leaf2orphans, key2leaf, nil
+	return leaf2orphans, nil
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index e0ca7a0..254cfee 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -20,6 +20,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildmappings"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
@@ -38,7 +39,7 @@ type rebuilder struct {
 	uuidMap      uuidmap.UUIDMap
 	csums        containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
 	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
-	key2leaf     containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr]
+	keyIO        keyio.Handle
 
 	augments map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]
 
@@ -64,7 +65,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Indexing orphans...")
-	leaf2orphans, key2leaf, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
+	leaf2orphans, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
 	if err != nil {
 		return nil, err
 	}
@@ -79,7 +80,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		uuidMap:      *scanData.uuidMap,
 		csums:        *csums,
 		leaf2orphans: leaf2orphans,
-		key2leaf:     *key2leaf,
+		keyIO:        *scanData.keyIO,
 
 		augments: make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]),
 	}
@@ -335,10 +336,10 @@ func (o *rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrf
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v %v ?}", treeID, objID, typ))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.key2leaf.Subrange(
-		func(k keyAndTree, _ btrfsvol.LogicalAddr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
-		func(_ keyAndTree, v btrfsvol.LogicalAddr) bool {
-			wants.InsertFrom(o.leaf2orphans[v])
+	o.keyIO.Keys.Subrange(
+		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
+		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
+			wants.InsertFrom(o.leaf2orphans[v.Node])
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
@@ -361,10 +362,10 @@ func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID b
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key=%v", treeID, tgt))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.key2leaf.Subrange(
-		func(k keyAndTree, _ btrfsvol.LogicalAddr) int { return tgt.Cmp(k.Key) },
-		func(_ keyAndTree, v btrfsvol.LogicalAddr) bool {
-			wants.InsertFrom(o.leaf2orphans[v])
+	o.keyIO.Keys.Subrange(
+		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { return tgt.Cmp(k.Key) },
+		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
+			wants.InsertFrom(o.leaf2orphans[v.Node])
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
@@ -392,20 +393,15 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key=%v +func", treeID, tgt))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.key2leaf.Subrange(
-		func(k keyAndTree, _ btrfsvol.LogicalAddr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
-		func(k keyAndTree, v btrfsvol.LogicalAddr) bool {
-			nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](o.raw, o.sb, v, btrfstree.NodeExpectations{
-				LAddr:      containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: v},
-				Generation: containers.Optional[btrfsprim.Generation]{OK: true, Val: o.graph.Nodes[v].Generation},
-			})
-			if err != nil {
-				o.ioErr(ctx, err)
+	o.keyIO.Keys.Subrange(
+		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
+		func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
+			itemBody, ok := o.keyIO.ReadItem(v)
+			if !ok {
+				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
 			}
-			for _, item := range nodeRef.Data.BodyLeaf {
-				if k.Key == item.Key && fn(item.Body) {
-					wants.InsertFrom(o.leaf2orphans[v])
-				}
+			if fn(itemBody) {
+				wants.InsertFrom(o.leaf2orphans[v.Node])
 			}
 			return true
 		})
@@ -441,18 +437,18 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 				continue
 			}
 			run := rbNode.Value.Sums
-			key := keyAndTree{
+			key := keyio.KeyAndTree{
 				Key:    rbNode.Value.Key,
 				TreeID: btrfsprim.CSUM_TREE_OBJECTID,
 			}
-			leaf, ok := o.key2leaf.Load(key)
+			itemPtr, ok := o.keyIO.Keys.Load(key)
 			if !ok {
 				// This is a panic because if we found it in `o.csums` then it has
 				// to be in some Node, and if we didn't find it from
 				// btrfs.LookupCSum(), then that Node must be an orphan.
 				panic(fmt.Errorf("should not happen: no orphan contains %v", key.Key))
 			}
-			o.wantAugment(ctx, key.TreeID, o.leaf2orphans[leaf])
+			o.wantAugment(ctx, key.TreeID, o.leaf2orphans[itemPtr.Node])
 
 			beg = run.Addr.Add(run.Size())
 		}
@@ -558,8 +554,8 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 		}
 		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("file extent for tree=%v inode=%v bytes [%v, %v)", treeID, ino, gap.Beg, gap.End))
 		wants := make(containers.Set[btrfsvol.LogicalAddr])
-		o.key2leaf.Subrange(
-			func(k keyAndTree, _ btrfsvol.LogicalAddr) int {
+		o.keyIO.Keys.Subrange(
+			func(k keyio.KeyAndTree, _ keyio.ItemPtr) int {
 				switch {
 				case min.Cmp(k.Key) < 0:
 					return 1
@@ -569,45 +565,37 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 					return 0
 				}
 			},
-			func(k keyAndTree, v btrfsvol.LogicalAddr) bool {
-				nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](o.raw, o.sb, v, btrfstree.NodeExpectations{
-					LAddr:      containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: v},
-					Generation: containers.Optional[btrfsprim.Generation]{OK: true, Val: o.graph.Nodes[v].Generation},
-				})
-				if err != nil {
-					o.ioErr(ctx, fmt.Errorf("error reading previously read node@%v: %w", v, err))
+			func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
+				itemBody, ok := o.keyIO.ReadItem(v)
+				if !ok {
+					o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
 				}
-				for _, item := range nodeRef.Data.BodyLeaf {
-					if k.Key != item.Key {
-						continue
+				switch itemBody := itemBody.(type) {
+				case btrfsitem.FileExtent:
+					itemBeg := int64(k.Offset)
+					itemSize, err := itemBody.Size()
+					if err != nil {
+						o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, k, err))
+						break
 					}
-					switch itemBody := item.Body.(type) {
-					case btrfsitem.FileExtent:
-						itemBeg := int64(item.Key.Offset)
-						itemSize, err := itemBody.Size()
-						if err != nil {
-							o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, item.Key, err))
-							continue
-						}
-						itemEnd := itemBeg + itemSize
-						// We're being greedy and "wanting" any extent that has any overlap with
-						// the gap.  But maybe instead we sould only want extents that are
-						// *entirely* within the gap.  I'll have to run it on real filesystems
-						// to see what works better.
-						//
-						// TODO(lukeshu): Re-evaluate whether being greedy here is the right
-						// thing.
-						if itemEnd > gap.Beg && itemBeg < gap.End {
-							wants.InsertFrom(o.leaf2orphans[v])
-						}
-					case btrfsitem.Error:
-						o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, item.Key, itemBody.Err))
-					default:
-						// This is a panic because the item decoder should not emit EXTENT_DATA
-						// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
-						// this code also being updated.
-						panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", itemBody))
+					itemEnd := itemBeg + itemSize
+					// We're being greedy and "wanting" any extent that has any overlap with
+					// the gap.  But maybe instead we sould only want extents that are
+					// *entirely* within the gap.  I'll have to run it on real filesystems
+					// to see what works better.
+					//
+					// TODO(lukeshu): Re-evaluate whether being greedy here is the right
+					// thing.
+					if itemEnd > gap.Beg && itemBeg < gap.End {
+						wants.InsertFrom(o.leaf2orphans[v.Node])
 					}
+				case btrfsitem.Error:
+					o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, k, itemBody.Err))
+				default:
+					// This is a panic because the item decoder should not emit EXTENT_DATA
+					// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
+					// this code also being updated.
+					panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", itemBody))
 				}
 				return true
 			})
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index 3575534..7e01693 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -16,6 +16,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
@@ -25,6 +26,7 @@ import (
 type scanResult struct {
 	uuidMap   *uuidmap.UUIDMap
 	nodeGraph *graph.Graph
+	keyIO     *keyio.Handle
 }
 
 type scanStats struct {
@@ -56,6 +58,7 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 		uuidMap:   uuidmap.New(),
 		nodeGraph: graph.New(*sb),
 	}
+	ret.keyIO = keyio.NewHandle(fs, *sb, ret.nodeGraph)
 
 	progress(done, total)
 	for _, devResults := range scanResults {
@@ -72,6 +75,7 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 			}
 
 			ret.nodeGraph.InsertNode(nodeRef)
+			ret.keyIO.InsertNode(nodeRef)
 
 			done++
 			progress(done, total)
-- 
cgit v1.2.3-54-g00ecf


From 7620e1948a4d1e17458d8cfc9ed306bb774a3274 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 04:30:56 -0700
Subject: rebuildnodes: Migrate to the new rebuilt-btrees system

---
 .../rebuildnodes/btrees/rebuilt_btrees.go          |  29 +-
 .../btrfsinspect/rebuildnodes/orphans.go           |  78 ----
 .../btrfsinspect/rebuildnodes/rebuild.go           | 396 +++++++++++++--------
 lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go   |  14 -
 .../btrfsinspect/rebuildnodes/uuidmap/scan.go      |  73 ----
 .../btrfsinspect/rebuildnodes/uuidmap/uuidmap.go   |  55 ---
 6 files changed, 273 insertions(+), 372 deletions(-)
 delete mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
 delete mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/scan.go
 delete mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/uuidmap.go

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index 613f3ca..23e974e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -51,6 +51,21 @@ func (tree *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID) bool {
 	}
 }
 
+// cowDistance returns how many COW-snapshots down the 'tree' is from
+// the 'parent'.
+func (tree *rebuiltTree) cowDistance(parentID btrfsprim.ObjID) (dist int, ok bool) {
+	for {
+		if parentID == tree.ID {
+			return dist, true
+		}
+		if tree.Parent == nil {
+			return 0, false
+		}
+		tree = tree.Parent
+		dist++
+	}
+}
+
 // RebuiltTrees is an abstraction for rebuilding and accessing
 // potentially broken btrees.
 //
@@ -70,12 +85,15 @@ func (tree *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID) bool {
 //
 //   - it does not keep track of errors encountered in a tree
 //
-// Additionally, it provides a piece of functionality that
+// Additionally, it provides some functionality that
 // btrfsutil.BrokenTrees does not:
 //
 //   - it provides a .LeafToRoots() method to advise on what
 //     additional roots should be added
 //
+//   - it provides a .COWDistance() method to compare how related two
+//     trees are
+//
 // A zero RebuiltTrees is invalid; it must be initialized with
 // NewRebuiltTrees().
 type RebuiltTrees struct {
@@ -380,6 +398,15 @@ func (ts *RebuiltTrees) LeafToRoots(ctx context.Context, treeID btrfsprim.ObjID,
 	return ret
 }
 
+// COWDistance returns how many COW-snapshots down from the 'child'
+// tree is from the 'parent' tree.
+//
+// It is invalid (panic) to call COWDistance for a tree without having
+// called AddTree for the child first.
+func (ts *RebuiltTrees) COWDistance(ctx context.Context, childID, parentID btrfsprim.ObjID) (dist int, ok bool) {
+	return ts.trees[childID].cowDistance(parentID)
+}
+
 // ListRoots returns a listing of all initialized trees and their root
 // nodes.
 //
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
deleted file mode 100644
index 517070d..0000000
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
+++ /dev/null
@@ -1,78 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package rebuildnodes
-
-import (
-	"context"
-	"fmt"
-	"time"
-
-	"github.com/datawire/dlib/dlog"
-
-	"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/btrfsprogs/btrfsinspect/rebuildnodes/graph"
-	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
-	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
-	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
-)
-
-func listRoots(graph graph.Graph, leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
-	ret := make(containers.Set[btrfsvol.LogicalAddr])
-	_listRoots(ret, graph, leaf)
-	return ret
-}
-
-func _listRoots(ret containers.Set[btrfsvol.LogicalAddr], graph graph.Graph, leaf btrfsvol.LogicalAddr) {
-	kps := graph.EdgesTo[leaf]
-	if len(kps) == 0 {
-		ret.Insert(leaf)
-	}
-	for _, kp := range kps {
-		_listRoots(ret, graph, kp.FromNode)
-	}
-}
-
-type crawlStats struct {
-	DoneNodes  int
-	TotalNodes int
-	FoundLeafs int
-}
-
-func (s crawlStats) String() string {
-	return fmt.Sprintf("... indexing orphaned leafs %v%% (%v/%v); found %d leaf nodes",
-		int(100*float64(s.DoneNodes)/float64(s.TotalNodes)),
-		s.DoneNodes, s.TotalNodes, s.FoundLeafs)
-}
-
-func indexOrphans(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph graph.Graph) (
-	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr],
-	err error,
-) {
-
-	crawlProgressWriter := textui.NewProgress[crawlStats](ctx, dlog.LogLevelInfo, 1*time.Second)
-	progress := func(done int) {
-		crawlProgressWriter.Set(crawlStats{
-			DoneNodes:  done,
-			TotalNodes: len(graph.Nodes),
-			FoundLeafs: len(leaf2orphans),
-		})
-	}
-	leaf2orphans = make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
-	for i, node := range maps.SortedKeys(graph.Nodes) {
-		progress(i)
-		if graph.Nodes[node].Level != 0 {
-			continue
-		}
-		if roots := listRoots(graph, node); !roots.Has(0) {
-			leaf2orphans[node] = roots
-		}
-	}
-	progress(len(graph.Nodes))
-	crawlProgressWriter.Done()
-
-	return leaf2orphans, nil
-}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 254cfee..2ddce88 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 	"sort"
+	"time"
 
 	"github.com/datawire/dlib/dlog"
 
@@ -19,30 +20,24 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildmappings"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
 type rebuilder struct {
-	raw   *btrfs.FS
-	inner interface {
-		btrfstree.TreeOperator
-		Augment(treeID btrfsprim.ObjID, nodeAddr btrfsvol.LogicalAddr) ([]btrfsprim.Key, error)
-	}
-	sb btrfstree.Superblock
-
-	graph        graph.Graph
-	uuidMap      uuidmap.UUIDMap
-	csums        containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
-	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
-	keyIO        keyio.Handle
+	sb      btrfstree.Superblock
+	rebuilt *btrees.RebuiltTrees
 
-	augments map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]
+	graph graph.Graph
+	csums containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
+	keyIO keyio.Handle
 
+	curKey          keyio.KeyAndTree
+	queue           []keyio.KeyAndTree
 	pendingAugments map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int
 }
 
@@ -64,31 +59,21 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		csums = new(containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum])
 	}
 
-	dlog.Info(ctx, "Indexing orphans...")
-	leaf2orphans, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
-	if err != nil {
-		return nil, err
-	}
-
 	dlog.Info(ctx, "Rebuilding node tree...")
 	o := &rebuilder{
-		raw:   fs,
-		inner: btrfsutil.NewBrokenTrees(ctx, fs),
-		sb:    *sb,
-
-		graph:        *scanData.nodeGraph,
-		uuidMap:      *scanData.uuidMap,
-		csums:        *csums,
-		leaf2orphans: leaf2orphans,
-		keyIO:        *scanData.keyIO,
+		sb: *sb,
 
-		augments: make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]),
+		graph: *scanData.nodeGraph,
+		csums: *csums,
+		keyIO: *scanData.keyIO,
 	}
+	o.rebuilt = btrees.NewRebuiltTrees(*sb, *scanData.nodeGraph, *scanData.keyIO,
+		o.cbAddedItem, o.cbLookupRoot, o.cbLookupUUID)
 	if err := o.rebuild(ctx); err != nil {
 		return nil, err
 	}
 
-	return o.augments, nil
+	return o.rebuilt.ListRoots(), nil
 }
 
 func (o *rebuilder) ioErr(ctx context.Context, err error) {
@@ -97,61 +82,153 @@ func (o *rebuilder) ioErr(ctx context.Context, err error) {
 	panic(err)
 }
 
+type rebuildStats struct {
+	PassNum int
+	Task    string
+	N, D    int
+}
+
+func (s rebuildStats) String() string {
+	pct := 100
+	if s.D > 0 {
+		pct = int(100 * float64(s.N) / float64(s.D))
+	}
+	return fmt.Sprintf("... pass %d: %s %v%% (%v/%v)",
+		s.PassNum, s.Task, pct, s.N, s.D)
+}
+
 func (o *rebuilder) rebuild(ctx context.Context) error {
 	passNum := 0
 	dlog.Infof(ctx, "... pass %d: scanning for implied items", passNum)
+	// Seed the queue
 	o.pendingAugments = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
-	btrfsutil.WalkAllTrees(ctx, o.inner, btrfsutil.WalkAllTreesHandler{
-		Err: func(*btrfsutil.WalkError) {},
-		TreeWalkHandler: btrfstree.TreeWalkHandler{
-			Item: func(path btrfstree.TreePath, item btrfstree.Item) error {
-				handleItem(o, ctx, path[0].FromTree, item)
-				return nil
-			},
-		},
-	})
-	for len(o.pendingAugments) > 0 {
-		// Apply the augments, keeping track of what keys are added to what tree.
-		dlog.Infof(ctx, "... pass %d: augmenting trees to add implied items", passNum)
-		newKeys := make(map[btrfsprim.ObjID][]btrfsprim.Key)
-		for _, treeID := range maps.SortedKeys(o.pendingAugments) {
-			dlog.Infof(ctx, "... ... augmenting tree %v:", treeID)
-			treeAugments := o.resolveTreeAugments(ctx, o.pendingAugments[treeID])
-			for _, nodeAddr := range maps.SortedKeys(treeAugments) {
-				added, err := o.inner.Augment(treeID, nodeAddr)
-				if err != nil {
-					dlog.Errorf(ctx, "error augmenting: %v", err)
-					continue
-				}
-				newKeys[treeID] = append(newKeys[treeID], added...)
-
-				set := o.augments[treeID]
-				if set == nil {
-					set = make(containers.Set[btrfsvol.LogicalAddr])
-					o.augments[treeID] = set
-				}
-				set.Insert(nodeAddr)
+	o.rebuilt.AddTree(ctx, btrfsprim.ROOT_TREE_OBJECTID)
+	o.rebuilt.AddTree(ctx, btrfsprim.CHUNK_TREE_OBJECTID)
+	o.rebuilt.AddTree(ctx, btrfsprim.TREE_LOG_OBJECTID)
+	o.rebuilt.AddTree(ctx, btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
+	for {
+		// Handle items in the queue
+		queue := o.queue
+		o.queue = nil
+		progressWriter := textui.NewProgress[rebuildStats](ctx, dlog.LogLevelInfo, 1*time.Second)
+		queueProgress := func(done int) {
+			progressWriter.Set(rebuildStats{
+				PassNum: passNum,
+				Task:    "processing item queue",
+				N:       done,
+				D:       len(queue),
+			})
+		}
+		for i, key := range queue {
+			queueProgress(i)
+			o.curKey = key
+			itemBody, ok := o.rebuilt.Load(ctx, key.TreeID, key.Key)
+			if !ok {
+				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
+			}
+			handleItem(o, ctx, key.TreeID, btrfstree.Item{
+				Key:  key.Key,
+				Body: itemBody,
+			})
+			if key.ItemType == btrfsitem.ROOT_ITEM_KEY {
+				o.rebuilt.AddTree(ctx, key.ObjectID)
 			}
 		}
-		// Clear the list of pending augments.
+		queueProgress(len(queue))
+		progressWriter.Done()
+
+		// Check if we can bail
+		if len(o.queue) == 0 && len(o.pendingAugments) == 0 {
+			break
+		}
+
+		// Apply augments that were requested while handling items from the queue
+		resolvedAugments := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], len(o.pendingAugments))
+		numAugments := 0
+		for _, treeID := range maps.SortedKeys(o.pendingAugments) {
+			resolvedAugments[treeID] = o.resolveTreeAugments(ctx, o.pendingAugments[treeID])
+			numAugments += len(resolvedAugments[treeID])
+		}
 		o.pendingAugments = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
-		passNum++
-		// Call handleItem() for each of the added keys.
-		dlog.Infof(ctx, "... pass %d: scanning for implied items", passNum)
-		for _, treeID := range maps.SortedKeys(newKeys) {
-			for _, key := range newKeys[treeID] {
-				item, err := o.inner.TreeLookup(treeID, key)
-				if err != nil {
-					o.ioErr(ctx, fmt.Errorf("error looking up already-inserted item: tree=%v key=%v: %w",
-						treeID, key, err))
-				}
-				handleItem(o, ctx, treeID, item)
+		progressWriter = textui.NewProgress[rebuildStats](ctx, dlog.LogLevelInfo, 1*time.Second)
+		numAugmented := 0
+		augmentProgress := func() {
+			progressWriter.Set(rebuildStats{
+				PassNum: passNum,
+				Task:    "applying augments",
+				N:       numAugmented,
+				D:       numAugments,
+			})
+		}
+		for _, treeID := range maps.SortedKeys(resolvedAugments) {
+			for _, nodeAddr := range maps.SortedKeys(resolvedAugments[treeID]) {
+				augmentProgress()
+				o.rebuilt.AddRoot(ctx, treeID, nodeAddr)
+				numAugmented++
 			}
 		}
+		augmentProgress()
+		progressWriter.Done()
+
+		passNum++
+		dlog.Infof(ctx, "... pass %d: scanning for implied items", passNum)
 	}
 	return nil
 }
 
+func (o *rebuilder) cbAddedItem(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key) {
+	o.queue = append(o.queue, keyio.KeyAndTree{
+		TreeID: tree,
+		Key:    key,
+	})
+}
+
+func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (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
+	}
+	itemBody, ok := o.rebuilt.Load(ctx, btrfsprim.ROOT_TREE_OBJECTID, key)
+	if !ok {
+		o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
+	}
+	switch itemBody := itemBody.(type) {
+	case btrfsitem.Root:
+		return 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
+	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.
+		panic(fmt.Errorf("should not happen: ROOT_ITEM item has unexpected type: %T", itemBody))
+	}
+}
+
+func (o *rebuilder) cbLookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool) {
+	key := btrfsitem.UUIDToKey(uuid)
+	if ok := o._wantOff(ctx, btrfsprim.UUID_TREE_OBJECTID, key.ObjectID, key.ItemType, key.Offset); !ok {
+		o.queue = append(o.queue, o.curKey)
+		return 0, false
+	}
+	itemBody, ok := o.rebuilt.Load(ctx, btrfsprim.UUID_TREE_OBJECTID, key)
+	if !ok {
+		o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
+	}
+	switch itemBody := itemBody.(type) {
+	case btrfsitem.UUIDMap:
+		return itemBody.ObjID, true
+	case btrfsitem.Error:
+		o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", btrfsprim.UUID_TREE_OBJECTID, key, itemBody.Err))
+		return 0, false
+	default:
+		// This is a panic because the item decoder should not emit UUID_SUBVOL items as anything but
+		// btrfsitem.UUIDMap or btrfsitem.Error without this code also being updated.
+		panic(fmt.Errorf("should not happen: UUID_SUBVOL item has unexpected type: %T", itemBody))
+	}
+}
+
 func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances []map[btrfsvol.LogicalAddr]int) containers.Set[btrfsvol.LogicalAddr] {
 	distances := make(map[btrfsvol.LogicalAddr]int)
 	generations := make(map[btrfsvol.LogicalAddr]btrfsprim.Generation)
@@ -268,37 +345,10 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-// _NodeFile is a subset of btrfstree.NodeFile.
-type _NodeFile interface {
-	ParentTree(btrfsprim.ObjID) (btrfsprim.ObjID, bool)
-}
-
-func treeDistance(fs _NodeFile, tree, leaf btrfsprim.ObjID) (int, bool) {
-	dist := 0
-	for {
-		if tree == leaf {
-			return dist, true
-		}
-
-		parentTree, ok := fs.ParentTree(tree)
-		if !ok {
-			// Failed to look up parent info.
-			return 0, false
-		}
-		if parentTree == 0 {
-			// End of the line.
-			return 0, false
-		}
-
-		tree = parentTree
-		dist++
-	}
-}
-
 func (o *rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, choices containers.Set[btrfsvol.LogicalAddr]) {
 	choicesWithDist := make(map[btrfsvol.LogicalAddr]int)
 	for choice := range choices {
-		if dist, ok := treeDistance(o.uuidMap, treeID, o.graph.Nodes[choice].Owner); ok {
+		if dist, ok := o.rebuilt.COWDistance(ctx, treeID, o.graph.Nodes[choice].Owner); ok {
 			choicesWithDist[choice] = dist
 		}
 	}
@@ -319,17 +369,25 @@ func (o *rebuilder) fsErr(ctx context.Context, e error) {
 
 // want implements rebuildCallbacks.
 func (o *rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) {
+	o._want(ctx, treeID, objID, typ)
+}
+func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) (key btrfsprim.Key, ok bool) {
+	if !o.rebuilt.AddTree(ctx, treeID) {
+		o.queue = append(o.queue, o.curKey)
+		return btrfsprim.Key{}, false
+	}
+
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	if _, err := o.inner.TreeSearch(treeID, func(key btrfsprim.Key, _ uint32) int {
+	if key, ok := o.rebuilt.Search(ctx, treeID, func(key btrfsprim.Key) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
-	}); err == nil {
-		return
+	}); ok {
+		return key, true
 	}
 
 	// OK, we need to insert it
@@ -339,14 +397,23 @@ func (o *rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrf
 	o.keyIO.Keys.Subrange(
 		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
 		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
-			wants.InsertFrom(o.leaf2orphans[v.Node])
+			wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
+	return btrfsprim.Key{}, false
 }
 
 // wantOff implements rebuildCallbacks.
 func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) {
+	o._wantOff(ctx, treeID, objID, typ, off)
+}
+func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) (ok bool) {
+	if !o.rebuilt.AddTree(ctx, treeID) {
+		o.queue = append(o.queue, o.curKey)
+		return false
+	}
+
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
@@ -354,8 +421,8 @@ func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID b
 		ItemType: typ,
 		Offset:   off,
 	}
-	if _, err := o.inner.TreeLookup(treeID, tgt); err == nil {
-		return
+	if _, ok := o.rebuilt.Load(ctx, treeID, tgt); ok {
+		return true
 	}
 
 	// OK, we need to insert it
@@ -365,26 +432,36 @@ func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID b
 	o.keyIO.Keys.Subrange(
 		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { return tgt.Cmp(k.Key) },
 		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
-			wants.InsertFrom(o.leaf2orphans[v.Node])
+			wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
+	return false
 }
 
 // wantFunc implements rebuildCallbacks.
 func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(btrfsitem.Item) bool) {
+	if !o.rebuilt.AddTree(ctx, treeID) {
+		o.queue = append(o.queue, o.curKey)
+		return
+	}
+
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	items, _ := o.inner.TreeSearchAll(treeID, func(key btrfsprim.Key, _ uint32) int {
+	keys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
 	})
-	for _, item := range items {
-		if fn(item.Body) {
+	for _, itemKey := range keys {
+		itemBody, ok := o.rebuilt.Load(ctx, treeID, itemKey)
+		if !ok {
+			o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", itemKey))
+		}
+		if fn(itemBody) {
 			return
 		}
 	}
@@ -398,10 +475,10 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 		func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
 			itemBody, ok := o.keyIO.ReadItem(v)
 			if !ok {
-				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
+				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v at %v", k, v))
 			}
 			if fn(itemBody) {
-				wants.InsertFrom(o.leaf2orphans[v.Node])
+				wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			}
 			return true
 		})
@@ -412,35 +489,42 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 //
 // interval is [beg, end)
 func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
-	for beg < end {
-		// check if we already have it
-		if run, err := btrfs.LookupCSum(o.inner, o.sb.ChecksumType, beg); err == nil {
-			// we already have it
-			beg = run.Addr.Add(run.Size())
-		} else {
-			// we need to insert it
-			ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("csum for laddr=%v", beg))
-			rbNode := o.csums.Search(func(item btrfsinspect.SysExtentCSum) int {
-				switch {
-				case item.Sums.Addr > beg:
-					return -1
-				case item.Sums.Addr.Add(item.Sums.Size()) <= beg:
-					return 1
-				default:
-					return 0
-				}
+	if !o.rebuilt.AddTree(ctx, btrfsprim.CSUM_TREE_OBJECTID) {
+		o.queue = append(o.queue, o.curKey)
+		return
+	}
 
-			})
-			if rbNode == nil {
-				o.wantAugment(ctx, btrfsprim.CSUM_TREE_OBJECTID, nil) // log an error
-				beg += btrfssum.BlockSize
-				continue
-			}
-			run := rbNode.Value.Sums
-			key := keyio.KeyAndTree{
-				Key:    rbNode.Value.Key,
-				TreeID: btrfsprim.CSUM_TREE_OBJECTID,
+	for beg < end {
+		// Figure out which key we want.
+
+		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("csum for laddr=%v", beg))
+		rbNode := o.csums.Search(func(item btrfsinspect.SysExtentCSum) int {
+			switch {
+			case item.Sums.Addr > beg:
+				return -1
+			case item.Sums.Addr.Add(item.Sums.Size()) <= beg:
+				return 1
+			default:
+				return 0
 			}
+
+		})
+		if rbNode == nil {
+			o.wantAugment(ctx, btrfsprim.CSUM_TREE_OBJECTID, nil) // log an error
+			beg += btrfssum.BlockSize
+			continue
+		}
+		run := rbNode.Value.Sums
+		key := keyio.KeyAndTree{
+			Key:    rbNode.Value.Key,
+			TreeID: btrfsprim.CSUM_TREE_OBJECTID,
+		}
+
+		// Check if we already have it.
+
+		// .Search is more efficient than .Load, because it doesn't load the body (and we don't need the body).
+		if _, ok := o.rebuilt.Search(ctx, key.TreeID, key.Key.Cmp); !ok {
+			// We need to insert it.
 			itemPtr, ok := o.keyIO.Keys.Load(key)
 			if !ok {
 				// This is a panic because if we found it in `o.csums` then it has
@@ -448,15 +532,20 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 				// btrfs.LookupCSum(), then that Node must be an orphan.
 				panic(fmt.Errorf("should not happen: no orphan contains %v", key.Key))
 			}
-			o.wantAugment(ctx, key.TreeID, o.leaf2orphans[itemPtr.Node])
-
-			beg = run.Addr.Add(run.Size())
+			o.wantAugment(ctx, key.TreeID, o.rebuilt.LeafToRoots(ctx, key.TreeID, itemPtr.Node))
 		}
+
+		beg = run.Addr.Add(run.Size())
 	}
 }
 
 // wantFileExt implements rebuildCallbacks.
 func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
+	if !o.rebuilt.AddTree(ctx, treeID) {
+		o.queue = append(o.queue, o.curKey)
+		return
+	}
+
 	min := btrfsprim.Key{
 		ObjectID: ino,
 		ItemType: btrfsitem.EXTENT_DATA_KEY,
@@ -467,7 +556,7 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 		ItemType: btrfsitem.EXTENT_DATA_KEY,
 		Offset:   uint64(size - 1),
 	}
-	exts, _ := o.inner.TreeSearchAll(treeID, func(key btrfsprim.Key, _ uint32) int {
+	extKeys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
 		switch {
 		case min.Cmp(key) < 0:
 			return 1
@@ -491,13 +580,18 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 		Beg: 0,
 		End: size,
 	})
-	for _, ext := range exts {
-		switch extBody := ext.Body.(type) {
+	for _, extKey := range extKeys {
+		extBody, ok := o.rebuilt.Load(ctx, treeID, extKey)
+		if !ok {
+			o.ioErr(ctx, fmt.Errorf("could not look up already-inserted item: tree=%v key=%v",
+				treeID, extKey))
+		}
+		switch extBody := extBody.(type) {
 		case btrfsitem.FileExtent:
-			extBeg := int64(ext.Key.Offset)
+			extBeg := int64(extKey.Offset)
 			extSize, err := extBody.Size()
 			if err != nil {
-				o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, ext.Key, err))
+				o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, extKey, err))
 				continue
 			}
 			extEnd := extBeg + extSize
@@ -532,7 +626,7 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 				})
 			}
 		case btrfsitem.Error:
-			o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, ext.Key, extBody.Err))
+			o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, extKey, extBody.Err))
 		default:
 			// This is a panic because the item decoder should not emit EXTENT_DATA
 			// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
@@ -587,7 +681,7 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 					// TODO(lukeshu): Re-evaluate whether being greedy here is the right
 					// thing.
 					if itemEnd > gap.Beg && itemBeg < gap.End {
-						wants.InsertFrom(o.leaf2orphans[v.Node])
+						wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 					}
 				case btrfsitem.Error:
 					o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, k, itemBody.Err))
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index 7e01693..8c519fc 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -17,14 +17,11 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
 type scanResult struct {
-	uuidMap   *uuidmap.UUIDMap
 	nodeGraph *graph.Graph
 	keyIO     *keyio.Handle
 }
@@ -55,7 +52,6 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	}
 
 	ret := &scanResult{
-		uuidMap:   uuidmap.New(),
 		nodeGraph: graph.New(*sb),
 	}
 	ret.keyIO = keyio.NewHandle(fs, *sb, ret.nodeGraph)
@@ -70,10 +66,6 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 				return nil, err
 			}
 
-			if err := ret.uuidMap.InsertNode(nodeRef); err != nil {
-				return nil, err
-			}
-
 			ret.nodeGraph.InsertNode(nodeRef)
 			ret.keyIO.InsertNode(nodeRef)
 
@@ -85,12 +77,6 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 		panic("should not happen")
 	}
 	progressWriter.Done()
-
-	missing := ret.uuidMap.MissingRootItems()
-	if len(missing) > 0 {
-		dlog.Errorf(ctx, "... could not find root items for %d trees: %v", len(missing), maps.SortedKeys(missing))
-	}
-
 	dlog.Info(ctx, "... done reading node data")
 
 	progressWriter = textui.NewProgress[scanStats](ctx, dlog.LogLevelInfo, 1*time.Second)
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/scan.go
deleted file mode 100644
index 98ed097..0000000
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/scan.go
+++ /dev/null
@@ -1,73 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package uuidmap
-
-import (
-	"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"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
-	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
-)
-
-func maybeSet[K, V comparable](name string, m map[K]V, k K, v V) error {
-	if other, conflict := m[k]; conflict && other != v {
-		return fmt.Errorf("conflict: %s %v can't have both %v and %v", name, k, other, v)
-	}
-	m[k] = v
-	return nil
-}
-
-func New() *UUIDMap {
-	ret := &UUIDMap{
-		ObjID2UUID: make(map[btrfsprim.ObjID]btrfsprim.UUID),
-		UUID2ObjID: make(map[btrfsprim.UUID]btrfsprim.ObjID),
-		TreeParent: make(map[btrfsprim.ObjID]btrfsprim.UUID),
-
-		SeenTrees: make(containers.Set[btrfsprim.ObjID]),
-	}
-
-	// These 4 trees are mentioned directly in the superblock, so
-	// they are always seen.
-	ret.SeenTrees.Insert(btrfsprim.ROOT_TREE_OBJECTID)
-	ret.SeenTrees.Insert(btrfsprim.CHUNK_TREE_OBJECTID)
-	ret.SeenTrees.Insert(btrfsprim.TREE_LOG_OBJECTID)
-	ret.SeenTrees.Insert(btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
-
-	return ret
-}
-
-func (m *UUIDMap) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
-	for _, item := range nodeRef.Data.BodyLeaf {
-		switch itemBody := item.Body.(type) {
-		case btrfsitem.Root:
-			if err := maybeSet("ObjID2UUID", m.ObjID2UUID, item.Key.ObjectID, itemBody.UUID); err != nil {
-				return err
-			}
-			if itemBody.UUID != (btrfsprim.UUID{}) {
-				if err := maybeSet("UUID2ObjID", m.UUID2ObjID, itemBody.UUID, item.Key.ObjectID); err != nil {
-					return err
-				}
-			}
-			if err := maybeSet("ParentUUID", m.TreeParent, item.Key.ObjectID, itemBody.ParentUUID); err != nil {
-				return err
-			}
-			m.SeenTrees.Insert(item.Key.ObjectID)
-		case btrfsitem.UUIDMap:
-			uuid := btrfsitem.KeyToUUID(item.Key)
-			if err := maybeSet("ObjID2UUID", m.ObjID2UUID, itemBody.ObjID, uuid); err != nil {
-				return err
-			}
-			if err := maybeSet("UUID2ObjID", m.UUID2ObjID, uuid, itemBody.ObjID); err != nil {
-				return err
-			}
-		}
-	}
-	m.SeenTrees.Insert(nodeRef.Data.Head.Owner)
-	return nil
-}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/uuidmap.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/uuidmap.go
deleted file mode 100644
index 32a34d1..0000000
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/uuidmap.go
+++ /dev/null
@@ -1,55 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package uuidmap
-
-import (
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
-	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-)
-
-type UUIDMap struct {
-	ObjID2UUID map[btrfsprim.ObjID]btrfsprim.UUID
-	UUID2ObjID map[btrfsprim.UUID]btrfsprim.ObjID
-	TreeParent map[btrfsprim.ObjID]btrfsprim.UUID
-
-	SeenTrees containers.Set[btrfsprim.ObjID]
-}
-
-func (m UUIDMap) MissingRootItems() containers.Set[btrfsprim.ObjID] {
-	missing := make(containers.Set[btrfsprim.ObjID])
-	for treeID := range m.SeenTrees {
-		if _, ok := m.ObjID2UUID[treeID]; !ok && treeID != btrfsprim.ROOT_TREE_OBJECTID {
-			missing.Insert(treeID)
-			continue
-		}
-		if _, ok := m.TreeParent[treeID]; !ok && treeID >= btrfsprim.FIRST_FREE_OBJECTID && treeID <= btrfsprim.LAST_FREE_OBJECTID {
-			missing.Insert(treeID)
-		}
-	}
-	return missing
-}
-
-// ParentTree implements btrfstree.NodeFile.
-func (m UUIDMap) ParentTree(tree btrfsprim.ObjID) (btrfsprim.ObjID, bool) {
-	if tree < btrfsprim.FIRST_FREE_OBJECTID || tree > btrfsprim.LAST_FREE_OBJECTID {
-		// no parent
-		return 0, true
-	}
-	parentUUID, ok := m.TreeParent[tree]
-	if !ok {
-		// could not look up parent info
-		return 0, false
-	}
-	if parentUUID == (btrfsprim.UUID{}) {
-		// no parent
-		return 0, true
-	}
-	parentObjID, ok := m.UUID2ObjID[parentUUID]
-	if !ok {
-		// could not look up parent info
-		return 0, false
-	}
-	return parentObjID, true
-}
-- 
cgit v1.2.3-54-g00ecf


From a476f616f1e4367a212ef9b03434f593b39692e5 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Thu, 22 Dec 2022 02:56:56 -0700
Subject: rebuildnodes: Improve logging

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 2ddce88..e05399e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -146,6 +146,7 @@ func (o *rebuilder) rebuild(ctx context.Context) error {
 		resolvedAugments := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], len(o.pendingAugments))
 		numAugments := 0
 		for _, treeID := range maps.SortedKeys(o.pendingAugments) {
+			dlog.Infof(ctx, "... ... augments for tree %v:", treeID)
 			resolvedAugments[treeID] = o.resolveTreeAugments(ctx, o.pendingAugments[treeID])
 			numAugments += len(resolvedAugments[treeID])
 		}
@@ -337,7 +338,12 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances
 	}
 
 	for i, list := range lists {
-		dlog.Infof(ctx, "... ... ... %d: %v: %v", i, list.Intersection(ret).TakeOne(), maps.SortedKeys(list))
+		chose := list.Intersection(ret)
+		if len(chose) == 0 {
+			dlog.Infof(ctx, "... ... ... lists[%d]: chose (none) from %v", i, maps.SortedKeys(list))
+		} else {
+			dlog.Infof(ctx, "... ... ... lists[%d]: chose %v from %v", i, chose.TakeOne(), maps.SortedKeys(list))
+		}
 	}
 
 	return ret
-- 
cgit v1.2.3-54-g00ecf


From ec31481b46fb772dfe2e977ef37fc0ace39fad8a Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Thu, 22 Dec 2022 16:30:27 -0700
Subject: rebuildnodes/btrees: Try to handle duplicate keys by generation

---
 .../rebuildnodes/btrees/rebuilt_btrees.go          | 56 ++++++++++++++--------
 1 file changed, 37 insertions(+), 19 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index 23e974e..cbe7df1 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -136,17 +136,18 @@ type rootStats struct {
 	TreeID   btrfsprim.ObjID
 	RootNode btrfsvol.LogicalAddr
 
-	DoneLeafs  int
-	TotalLeafs int
-	AddedItems int
+	DoneLeafs     int
+	TotalLeafs    int
+	AddedItems    int
+	ReplacedItems int
 }
 
 func (s rootStats) String() string {
-	return fmt.Sprintf("tree %v: adding root node@%v: %v%% (%v/%v) (added %v items)",
+	return fmt.Sprintf("tree %v: adding root node@%v: %v%% (%v/%v) (added %v items, replaced %v items)",
 		s.TreeID, s.RootNode,
 		int(100*float64(s.DoneLeafs)/float64(s.TotalLeafs)),
 		s.DoneLeafs, s.TotalLeafs,
-		s.AddedItems)
+		s.AddedItems, s.ReplacedItems)
 }
 
 // AddRoot adds an additional root node to an existing tree.  It is
@@ -161,13 +162,15 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 
 	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, 1*time.Second)
 	numAdded := 0
+	numReplaced := 0
 	progress := func(done int) {
 		progressWriter.Set(rootStats{
-			TreeID:     treeID,
-			RootNode:   rootNode,
-			DoneLeafs:  done,
-			TotalLeafs: len(tree.leafToRoots),
-			AddedItems: numAdded,
+			TreeID:        treeID,
+			RootNode:      rootNode,
+			DoneLeafs:     done,
+			TotalLeafs:    len(tree.leafToRoots),
+			AddedItems:    numAdded,
+			ReplacedItems: numReplaced,
 		})
 	}
 	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
@@ -177,17 +180,32 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 			continue
 		}
 		for j, item := range ts.keyIO.ReadNode(leaf).Data.BodyLeaf {
-			if _, exists := tree.Items.Load(item.Key); exists {
-				// This is a panic because I'm not really sure what the best way to
-				// 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", item.Key, treeID))
-			}
-			tree.Items.Store(item.Key, keyio.ItemPtr{
+			newPtr := keyio.ItemPtr{
 				Node: leaf,
 				Idx:  j,
-			})
-			numAdded++
+			}
+			if oldPtr, exists := tree.Items.Load(item.Key); !exists {
+				tree.Items.Store(item.Key, 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)
+					numReplaced++
+				case oldGen > newGen:
+					// keep the old one
+				case oldGen == newGen:
+					// This is a panic because I'm not really sure what the best way to
+					// 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,
+						oldPtr, ts.graph.Nodes[oldPtr.Node],
+						newPtr, ts.graph.Nodes[newPtr.Node]))
+				}
+			}
 			ts.cbAddedItem(ctx, treeID, item.Key)
 			progress(i)
 		}
-- 
cgit v1.2.3-54-g00ecf


From 3e22a30a7febd9890163966d01dd2b34656c9e39 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Thu, 22 Dec 2022 18:24:03 -0700
Subject: rebuildnodes: Handle the same node being added twice

---
 .../btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go   | 16 ++++++++++++----
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go      | 16 +++++++++-------
 2 files changed, 21 insertions(+), 11 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index cbe7df1..a4ccd5e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -34,6 +34,7 @@ type rebuiltTree struct {
 
 	// mutable
 	Roots containers.Set[btrfsvol.LogicalAddr]
+	Leafs containers.Set[btrfsvol.LogicalAddr]
 	Items containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
 }
 
@@ -175,10 +176,10 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 	}
 	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
 		progress(i)
-		roots := tree.leafToRoots[leaf]
-		if !roots.Has(rootNode) {
+		if tree.Leafs.Has(leaf) || !tree.leafToRoots[leaf].Has(rootNode) {
 			continue
 		}
+		tree.Leafs.Insert(leaf)
 		for j, item := range ts.keyIO.ReadNode(leaf).Data.BodyLeaf {
 			newPtr := keyio.ItemPtr{
 				Node: leaf,
@@ -234,6 +235,7 @@ func (ts *RebuiltTrees) addTree(ctx context.Context, treeID btrfsprim.ObjID, sta
 	tree := &rebuiltTree{
 		ID:    treeID,
 		Roots: make(containers.Set[btrfsvol.LogicalAddr]),
+		Leafs: make(containers.Set[btrfsvol.LogicalAddr]),
 	}
 	var root btrfsvol.LogicalAddr
 	switch treeID {
@@ -404,11 +406,17 @@ func (ts *RebuiltTrees) LeafToRoots(ctx context.Context, treeID btrfsprim.ObjID,
 	if !ts.AddTree(ctx, treeID) {
 		return nil
 	}
+	if ts.graph.Nodes[leaf].Level != 0 {
+		panic(fmt.Errorf("should not happen: NodeToRoots(tree=%v, leaf=%v): not a leaf",
+			treeID, leaf))
+	}
 	ret := make(containers.Set[btrfsvol.LogicalAddr])
 	for root := range ts.trees[treeID].leafToRoots[leaf] {
-		if !ts.trees[treeID].Roots.Has(root) {
-			ret.Insert(root)
+		if ts.trees[treeID].Roots.Has(root) {
+			panic(fmt.Errorf("should not happen: NodeToRoots(tree=%v, leaf=%v): tree contains root=%v but not leaf",
+				treeID, leaf, root))
 		}
+		ret.Insert(root)
 	}
 	if len(ret) == 0 {
 		return nil
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index e05399e..4e60632 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -352,16 +352,18 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 func (o *rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, choices containers.Set[btrfsvol.LogicalAddr]) {
-	choicesWithDist := make(map[btrfsvol.LogicalAddr]int)
-	for choice := range choices {
-		if dist, ok := o.rebuilt.COWDistance(ctx, treeID, o.graph.Nodes[choice].Owner); ok {
-			choicesWithDist[choice] = dist
-		}
-	}
-	if len(choicesWithDist) == 0 {
+	if len(choices) == 0 {
 		dlog.Errorf(ctx, "augment(tree=%v): could not find wanted item", treeID)
 		return
 	}
+	choicesWithDist := make(map[btrfsvol.LogicalAddr]int, len(choices))
+	for choice := range choices {
+		dist, ok := o.rebuilt.COWDistance(ctx, treeID, o.graph.Nodes[choice].Owner)
+		if !ok {
+			panic(fmt.Errorf("should not happen: .wantAugment called for tree=%v with invalid choice=%v", treeID, choice))
+		}
+		choicesWithDist[choice] = dist
+	}
 	dlog.Infof(ctx, "augment(tree=%v): %v", treeID, maps.SortedKeys(choicesWithDist))
 	o.pendingAugments[treeID] = append(o.pendingAugments[treeID], choicesWithDist)
 }
-- 
cgit v1.2.3-54-g00ecf


From fc6f8871c39d3a74153e4ce8fd440c256f65b650 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 16:03:27 -0700
Subject: rebuildnodes/btrees: Track the tree's root item offset

---
 .../btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go     | 14 ++++++++------
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go        |  8 ++++----
 2 files changed, 12 insertions(+), 10 deletions(-)

(limited to 'lib')

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.
-- 
cgit v1.2.3-54-g00ecf


From 51d6ccf1e5af4d88bfe6842b06fa1bc2d5cc732c Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 17:10:47 -0700
Subject: rebuildnodes: Have keyio.Handle always be a pointer, fuss with
 ScanDevices return type

---
 .../rebuildnodes/btrees/rebuilt_btrees.go          |  4 ++--
 .../btrfsinspect/rebuildnodes/rebuild.go           | 10 ++++----
 lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go   | 27 ++++++++--------------
 3 files changed, 17 insertions(+), 24 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index 3a6c35b..1d5ba14 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -102,7 +102,7 @@ type RebuiltTrees struct {
 	// static
 	sb    btrfstree.Superblock
 	graph pkggraph.Graph
-	keyIO keyio.Handle
+	keyIO *keyio.Handle
 
 	// static callbacks
 	cbAddedItem  func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key)
@@ -116,7 +116,7 @@ type RebuiltTrees struct {
 // NewRebuiltTrees returns a new RebuiltTrees instance.  All of the
 // callbacks must be non-nil.
 func NewRebuiltTrees(
-	sb btrfstree.Superblock, graph pkggraph.Graph, keyIO keyio.Handle,
+	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) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool),
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool),
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 180edab..fecf5e6 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -34,7 +34,7 @@ type rebuilder struct {
 
 	graph graph.Graph
 	csums containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
-	keyIO keyio.Handle
+	keyIO *keyio.Handle
 
 	curKey          keyio.KeyAndTree
 	queue           []keyio.KeyAndTree
@@ -42,7 +42,7 @@ type rebuilder struct {
 }
 
 func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], error) {
-	scanData, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
+	nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
 	if err != nil {
 		return nil, err
 	}
@@ -63,11 +63,11 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	o := &rebuilder{
 		sb: *sb,
 
-		graph: *scanData.nodeGraph,
+		graph: nodeGraph,
 		csums: *csums,
-		keyIO: *scanData.keyIO,
+		keyIO: keyIO,
 	}
-	o.rebuilt = btrees.NewRebuiltTrees(*sb, *scanData.nodeGraph, *scanData.keyIO,
+	o.rebuilt = btrees.NewRebuiltTrees(*sb, nodeGraph, keyIO,
 		o.cbAddedItem, o.cbLookupRoot, o.cbLookupUUID)
 	if err := o.rebuild(ctx); err != nil {
 		return nil, err
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index 8c519fc..93d3e0c 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -21,11 +21,6 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
-type scanResult struct {
-	nodeGraph *graph.Graph
-	keyIO     *keyio.Handle
-}
-
 type scanStats struct {
 	N, D int
 }
@@ -36,12 +31,12 @@ func (s scanStats) String() string {
 		s.N, s.D)
 }
 
-func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (*scanResult, error) {
+func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (graph.Graph, *keyio.Handle, error) {
 	dlog.Infof(ctx, "Reading node data from FS...")
 
 	sb, err := fs.Superblock()
 	if err != nil {
-		return nil, err
+		return graph.Graph{}, nil, err
 	}
 
 	total := countNodes(scanResults)
@@ -51,10 +46,8 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 		progressWriter.Set(scanStats{N: done, D: total})
 	}
 
-	ret := &scanResult{
-		nodeGraph: graph.New(*sb),
-	}
-	ret.keyIO = keyio.NewHandle(fs, *sb, ret.nodeGraph)
+	nodeGraph := graph.New(*sb)
+	keyIO := keyio.NewHandle(fs, *sb, nodeGraph)
 
 	progress(done, total)
 	for _, devResults := range scanResults {
@@ -63,11 +56,11 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 			})
 			if err != nil {
-				return nil, err
+				return graph.Graph{}, nil, err
 			}
 
-			ret.nodeGraph.InsertNode(nodeRef)
-			ret.keyIO.InsertNode(nodeRef)
+			nodeGraph.InsertNode(nodeRef)
+			keyIO.InsertNode(nodeRef)
 
 			done++
 			progress(done, total)
@@ -81,11 +74,11 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 
 	progressWriter = textui.NewProgress[scanStats](ctx, dlog.LogLevelInfo, 1*time.Second)
 	dlog.Infof(ctx, "Checking keypointers for dead-ends...")
-	if err := ret.nodeGraph.FinalCheck(fs, *sb, progress); err != nil {
-		return nil, err
+	if err := nodeGraph.FinalCheck(fs, *sb, progress); err != nil {
+		return graph.Graph{}, nil, err
 	}
 	progressWriter.Done()
 	dlog.Info(ctx, "... done checking keypointers")
 
-	return ret, nil
+	return *nodeGraph, keyIO, nil
 }
-- 
cgit v1.2.3-54-g00ecf


From 8fc3c78924da1dedd3adce3b923adf58e5ca732a Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 17:50:46 -0700
Subject: rebuildnodes: Have the graph store keys; avoid I/O when indexing a
 tree

---
 .../btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go       | 12 ++++++------
 lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go      | 11 +++++++++--
 2 files changed, 15 insertions(+), 8 deletions(-)

(limited to 'lib')

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{
-- 
cgit v1.2.3-54-g00ecf


From 6330b26a9f9c7769c48e2bc90e065457f8e138ab Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 16:48:42 -0700
Subject: rebuildnodes: Have the key index belong to the btree, and be smarter

---
 .../rebuildnodes/btrees/rebuilt_btrees.go          | 106 ++++++++++++++-------
 .../btrfsinspect/rebuildnodes/keyio/keyio.go       |  21 +---
 .../btrfsinspect/rebuildnodes/rebuild.go           |  45 +++++----
 lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go   |  13 +--
 4 files changed, 103 insertions(+), 82 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index f919b37..50c75a4 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -32,6 +32,7 @@ type rebuiltTree struct {
 
 	// all leafs (lvl=0) that pass .isOwnerOK, even if not in the tree
 	leafToRoots map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
+	keys        containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
 
 	// mutable
 	Roots containers.Set[btrfsvol.LogicalAddr]
@@ -41,14 +42,14 @@ type rebuiltTree struct {
 
 // isOwnerOK returns whether it is permissible for a node with
 // .Head.Owner=owner to be in this tree.
-func (tree *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID) bool {
+func (tree *rebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generation) bool {
 	for {
-		if tree == nil {
-			return false
-		}
 		if owner == tree.ID {
 			return true
 		}
+		if tree.Parent == nil || gen >= tree.ParentGen {
+			return false
+		}
 		tree = tree.Parent
 	}
 }
@@ -68,6 +69,38 @@ func (tree *rebuiltTree) cowDistance(parentID btrfsprim.ObjID) (dist int, ok boo
 	}
 }
 
+func (tree *rebuiltTree) shouldReplace(graph pkggraph.Graph, oldNode, newNode btrfsvol.LogicalAddr) bool {
+	oldDist, _ := tree.cowDistance(graph.Nodes[oldNode].Owner)
+	newDist, _ := tree.cowDistance(graph.Nodes[newNode].Owner)
+	switch {
+	case newDist < oldDist:
+		// Replace the old one with the new lower-dist one.
+		return true
+	case newDist > oldDist:
+		// Retain the old lower-dist one.
+		return false
+	default:
+		oldGen := graph.Nodes[oldNode].Generation
+		newGen := graph.Nodes[newNode].Generation
+		switch {
+		case newGen > oldGen:
+			// Replace the old one with the new higher-gen one.
+			return true
+		case newGen < oldGen:
+			// Retain the old higher-gen one.
+			return false
+		default:
+			// This is a panic because I'm not really sure what the best way to
+			// 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 nodes in tree=%v: old=%v=%v ; new=%v=%v",
+				tree.ID,
+				oldNode, graph.Nodes[oldNode],
+				newNode, graph.Nodes[newNode]))
+		}
+	}
+}
+
 // RebuiltTrees is an abstraction for rebuilding and accessing
 // potentially broken btrees.
 //
@@ -189,24 +222,9 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 			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(itemKey, newPtr)
-					numReplaced++
-				case oldGen > newGen:
-					// keep the old one
-				case oldGen == newGen:
-					// This is a panic because I'm not really sure what the best way to
-					// 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",
-						itemKey, treeID,
-						oldPtr, ts.graph.Nodes[oldPtr.Node],
-						newPtr, ts.graph.Nodes[newPtr.Node]))
-				}
+			} else if tree.shouldReplace(ts.graph, oldPtr.Node, newPtr.Node) {
+				tree.Items.Store(itemKey, newPtr)
+				numReplaced++
 			}
 			ts.cbAddedItem(ctx, treeID, itemKey)
 			progress(i)
@@ -327,26 +345,42 @@ func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAd
 		return
 	}
 	if slices.Contains(node, stack) {
-		return
+		panic("loop")
 	}
-	if !tree.isOwnerOK(graph.Nodes[node].Owner) {
+	if !tree.isOwnerOK(graph.Nodes[node].Owner, graph.Nodes[node].Generation) {
 		index[node] = nil
 		return
 	}
+
+	// tree.leafToRoots
+	stack = append(stack, node)
+	var roots containers.Set[btrfsvol.LogicalAddr]
 	kps := slices.RemoveAllFunc(graph.EdgesTo[node], func(kp *pkggraph.Edge) bool {
-		return !tree.isOwnerOK(graph.Nodes[kp.FromNode].Owner)
+		return !tree.isOwnerOK(graph.Nodes[kp.FromNode].Owner, graph.Nodes[kp.FromNode].Generation)
 	})
-	if len(kps) == 0 {
-		index[node] = containers.NewSet[btrfsvol.LogicalAddr](node)
-		return
-	}
-	stack = append(stack, node)
-	roots := make(containers.Set[btrfsvol.LogicalAddr])
 	for _, kp := range kps {
 		tree.indexNode(graph, kp.FromNode, index, progress, stack)
-		roots.InsertFrom(index[kp.FromNode])
+		if len(index[kp.FromNode]) > 0 {
+			if roots == nil {
+				roots = make(containers.Set[btrfsvol.LogicalAddr])
+			}
+			roots.InsertFrom(index[kp.FromNode])
+		}
+	}
+	if roots == nil {
+		roots = containers.NewSet[btrfsvol.LogicalAddr](node)
 	}
 	index[node] = roots
+
+	// tree.keys
+	for i, key := range graph.Nodes[node].Items {
+		if oldPtr, ok := tree.keys.Load(key); !ok || tree.shouldReplace(graph, oldPtr.Node, node) {
+			tree.keys.Store(key, keyio.ItemPtr{
+				Node: node,
+				Idx:  i,
+			})
+		}
+	}
 }
 
 // Load reads an item from a tree.
@@ -426,6 +460,14 @@ func (ts *RebuiltTrees) LeafToRoots(ctx context.Context, treeID btrfsprim.ObjID,
 	return ret
 }
 
+// Keys returns a map of all keys in node that would be valid in this tree.
+//
+// It is invalid (panic) to call Keys for a tree without having called
+// AddTree first.
+func (ts *RebuiltTrees) Keys(treeID btrfsprim.ObjID) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
+	return &ts.trees[treeID].keys
+}
+
 // COWDistance returns how many COW-snapshots down from the 'child'
 // tree is from the 'parent' tree.
 //
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
index 875b616..5ae9ce2 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
@@ -38,16 +38,14 @@ func (ptr ItemPtr) String() string {
 }
 
 type Handle struct {
-	Keys containers.SortedMap[KeyAndTree, ItemPtr]
-
 	rawFile diskio.File[btrfsvol.LogicalAddr]
 	sb      btrfstree.Superblock
-	graph   *graph.Graph
+	graph   graph.Graph
 
 	cache *containers.LRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]]
 }
 
-func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph *graph.Graph) *Handle {
+func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph graph.Graph) *Handle {
 	return &Handle{
 		rawFile: file,
 		sb:      sb,
@@ -57,21 +55,6 @@ func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock,
 	}
 }
 
-func (o *Handle) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
-	for i, item := range nodeRef.Data.BodyLeaf {
-		k := KeyAndTree{
-			Key:    item.Key,
-			TreeID: nodeRef.Data.Head.Owner,
-		}
-		if cur, ok := o.Keys.Load(k); !ok || o.graph.Nodes[cur.Node].Generation < nodeRef.Data.Head.Generation {
-			o.Keys.Store(k, ItemPtr{
-				Node: nodeRef.Addr,
-				Idx:  i,
-			})
-		}
-	}
-}
-
 func (o *Handle) ReadNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node] {
 	if cached, ok := o.cache.Get(laddr); ok {
 		return cached
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index fecf5e6..982c27d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -42,7 +42,7 @@ type rebuilder struct {
 }
 
 func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], error) {
-	nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
+	nodeGraph, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
 	if err != nil {
 		return nil, err
 	}
@@ -60,6 +60,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Rebuilding node tree...")
+	keyIO := keyio.NewHandle(fs, *sb, nodeGraph)
 	o := &rebuilder{
 		sb: *sb,
 
@@ -402,9 +403,9 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btr
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v %v ?}", treeID, objID, typ))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.keyIO.Keys.Subrange(
-		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
-		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
+	o.rebuilt.Keys(treeID).Subrange(
+		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Cmp(k) },
+		func(_ btrfsprim.Key, v keyio.ItemPtr) bool {
 			wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			return true
 		})
@@ -437,9 +438,9 @@ func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key=%v", treeID, tgt))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.keyIO.Keys.Subrange(
-		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { return tgt.Cmp(k.Key) },
-		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
+	o.rebuilt.Keys(treeID).Subrange(
+		func(k btrfsprim.Key, _ keyio.ItemPtr) int { return tgt.Cmp(k) },
+		func(_ btrfsprim.Key, v keyio.ItemPtr) bool {
 			wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			return true
 		})
@@ -478,9 +479,9 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key=%v +func", treeID, tgt))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.keyIO.Keys.Subrange(
-		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
-		func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
+	o.rebuilt.Keys(treeID).Subrange(
+		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Cmp(k) },
+		func(k btrfsprim.Key, v keyio.ItemPtr) bool {
 			itemBody, ok := o.keyIO.ReadItem(v)
 			if !ok {
 				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v at %v", k, v))
@@ -523,24 +524,22 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 			continue
 		}
 		run := rbNode.Value.Sums
-		key := keyio.KeyAndTree{
-			Key:    rbNode.Value.Key,
-			TreeID: btrfsprim.CSUM_TREE_OBJECTID,
-		}
+		key := rbNode.Value.Key
+		treeID := btrfsprim.CSUM_TREE_OBJECTID
 
 		// Check if we already have it.
 
 		// .Search is more efficient than .Load, because it doesn't load the body (and we don't need the body).
-		if _, ok := o.rebuilt.Search(ctx, key.TreeID, key.Key.Cmp); !ok {
+		if _, ok := o.rebuilt.Search(ctx, treeID, key.Cmp); !ok {
 			// We need to insert it.
-			itemPtr, ok := o.keyIO.Keys.Load(key)
+			itemPtr, ok := o.rebuilt.Keys(btrfsprim.CSUM_TREE_OBJECTID).Load(key)
 			if !ok {
 				// This is a panic because if we found it in `o.csums` then it has
 				// to be in some Node, and if we didn't find it from
 				// btrfs.LookupCSum(), then that Node must be an orphan.
-				panic(fmt.Errorf("should not happen: no orphan contains %v", key.Key))
+				panic(fmt.Errorf("should not happen: no orphan contains %v", key))
 			}
-			o.wantAugment(ctx, key.TreeID, o.rebuilt.LeafToRoots(ctx, key.TreeID, itemPtr.Node))
+			o.wantAugment(ctx, treeID, o.rebuilt.LeafToRoots(ctx, treeID, itemPtr.Node))
 		}
 
 		beg = run.Addr.Add(run.Size())
@@ -656,18 +655,18 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 		}
 		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("file extent for tree=%v inode=%v bytes [%v, %v)", treeID, ino, gap.Beg, gap.End))
 		wants := make(containers.Set[btrfsvol.LogicalAddr])
-		o.keyIO.Keys.Subrange(
-			func(k keyio.KeyAndTree, _ keyio.ItemPtr) int {
+		o.rebuilt.Keys(treeID).Subrange(
+			func(k btrfsprim.Key, _ keyio.ItemPtr) int {
 				switch {
-				case min.Cmp(k.Key) < 0:
+				case min.Cmp(k) < 0:
 					return 1
-				case max.Cmp(k.Key) > 0:
+				case max.Cmp(k) > 0:
 					return -1
 				default:
 					return 0
 				}
 			},
-			func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
+			func(k btrfsprim.Key, v keyio.ItemPtr) bool {
 				itemBody, ok := o.keyIO.ReadItem(v)
 				if !ok {
 					o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index 93d3e0c..bcf2f5b 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -16,7 +16,6 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
@@ -31,12 +30,12 @@ func (s scanStats) String() string {
 		s.N, s.D)
 }
 
-func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (graph.Graph, *keyio.Handle, error) {
+func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (graph.Graph, error) {
 	dlog.Infof(ctx, "Reading node data from FS...")
 
 	sb, err := fs.Superblock()
 	if err != nil {
-		return graph.Graph{}, nil, err
+		return graph.Graph{}, err
 	}
 
 	total := countNodes(scanResults)
@@ -47,7 +46,6 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	}
 
 	nodeGraph := graph.New(*sb)
-	keyIO := keyio.NewHandle(fs, *sb, nodeGraph)
 
 	progress(done, total)
 	for _, devResults := range scanResults {
@@ -56,11 +54,10 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 			})
 			if err != nil {
-				return graph.Graph{}, nil, err
+				return graph.Graph{}, err
 			}
 
 			nodeGraph.InsertNode(nodeRef)
-			keyIO.InsertNode(nodeRef)
 
 			done++
 			progress(done, total)
@@ -75,10 +72,10 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	progressWriter = textui.NewProgress[scanStats](ctx, dlog.LogLevelInfo, 1*time.Second)
 	dlog.Infof(ctx, "Checking keypointers for dead-ends...")
 	if err := nodeGraph.FinalCheck(fs, *sb, progress); err != nil {
-		return graph.Graph{}, nil, err
+		return graph.Graph{}, err
 	}
 	progressWriter.Done()
 	dlog.Info(ctx, "... done checking keypointers")
 
-	return *nodeGraph, keyIO, nil
+	return *nodeGraph, nil
 }
-- 
cgit v1.2.3-54-g00ecf


From 7ca8e6805164a5855c4af6ec5103c8e0cbcab89c Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 18:00:47 -0700
Subject: rebuildnodes: Move keyio.KeyAndTree→keyAndtree
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../btrfsinspect/rebuildnodes/keyio/keyio.go           | 12 ------------
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go    | 18 +++++++++++++++---
 2 files changed, 15 insertions(+), 15 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
index 5ae9ce2..fa69c78 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
@@ -16,18 +16,6 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 )
 
-type KeyAndTree struct {
-	btrfsprim.Key
-	TreeID btrfsprim.ObjID
-}
-
-func (a KeyAndTree) Cmp(b KeyAndTree) int {
-	if d := a.Key.Cmp(b.Key); d != 0 {
-		return d
-	}
-	return containers.NativeCmp(a.TreeID, b.TreeID)
-}
-
 type ItemPtr struct {
 	Node btrfsvol.LogicalAddr
 	Idx  int
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 982c27d..00dbad5 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -28,6 +28,18 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
+type keyAndTree struct {
+	btrfsprim.Key
+	TreeID btrfsprim.ObjID
+}
+
+func (a keyAndTree) Cmp(b keyAndTree) int {
+	if d := a.Key.Cmp(b.Key); d != 0 {
+		return d
+	}
+	return containers.NativeCmp(a.TreeID, b.TreeID)
+}
+
 type rebuilder struct {
 	sb      btrfstree.Superblock
 	rebuilt *btrees.RebuiltTrees
@@ -36,8 +48,8 @@ type rebuilder struct {
 	csums containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
 	keyIO *keyio.Handle
 
-	curKey          keyio.KeyAndTree
-	queue           []keyio.KeyAndTree
+	curKey          keyAndTree
+	queue           []keyAndTree
 	pendingAugments map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int
 }
 
@@ -179,7 +191,7 @@ func (o *rebuilder) rebuild(ctx context.Context) error {
 }
 
 func (o *rebuilder) cbAddedItem(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key) {
-	o.queue = append(o.queue, keyio.KeyAndTree{
+	o.queue = append(o.queue, keyAndTree{
 		TreeID: tree,
 		Key:    key,
 	})
-- 
cgit v1.2.3-54-g00ecf


From 5fcb6f7ab77721faf0f89e00bf3087baba2d566c Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 19:41:39 -0700
Subject: rebuildnodes: Ignore the LOG_TREE for now

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 00dbad5..076710e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -117,7 +117,7 @@ func (o *rebuilder) rebuild(ctx context.Context) error {
 	o.pendingAugments = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
 	o.rebuilt.AddTree(ctx, btrfsprim.ROOT_TREE_OBJECTID)
 	o.rebuilt.AddTree(ctx, btrfsprim.CHUNK_TREE_OBJECTID)
-	o.rebuilt.AddTree(ctx, btrfsprim.TREE_LOG_OBJECTID)
+	//o.rebuilt.AddTree(ctx, btrfsprim.TREE_LOG_OBJECTID) // TODO(lukeshu): Special LOG_TREE handling
 	o.rebuilt.AddTree(ctx, btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
 	for {
 		// Handle items in the queue
-- 
cgit v1.2.3-54-g00ecf


From b3264cfa1a48998bc9406c18339801c21a468f05 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 19:54:56 -0700
Subject: rebuildnodes: Less noisy logging about csums

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 076710e..1daa83d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -510,11 +510,13 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 //
 // interval is [beg, end)
 func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
-	if !o.rebuilt.AddTree(ctx, btrfsprim.CSUM_TREE_OBJECTID) {
+	treeID := btrfsprim.CSUM_TREE_OBJECTID
+	if !o.rebuilt.AddTree(ctx, treeID) {
 		o.queue = append(o.queue, o.curKey)
 		return
 	}
 
+	last := beg
 	for beg < end {
 		// Figure out which key we want.
 
@@ -531,20 +533,20 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 
 		})
 		if rbNode == nil {
-			o.wantAugment(ctx, btrfsprim.CSUM_TREE_OBJECTID, nil) // log an error
 			beg += btrfssum.BlockSize
 			continue
+		} else if last < beg {
+			dlog.Errorf(ctx, "augment(tree=%v): could not find csum items for %v-%v", treeID, last, beg)
 		}
 		run := rbNode.Value.Sums
 		key := rbNode.Value.Key
-		treeID := btrfsprim.CSUM_TREE_OBJECTID
 
 		// Check if we already have it.
 
 		// .Search is more efficient than .Load, because it doesn't load the body (and we don't need the body).
 		if _, ok := o.rebuilt.Search(ctx, treeID, key.Cmp); !ok {
 			// We need to insert it.
-			itemPtr, ok := o.rebuilt.Keys(btrfsprim.CSUM_TREE_OBJECTID).Load(key)
+			itemPtr, ok := o.rebuilt.Keys(treeID).Load(key)
 			if !ok {
 				// This is a panic because if we found it in `o.csums` then it has
 				// to be in some Node, and if we didn't find it from
@@ -555,6 +557,10 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 		}
 
 		beg = run.Addr.Add(run.Size())
+		last = beg
+	}
+	if last < beg {
+		dlog.Errorf(ctx, "augment(tree=%v): could not find csum items for %v-%v", treeID, last, beg)
 	}
 }
 
-- 
cgit v1.2.3-54-g00ecf


From a5b98786cb30d759bebf4765c2a959003a63b499 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 20:58:53 -0700
Subject: rebuildnodes/keyio: Don't export readNode()

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
index fa69c78..916dc53 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
@@ -43,7 +43,7 @@ func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock,
 	}
 }
 
-func (o *Handle) ReadNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node] {
+func (o *Handle) readNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node] {
 	if cached, ok := o.cache.Get(laddr); ok {
 		return cached
 	}
@@ -80,7 +80,7 @@ func (o *Handle) ReadItem(ptr ItemPtr) (item btrfsitem.Item, ok bool) {
 	if o.graph.Nodes[ptr.Node].Level != 0 || ptr.Idx < 0 {
 		return nil, false
 	}
-	items := o.ReadNode(ptr.Node).Data.BodyLeaf
+	items := o.readNode(ptr.Node).Data.BodyLeaf
 	if ptr.Idx >= len(items) {
 		return nil, false
 	}
-- 
cgit v1.2.3-54-g00ecf


From ce1ea3ba56fb7738b4ca567930e357473991e7dd Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 22:32:52 -0700
Subject: rebuildnodes: Rework the wantCsum and wantFileExt algorithms

---
 .../btrfsinspect/rebuildnodes/rebuild.go           | 348 ++++++++++-----------
 1 file changed, 169 insertions(+), 179 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 1daa83d..5badc33 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -15,11 +15,9 @@ import (
 	"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/btrfssum"
 	"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/btrfsprogs/btrfsinspect"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildmappings"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
@@ -45,7 +43,6 @@ type rebuilder struct {
 	rebuilt *btrees.RebuiltTrees
 
 	graph graph.Graph
-	csums containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
 	keyIO *keyio.Handle
 
 	curKey          keyAndTree
@@ -65,19 +62,12 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		return nil, err
 	}
 
-	dlog.Info(ctx, "Indexing checksums...")
-	csums, _ := rebuildmappings.ExtractLogicalSums(ctx, nodeScanResults)
-	if csums == nil {
-		csums = new(containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum])
-	}
-
 	dlog.Info(ctx, "Rebuilding node tree...")
 	keyIO := keyio.NewHandle(fs, *sb, nodeGraph)
 	o := &rebuilder{
 		sb: *sb,
 
 		graph: nodeGraph,
-		csums: *csums,
 		keyIO: keyIO,
 	}
 	o.rebuilt = btrees.NewRebuiltTrees(*sb, nodeGraph, keyIO,
@@ -506,219 +496,219 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 	o.wantAugment(ctx, treeID, wants)
 }
 
-// func implements rebuildCallbacks.
-//
-// interval is [beg, end)
-func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
-	treeID := btrfsprim.CSUM_TREE_OBJECTID
+func (o *rebuilder) _wantRange(
+	ctx context.Context,
+	treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType,
+	sizeFn func(btrfsprim.Key) (uint64, error),
+	beg, end uint64,
+) {
 	if !o.rebuilt.AddTree(ctx, treeID) {
 		o.queue = append(o.queue, o.curKey)
 		return
 	}
 
-	last := beg
-	for beg < end {
-		// Figure out which key we want.
-
-		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("csum for laddr=%v", beg))
-		rbNode := o.csums.Search(func(item btrfsinspect.SysExtentCSum) int {
-			switch {
-			case item.Sums.Addr > beg:
-				return -1
-			case item.Sums.Addr.Add(item.Sums.Size()) <= beg:
-				return 1
-			default:
-				return 0
-			}
-
-		})
-		if rbNode == nil {
-			beg += btrfssum.BlockSize
-			continue
-		} else if last < beg {
-			dlog.Errorf(ctx, "augment(tree=%v): could not find csum items for %v-%v", treeID, last, beg)
-		}
-		run := rbNode.Value.Sums
-		key := rbNode.Value.Key
-
-		// Check if we already have it.
-
-		// .Search is more efficient than .Load, because it doesn't load the body (and we don't need the body).
-		if _, ok := o.rebuilt.Search(ctx, treeID, key.Cmp); !ok {
-			// We need to insert it.
-			itemPtr, ok := o.rebuilt.Keys(treeID).Load(key)
-			if !ok {
-				// This is a panic because if we found it in `o.csums` then it has
-				// to be in some Node, and if we didn't find it from
-				// btrfs.LookupCSum(), then that Node must be an orphan.
-				panic(fmt.Errorf("should not happen: no orphan contains %v", key))
-			}
-			o.wantAugment(ctx, treeID, o.rebuilt.LeafToRoots(ctx, treeID, itemPtr.Node))
-		}
-
-		beg = run.Addr.Add(run.Size())
-		last = beg
-	}
-	if last < beg {
-		dlog.Errorf(ctx, "augment(tree=%v): could not find csum items for %v-%v", treeID, last, beg)
-	}
-}
-
-// wantFileExt implements rebuildCallbacks.
-func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
-	if !o.rebuilt.AddTree(ctx, treeID) {
-		o.queue = append(o.queue, o.curKey)
-		return
-	}
-
-	min := btrfsprim.Key{
-		ObjectID: ino,
-		ItemType: btrfsitem.EXTENT_DATA_KEY,
-		Offset:   0,
+	// Step 1: Build a listing of the runs that we do have.
+	runMin := btrfsprim.Key{
+		ObjectID: objID,
+		ItemType: typ,
+		Offset:   0, // *NOT* `beg`
 	}
-	max := btrfsprim.Key{
-		ObjectID: ino,
-		ItemType: btrfsitem.EXTENT_DATA_KEY,
-		Offset:   uint64(size - 1),
+	runMax := btrfsprim.Key{
+		ObjectID: objID,
+		ItemType: typ,
+		Offset:   end - 1,
 	}
-	extKeys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
+	runKeys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
 		switch {
-		case min.Cmp(key) < 0:
+		case runMin.Cmp(key) < 0:
 			return 1
-		case max.Cmp(key) > 0:
+		case runMax.Cmp(key) > 0:
 			return -1
 		default:
 			return 0
 		}
 	})
 
+	// Step 2: Build a listing of the gaps.
+	//
+	// Start with a gap of the whole range, then subtract each run
+	// from it.
 	type gap struct {
 		// range is [Beg,End)
-		Beg, End int64
+		Beg, End uint64
 	}
-	gaps := &containers.RBTree[containers.NativeOrdered[int64], gap]{
-		KeyFn: func(gap gap) containers.NativeOrdered[int64] {
-			return containers.NativeOrdered[int64]{Val: gap.Beg}
+	gaps := &containers.RBTree[containers.NativeOrdered[uint64], gap]{
+		KeyFn: func(gap gap) containers.NativeOrdered[uint64] {
+			return containers.NativeOrdered[uint64]{Val: gap.Beg}
 		},
 	}
 	gaps.Insert(gap{
-		Beg: 0,
-		End: size,
+		Beg: beg,
+		End: end,
 	})
-	for _, extKey := range extKeys {
-		extBody, ok := o.rebuilt.Load(ctx, treeID, extKey)
-		if !ok {
-			o.ioErr(ctx, fmt.Errorf("could not look up already-inserted item: tree=%v key=%v",
-				treeID, extKey))
+	for _, runKey := range runKeys {
+		runSize, err := sizeFn(runKey)
+		if err != nil {
+			o.fsErr(ctx, fmt.Errorf("tree=%v key=%v: %w", treeID, runKey, err))
+		}
+		if runSize == 0 {
+			continue
+		}
+		runBeg := runKey.Offset
+		runEnd := runBeg + runSize
+		if runEnd <= beg {
+			continue
 		}
-		switch extBody := extBody.(type) {
-		case btrfsitem.FileExtent:
-			extBeg := int64(extKey.Offset)
-			extSize, err := extBody.Size()
-			if err != nil {
-				o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, extKey, err))
-				continue
+		overlappingGaps := gaps.SearchRange(func(gap gap) int {
+			switch {
+			case gap.End <= runBeg:
+				return 1
+			case runEnd <= gap.Beg:
+				return -1
+			default:
+				return 0
 			}
-			extEnd := extBeg + extSize
-			overlappingGaps := gaps.SearchRange(func(gap gap) int {
-				switch {
-				case gap.End <= extBeg:
-					return 1
-				case extEnd <= gap.Beg:
-					return -1
-				default:
-					return 0
-				}
+		})
+		if len(overlappingGaps) == 0 {
+			continue
+		}
+		gapsBeg := overlappingGaps[0].Beg
+		gapsEnd := overlappingGaps[len(overlappingGaps)-1].End
+		for _, gap := range overlappingGaps {
+			gaps.Delete(containers.NativeOrdered[uint64]{Val: gap.Beg})
+		}
+		if gapsBeg < runBeg {
+			gaps.Insert(gap{
+				Beg: gapsBeg,
+				End: runBeg,
+			})
+		}
+		if gapsEnd > runEnd {
+			gaps.Insert(gap{
+				Beg: runEnd,
+				End: gapsEnd,
 			})
-			if len(overlappingGaps) == 0 {
-				continue
-			}
-			beg := overlappingGaps[0].Beg
-			end := overlappingGaps[len(overlappingGaps)-1].End
-			for _, gap := range overlappingGaps {
-				gaps.Delete(containers.NativeOrdered[int64]{Val: gap.Beg})
-			}
-			if beg < extBeg {
-				gaps.Insert(gap{
-					Beg: beg,
-					End: extBeg,
-				})
-			}
-			if end > extEnd {
-				gaps.Insert(gap{
-					Beg: extEnd,
-					End: end,
-				})
-			}
-		case btrfsitem.Error:
-			o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, extKey, extBody.Err))
-		default:
-			// This is a panic because the item decoder should not emit EXTENT_DATA
-			// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
-			// this code also being updated.
-			panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", extBody))
 		}
 	}
+
+	// Step 2: Fill each gap.
 	_ = gaps.Walk(func(rbNode *containers.RBNode[gap]) error {
 		gap := rbNode.Value
-		min := btrfsprim.Key{
-			ObjectID: ino,
-			ItemType: btrfsitem.EXTENT_DATA_KEY,
-			Offset:   0,
+		last := gap.Beg
+		runMin := btrfsprim.Key{
+			ObjectID: objID,
+			ItemType: typ,
+			Offset:   0, // *NOT* `gap.Beg`
 		}
-		max := btrfsprim.Key{
-			ObjectID: ino,
-			ItemType: btrfsitem.EXTENT_DATA_KEY,
-			Offset:   uint64(gap.End - 1),
+		runMax := btrfsprim.Key{
+			ObjectID: objID,
+			ItemType: typ,
+			Offset:   gap.End - 1,
 		}
-		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("file extent for tree=%v inode=%v bytes [%v, %v)", treeID, ino, gap.Beg, gap.End))
-		wants := make(containers.Set[btrfsvol.LogicalAddr])
 		o.rebuilt.Keys(treeID).Subrange(
-			func(k btrfsprim.Key, _ keyio.ItemPtr) int {
+			func(key btrfsprim.Key, _ keyio.ItemPtr) int {
 				switch {
-				case min.Cmp(k) < 0:
+				case runMin.Cmp(key) < 0:
 					return 1
-				case max.Cmp(k) > 0:
+				case runMax.Cmp(key) > 0:
 					return -1
 				default:
 					return 0
 				}
 			},
 			func(k btrfsprim.Key, v keyio.ItemPtr) bool {
-				itemBody, ok := o.keyIO.ReadItem(v)
-				if !ok {
-					o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
+				runSize, err := sizeFn(k)
+				if err != nil {
+					o.fsErr(ctx, fmt.Errorf("tree=%v key=%v: %w", treeID, k, err))
+					return true
 				}
-				switch itemBody := itemBody.(type) {
-				case btrfsitem.FileExtent:
-					itemBeg := int64(k.Offset)
-					itemSize, err := itemBody.Size()
-					if err != nil {
-						o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, k, err))
-						break
-					}
-					itemEnd := itemBeg + itemSize
-					// We're being greedy and "wanting" any extent that has any overlap with
-					// the gap.  But maybe instead we sould only want extents that are
-					// *entirely* within the gap.  I'll have to run it on real filesystems
-					// to see what works better.
-					//
-					// TODO(lukeshu): Re-evaluate whether being greedy here is the right
-					// thing.
-					if itemEnd > gap.Beg && itemBeg < gap.End {
-						wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
-					}
-				case btrfsitem.Error:
-					o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, k, itemBody.Err))
-				default:
-					// This is a panic because the item decoder should not emit EXTENT_DATA
-					// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
-					// this code also being updated.
-					panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", itemBody))
+				if runSize == 0 {
+					return true
+				}
+				runBeg := k.Offset
+				runEnd := runBeg + runSize
+				if runEnd <= gap.Beg {
+					return true
+				}
+
+				// TODO: This is dumb and greedy.
+				if last < runBeg {
+					// log an error
+					o.wantAugment(dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v, %v, %v-%v}",
+						treeID, objID, typ, last, runBeg)),
+						treeID, nil)
 				}
+				o.wantAugment(dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v, %v, %v-%v}",
+					treeID, objID, typ, gap.Beg, gap.End)),
+					treeID, o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
+				last = runEnd
+
 				return true
 			})
-		o.wantAugment(ctx, treeID, wants)
+		if last < gap.End {
+			// log an error
+			o.wantAugment(dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v, %v, %v-%v}",
+				treeID, objID, typ, last, gap.End)),
+				treeID, nil)
+		}
 		return nil
 	})
 }
+
+// func implements rebuildCallbacks.
+//
+// interval is [beg, end)
+func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
+	const treeID = btrfsprim.CSUM_TREE_OBJECTID
+	o._wantRange(ctx, treeID, btrfsprim.EXTENT_CSUM_OBJECTID, btrfsprim.EXTENT_CSUM_KEY,
+		func(key btrfsprim.Key) (uint64, error) {
+			ptr, ok := o.rebuilt.Keys(treeID).Load(key)
+			if !ok {
+				panic(fmt.Errorf("should not happen: could not load key: %v", key))
+			}
+			body, ok := o.keyIO.ReadItem(ptr)
+			if !ok {
+				panic(fmt.Errorf("should not happen: could not read item: %v", key))
+			}
+			switch body := body.(type) {
+			case btrfsitem.ExtentCSum:
+				return uint64(body.Size()), nil
+			case btrfsitem.Error:
+				return 0, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, key, body.Err)
+			default:
+				// This is a panic because the item decoder should not emit EXTENT_CSUM
+				// items as anything but btrfsitem.ExtentCSum or btrfsitem.Error without
+				// this code also being updated.
+				panic(fmt.Errorf("should not happen: EXTENT_CSUM item has unexpected type: %T", body))
+			}
+		},
+		uint64(beg), uint64(end))
+}
+
+// wantFileExt implements rebuildCallbacks.
+func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
+	o._wantRange(ctx, treeID, ino, btrfsprim.EXTENT_DATA_KEY,
+		func(key btrfsprim.Key) (uint64, error) {
+			ptr, ok := o.rebuilt.Keys(treeID).Load(key)
+			if !ok {
+				panic(fmt.Errorf("should not happen: could not load key: %v", key))
+			}
+			body, ok := o.keyIO.ReadItem(ptr)
+			if !ok {
+				panic(fmt.Errorf("should not happen: could not read item: %v", key))
+			}
+			switch body := body.(type) {
+			case btrfsitem.FileExtent:
+				size, err := body.Size()
+				return uint64(size), err
+			case btrfsitem.Error:
+				return 0, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, key, body.Err)
+			default:
+				// This is a panic because the item decoder should not emit EXTENT_DATA
+				// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
+				// this code also being updated.
+				panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", body))
+			}
+		},
+		0, uint64(size))
+}
-- 
cgit v1.2.3-54-g00ecf


From f72d881090d5a806cbbc5b3c3fc66683a28474fd Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 22:34:04 -0700
Subject: Revert "scandevices: Include a Key in SysExtentCSum"

This reverts commit a71e1969a1b24500ce9885c4a3f1aeee40c421a8.
---
 lib/btrfsprogs/btrfsinspect/scandevices.go | 2 --
 1 file changed, 2 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/scandevices.go b/lib/btrfsprogs/btrfsinspect/scandevices.go
index 4834826..b6c929b 100644
--- a/lib/btrfsprogs/btrfsinspect/scandevices.go
+++ b/lib/btrfsprogs/btrfsinspect/scandevices.go
@@ -75,7 +75,6 @@ type SysDevExtent struct {
 }
 
 type SysExtentCSum struct {
-	Key        btrfsprim.Key
 	Generation btrfsprim.Generation
 	Sums       btrfsitem.ExtentCSum
 }
@@ -227,7 +226,6 @@ func ScanOneDevice(ctx context.Context, dev *btrfs.Device, sb btrfstree.Superblo
 						//dlog.Tracef(ctx, "... dev[%q] node@%v: item %v: found csums",
 						//	dev.Name(), nodeRef.Addr, i)
 						result.FoundExtentCSums = append(result.FoundExtentCSums, SysExtentCSum{
-							Key:        item.Key,
 							Generation: nodeRef.Data.Head.Generation,
 							Sums:       sums,
 						})
-- 
cgit v1.2.3-54-g00ecf


From d1173de57640e67f4bdc4a0cdc923ff8f509f7df Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 22:34:16 -0700
Subject: Revert "rebuildmappings: Publicly expose the logicalsums RBTree"

This reverts commit 7d77b89f355714da88bb1bbde5e343747bfb6a81.
---
 .../btrfsinspect/rebuildmappings/logicalsums.go         | 17 ++++-------------
 .../btrfsinspect/rebuildmappings/rebuildmappings.go     |  2 +-
 2 files changed, 5 insertions(+), 14 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go b/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go
index 9b2da93..537a970 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go
@@ -20,7 +20,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
-func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevicesResult) (addrspace *containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum], sumSize int) {
+func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevicesResult) btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr] {
 	var records []btrfsinspect.SysExtentCSum
 	for _, devResults := range scanResults {
 		records = append(records, devResults.FoundExtentCSums...)
@@ -38,9 +38,9 @@ func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevice
 		}
 	})
 	if len(records) == 0 {
-		return nil, 0
+		return btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr]{}
 	}
-	sumSize = records[0].Sums.ChecksumSize
+	sumSize := records[0].Sums.ChecksumSize
 
 	// Now build them in to a coherent address space.
 	//
@@ -53,7 +53,7 @@ func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevice
 	// "AAAAAAA" shouldn't be present, and if we just discard "BBBBBBBB"
 	// because it conflicts with "CCCCCCC", then we would erroneously
 	// include "AAAAAAA".
-	addrspace = &containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]{
+	addrspace := &containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]{
 		KeyFn: func(item btrfsinspect.SysExtentCSum) containers.NativeOrdered[btrfsvol.LogicalAddr] {
 			return containers.NativeOrdered[btrfsvol.LogicalAddr]{Val: item.Sums.Addr}
 		},
@@ -148,15 +148,6 @@ func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevice
 		}
 	}
 
-	return addrspace, sumSize
-}
-
-func ExtractAndFlattenLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevicesResult) btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr] {
-	addrspace, sumSize := ExtractLogicalSums(ctx, scanResults)
-	if addrspace == nil {
-		return btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr]{}
-	}
-
 	// Now flatten that RBTree in to a SumRunWithGaps.
 	var flattened btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr]
 	var curAddr btrfsvol.LogicalAddr
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildmappings/rebuildmappings.go b/lib/btrfsprogs/btrfsinspect/rebuildmappings/rebuildmappings.go
index c391053..f1959b4 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildmappings/rebuildmappings.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildmappings/rebuildmappings.go
@@ -145,7 +145,7 @@ func RebuildMappings(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect
 
 	dlog.Infof(ctx, "5/6: Searching for %d block groups in checksum map (exact)...", len(bgs))
 	physicalSums := ExtractPhysicalSums(scanResults)
-	logicalSums := ExtractAndFlattenLogicalSums(ctx, scanResults)
+	logicalSums := ExtractLogicalSums(ctx, scanResults)
 	if err := matchBlockGroupSums(ctx, fs, bgs, physicalSums, logicalSums); err != nil {
 		return err
 	}
-- 
cgit v1.2.3-54-g00ecf


From bc06b340b71a07a0bb500e1bb7e81ece19f1a9c5 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sat, 24 Dec 2022 04:38:24 -0700
Subject: rebuildnodes: Track item sizes at startup, to avoid extra i/o

---
 .../btrfsinspect/rebuildnodes/keyio/keyio.go       | 46 ++++++++++++++++-
 .../btrfsinspect/rebuildnodes/rebuild.go           | 59 +++++-----------------
 lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go   | 15 ++++--
 3 files changed, 67 insertions(+), 53 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
index 916dc53..d92495e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
@@ -25,24 +25,66 @@ func (ptr ItemPtr) String() string {
 	return fmt.Sprintf("node@%v[%v]", ptr.Node, ptr.Idx)
 }
 
+type SizeAndErr struct {
+	Size uint64
+	Err  error
+}
+
 type Handle struct {
 	rawFile diskio.File[btrfsvol.LogicalAddr]
 	sb      btrfstree.Superblock
 	graph   graph.Graph
 
+	Sizes map[ItemPtr]SizeAndErr
+
 	cache *containers.LRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]]
 }
 
-func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph graph.Graph) *Handle {
+func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock) *Handle {
 	return &Handle{
 		rawFile: file,
 		sb:      sb,
-		graph:   graph,
+
+		Sizes: make(map[ItemPtr]SizeAndErr),
 
 		cache: containers.NewLRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]](8),
 	}
 }
 
+func (o *Handle) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
+	for i, item := range nodeRef.Data.BodyLeaf {
+		ptr := ItemPtr{
+			Node: nodeRef.Addr,
+			Idx:  i,
+		}
+		switch itemBody := item.Body.(type) {
+		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.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, nodeRef.Data.Head.Owner, item.Key, itemBody.Err),
+				}
+			}
+		}
+	}
+}
+
+func (o *Handle) SetGraph(graph graph.Graph) {
+	o.graph = graph
+}
+
 func (o *Handle) readNode(laddr btrfsvol.LogicalAddr) *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node] {
 	if cached, ok := o.cache.Get(laddr); ok {
 		return cached
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 5badc33..66eaf1a 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -51,7 +51,7 @@ type rebuilder struct {
 }
 
 func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], error) {
-	nodeGraph, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
+	nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
 	if err != nil {
 		return nil, err
 	}
@@ -63,7 +63,6 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Rebuilding node tree...")
-	keyIO := keyio.NewHandle(fs, *sb, nodeGraph)
 	o := &rebuilder{
 		sb: *sb,
 
@@ -499,7 +498,6 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 func (o *rebuilder) _wantRange(
 	ctx context.Context,
 	treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType,
-	sizeFn func(btrfsprim.Key) (uint64, error),
 	beg, end uint64,
 ) {
 	if !o.rebuilt.AddTree(ctx, treeID) {
@@ -507,6 +505,18 @@ func (o *rebuilder) _wantRange(
 		return
 	}
 
+	sizeFn := func(key btrfsprim.Key) (uint64, error) {
+		ptr, ok := o.rebuilt.Keys(treeID).Load(key)
+		if !ok {
+			panic(fmt.Errorf("should not happen: could not load key: %v", key))
+		}
+		sizeAndErr, ok := o.keyIO.Sizes[ptr]
+		if !ok {
+			panic(fmt.Errorf("should not happen: %v item did not have a size recorded", typ))
+		}
+		return sizeAndErr.Size, sizeAndErr.Err
+	}
+
 	// Step 1: Build a listing of the runs that we do have.
 	runMin := btrfsprim.Key{
 		ObjectID: objID,
@@ -661,54 +671,11 @@ func (o *rebuilder) _wantRange(
 func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
 	const treeID = btrfsprim.CSUM_TREE_OBJECTID
 	o._wantRange(ctx, treeID, btrfsprim.EXTENT_CSUM_OBJECTID, btrfsprim.EXTENT_CSUM_KEY,
-		func(key btrfsprim.Key) (uint64, error) {
-			ptr, ok := o.rebuilt.Keys(treeID).Load(key)
-			if !ok {
-				panic(fmt.Errorf("should not happen: could not load key: %v", key))
-			}
-			body, ok := o.keyIO.ReadItem(ptr)
-			if !ok {
-				panic(fmt.Errorf("should not happen: could not read item: %v", key))
-			}
-			switch body := body.(type) {
-			case btrfsitem.ExtentCSum:
-				return uint64(body.Size()), nil
-			case btrfsitem.Error:
-				return 0, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, key, body.Err)
-			default:
-				// This is a panic because the item decoder should not emit EXTENT_CSUM
-				// items as anything but btrfsitem.ExtentCSum or btrfsitem.Error without
-				// this code also being updated.
-				panic(fmt.Errorf("should not happen: EXTENT_CSUM item has unexpected type: %T", body))
-			}
-		},
 		uint64(beg), uint64(end))
 }
 
 // wantFileExt implements rebuildCallbacks.
 func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
 	o._wantRange(ctx, treeID, ino, btrfsprim.EXTENT_DATA_KEY,
-		func(key btrfsprim.Key) (uint64, error) {
-			ptr, ok := o.rebuilt.Keys(treeID).Load(key)
-			if !ok {
-				panic(fmt.Errorf("should not happen: could not load key: %v", key))
-			}
-			body, ok := o.keyIO.ReadItem(ptr)
-			if !ok {
-				panic(fmt.Errorf("should not happen: could not read item: %v", key))
-			}
-			switch body := body.(type) {
-			case btrfsitem.FileExtent:
-				size, err := body.Size()
-				return uint64(size), err
-			case btrfsitem.Error:
-				return 0, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, key, body.Err)
-			default:
-				// This is a panic because the item decoder should not emit EXTENT_DATA
-				// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
-				// this code also being updated.
-				panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", body))
-			}
-		},
 		0, uint64(size))
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index bcf2f5b..5c2d0fd 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -16,6 +16,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
@@ -30,12 +31,12 @@ func (s scanStats) String() string {
 		s.N, s.D)
 }
 
-func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (graph.Graph, error) {
+func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (graph.Graph, *keyio.Handle, error) {
 	dlog.Infof(ctx, "Reading node data from FS...")
 
 	sb, err := fs.Superblock()
 	if err != nil {
-		return graph.Graph{}, err
+		return graph.Graph{}, nil, err
 	}
 
 	total := countNodes(scanResults)
@@ -46,6 +47,7 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	}
 
 	nodeGraph := graph.New(*sb)
+	keyIO := keyio.NewHandle(fs, *sb)
 
 	progress(done, total)
 	for _, devResults := range scanResults {
@@ -54,10 +56,11 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 			})
 			if err != nil {
-				return graph.Graph{}, err
+				return graph.Graph{}, nil, err
 			}
 
 			nodeGraph.InsertNode(nodeRef)
+			keyIO.InsertNode(nodeRef)
 
 			done++
 			progress(done, total)
@@ -72,10 +75,12 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	progressWriter = textui.NewProgress[scanStats](ctx, dlog.LogLevelInfo, 1*time.Second)
 	dlog.Infof(ctx, "Checking keypointers for dead-ends...")
 	if err := nodeGraph.FinalCheck(fs, *sb, progress); err != nil {
-		return graph.Graph{}, err
+		return graph.Graph{}, nil, err
 	}
 	progressWriter.Done()
 	dlog.Info(ctx, "... done checking keypointers")
 
-	return *nodeGraph, nil
+	keyIO.SetGraph(*nodeGraph)
+
+	return *nodeGraph, keyIO, nil
 }
-- 
cgit v1.2.3-54-g00ecf