rebuildnodes/btrees: Refactor to split the forrest from the trees

4 files changed, 498 insertions, 508 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
new file mode 100644
index 0000000..791e646
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -0,0 +1,173 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrees
+
+import (
+	"context"
+
+	"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"
+	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/slices"
+)
+
+// RebuiltForrest is an abstraction for rebuilding and accessing
+// potentially broken btrees.
+//
+// It is conceptually a btrfstree.TreeOperator, and adds similar
+// broken-tree handling to btrfsutil.BrokenForrest.  However, the API
+// is different thant btrfstree.TreeOperator, and is much more
+// efficient than btrfsutil.BrokenForrest.
+//
+// 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
+//
+// Additionally, it provides some functionality that
+// btrfsutil.BrokenForrest 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 RebuiltForrest is invalid; it must be initialized with
+// NewRebuiltForrest().
+type RebuiltForrest struct {
+	// static
+	sb    btrfstree.Superblock
+	graph pkggraph.Graph
+	keyIO *keyio.Handle
+
+	// static callbacks
+	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)
+
+	// mutable
+	trees map[btrfsprim.ObjID]*RebuiltTree
+}
+
+// NewRebuiltForrest returns a new RebuiltForrest instance.  All of
+// the callbacks must be non-nil.
+func NewRebuiltForrest(
+	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),
+) *RebuiltForrest {
+	return &RebuiltForrest{
+		sb:    sb,
+		graph: graph,
+		keyIO: keyIO,
+
+		cbAddedItem:  cbAddedItem,
+		cbLookupRoot: cbLookupRoot,
+		cbLookupUUID: cbLookupUUID,
+
+		trees: make(map[btrfsprim.ObjID]*RebuiltTree),
+	}
+}
+
+// Tree returns a given tree, initializing it if nescessary.  If it is
+// unable to initialize the tree, then nil is returned, and nothing is
+// done to the forrest.
+//
+// The tree is initialized with the normal root node of the tree.
+func (ts *RebuiltForrest) Tree(ctx context.Context, treeID btrfsprim.ObjID) *RebuiltTree {
+	if !ts.addTree(ctx, treeID, nil) {
+		return nil
+	}
+	return ts.trees[treeID]
+}
+
+func (ts *RebuiltForrest) 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
+	}
+	stack = append(stack, treeID)
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree", stack)
+	dlog.Info(ctx, "adding tree...")
+
+	tree := &RebuiltTree{
+		ID:      treeID,
+		Roots:   make(containers.Set[btrfsvol.LogicalAddr]),
+		Leafs:   make(containers.Set[btrfsvol.LogicalAddr]),
+		forrest: ts,
+	}
+	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:
+		if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
+			return false
+		}
+		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.UUID_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]
+		}
+	}
+	tree.indexLeafs(ctx)
+
+	ts.trees[treeID] = tree
+	if root != 0 {
+		tree.AddRoot(ctx, root)
+	}
+
+	return true
+}
+
+// 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
+// RebuiltForrest's internal set!
+func (ts *RebuiltForrest) 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
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
deleted file mode 100644
index ffe225a..0000000
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ /dev/null
@@ -1,482 +0,0 @@
-// Copyright (C) 2022-2023  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"
-	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/maps"
-	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
-	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
-)
-
-type rebuiltTree struct {
-	// static
-	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]
-	keys        containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
-
-	// mutable
-	Roots containers.Set[btrfsvol.LogicalAddr]
-	Leafs containers.Set[btrfsvol.LogicalAddr]
-	Items containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
-}
-
-// 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, gen btrfsprim.Generation) bool {
-	for {
-		if owner == tree.ID {
-			return true
-		}
-		if tree.Parent == nil || gen >= tree.ParentGen {
-			return false
-		}
-		tree = tree.Parent
-	}
-}
-
-// 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++
-	}
-}
-
-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.
-//
-// 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
-//
-// 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 {
-	// static
-	sb    btrfstree.Superblock
-	graph pkggraph.Graph
-	keyIO *keyio.Handle
-
-	// static callbacks
-	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)
-
-	// mutable
-	trees map[btrfsprim.ObjID]*rebuiltTree
-}
-
-// 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,
-	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),
-) *RebuiltTrees {
-	return &RebuiltTrees{
-		sb:    sb,
-		graph: graph,
-		keyIO: keyIO,
-
-		cbAddedItem:  cbAddedItem,
-		cbLookupRoot: cbLookupRoot,
-		cbLookupUUID: cbLookupUUID,
-
-		trees: make(map[btrfsprim.ObjID]*rebuiltTree),
-	}
-}
-
-type rootStats struct {
-	Leafs         textui.Portion[int]
-	AddedItems    int
-	ReplacedItems int
-}
-
-func (s rootStats) String() string {
-	return textui.Sprintf("%v (added %v items, replaced %v items)",
-		s.Leafs, s.AddedItems, s.ReplacedItems)
-}
-
-// 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) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-root", fmt.Sprintf("tree=%v rootNode=%v", treeID, rootNode))
-	tree := ts.trees[treeID]
-	tree.Roots.Insert(rootNode)
-
-	var stats rootStats
-	stats.Leafs.D = len(tree.leafToRoots)
-	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
-		stats.Leafs.N = i
-		progressWriter.Set(stats)
-		if tree.Leafs.Has(leaf) || !tree.leafToRoots[leaf].Has(rootNode) {
-			continue
-		}
-		tree.Leafs.Insert(leaf)
-		for j, itemKey := range ts.graph.Nodes[leaf].Items {
-			newPtr := keyio.ItemPtr{
-				Node: leaf,
-				Idx:  j,
-			}
-			if oldPtr, exists := tree.Items.Load(itemKey); !exists {
-				tree.Items.Store(itemKey, newPtr)
-				stats.AddedItems++
-			} else if tree.shouldReplace(ts.graph, oldPtr.Node, newPtr.Node) {
-				tree.Items.Store(itemKey, newPtr)
-				stats.ReplacedItems++
-			}
-			ts.cbAddedItem(ctx, treeID, itemKey)
-			progressWriter.Set(stats)
-		}
-	}
-	stats.Leafs.N = len(tree.leafToRoots)
-	progressWriter.Set(stats)
-	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
-	}
-	stack = append(stack, treeID)
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree", stack)
-	dlog.Info(ctx, "adding tree...")
-
-	tree := &rebuiltTree{
-		ID:    treeID,
-		Roots: make(containers.Set[btrfsvol.LogicalAddr]),
-		Leafs: 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:
-		if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
-			return false
-		}
-		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.UUID_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]
-		}
-	}
-	tree.indexLeafs(ctx, ts.graph)
-
-	ts.trees[treeID] = tree
-	if root != 0 {
-		ts.AddRoot(ctx, treeID, root)
-	}
-
-	return true
-}
-
-func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree.substep", "index-nodes")
-
-	nodeToRoots := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
-
-	var stats textui.Portion[int]
-	stats.D = len(graph.Nodes)
-	progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	progress := func() {
-		stats.N = len(nodeToRoots)
-		progressWriter.Set(stats)
-	}
-
-	progress()
-	for _, node := range maps.SortedKeys(graph.Nodes) {
-		tree.indexNode(ctx, 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(ctx context.Context, graph pkggraph.Graph, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) {
-	defer progress()
-	if err := ctx.Err(); err != nil {
-		return
-	}
-	if _, done := index[node]; done {
-		return
-	}
-	if slices.Contains(node, stack) {
-		// This is a panic because graph.FinalCheck() should
-		// have already checked for loops.
-		panic("loop")
-	}
-	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, graph.Nodes[kp.FromNode].Generation)
-	})
-	for _, kp := range kps {
-		tree.indexNode(ctx, graph, kp.FromNode, index, progress, stack)
-		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,
-			})
-		}
-	}
-}
-
-// Resolve a key to a keyio.ItemPtr.
-//
-// It is not nescessary to call AddTree for that tree first; Resolve will
-// call it for you.
-func (ts *RebuiltTrees) Resolve(ctx context.Context, treeID btrfsprim.ObjID, key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
-	if !ts.AddTree(ctx, treeID) {
-		return keyio.ItemPtr{}, false
-	}
-	return ts.trees[treeID].Items.Load(key)
-}
-
-// 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) {
-	ptr, ok := ts.Resolve(ctx, treeID, key)
-	if !ok {
-		return nil, false
-	}
-	return ts.keyIO.ReadItem(ctx, ptr)
-}
-
-// 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, _ keyio.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, _ keyio.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
-}
-
-// 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
-	}
-	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) {
-			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
-	}
-	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.
-//
-// 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.
-//
-// 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
-}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
new file mode 100644
index 0000000..1fc1f52
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -0,0 +1,299 @@
+// 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/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/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
+)
+
+type RebuiltTree struct {
+	// static
+	ID        btrfsprim.ObjID
+	UUID      btrfsprim.UUID
+	Parent    *RebuiltTree
+	ParentGen btrfsprim.Generation // offset of this tree's root item
+	forrest   *RebuiltForrest
+
+	// 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]
+	Leafs containers.Set[btrfsvol.LogicalAddr]
+	Items containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
+}
+
+// initializaton (called by `RebuiltForrest.Tree()`) ///////////////////////////////////////////////////////////////////
+
+func (tree *RebuiltTree) indexLeafs(ctx context.Context) {
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree.substep", "index-nodes")
+
+	nodeToRoots := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+
+	var stats textui.Portion[int]
+	stats.D = len(tree.forrest.graph.Nodes)
+	progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	progress := func() {
+		stats.N = len(nodeToRoots)
+		progressWriter.Set(stats)
+	}
+
+	progress()
+	for _, node := range maps.SortedKeys(tree.forrest.graph.Nodes) {
+		tree.indexNode(ctx, node, nodeToRoots, progress, nil)
+	}
+	progressWriter.Done()
+
+	tree.leafToRoots = make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+	for node, roots := range nodeToRoots {
+		if tree.forrest.graph.Nodes[node].Level == 0 && len(roots) > 0 {
+			tree.leafToRoots[node] = roots
+		}
+	}
+}
+
+func (tree *RebuiltTree) indexNode(ctx context.Context, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) {
+	defer progress()
+	if err := ctx.Err(); err != nil {
+		return
+	}
+	if _, done := index[node]; done {
+		return
+	}
+	if slices.Contains(node, stack) {
+		// This is a panic because tree.forrest.graph.FinalCheck() should
+		// have already checked for loops.
+		panic("loop")
+	}
+	if !tree.isOwnerOK(tree.forrest.graph.Nodes[node].Owner, tree.forrest.graph.Nodes[node].Generation) {
+		index[node] = nil
+		return
+	}
+
+	// tree.leafToRoots
+	stack = append(stack, node)
+	var roots containers.Set[btrfsvol.LogicalAddr]
+	kps := slices.RemoveAllFunc(tree.forrest.graph.EdgesTo[node], func(kp *pkggraph.Edge) bool {
+		return !tree.isOwnerOK(tree.forrest.graph.Nodes[kp.FromNode].Owner, tree.forrest.graph.Nodes[kp.FromNode].Generation)
+	})
+	for _, kp := range kps {
+		tree.indexNode(ctx, kp.FromNode, index, progress, stack)
+		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 tree.forrest.graph.Nodes[node].Items {
+		if oldPtr, ok := tree.keys.Load(key); !ok || tree.shouldReplace(oldPtr.Node, node) {
+			tree.keys.Store(key, keyio.ItemPtr{
+				Node: node,
+				Idx:  i,
+			})
+		}
+	}
+}
+
+// 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, gen btrfsprim.Generation) bool {
+	for {
+		if owner == tree.ID {
+			return true
+		}
+		if tree.Parent == nil || gen >= tree.ParentGen {
+			return false
+		}
+		tree = tree.Parent
+	}
+}
+
+// .AddRoot() //////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+type rootStats struct {
+	Leafs         textui.Portion[int]
+	AddedItems    int
+	ReplacedItems int
+}
+
+func (s rootStats) String() string {
+	return textui.Sprintf("%v (added %v items, replaced %v items)",
+		s.Leafs, s.AddedItems, s.ReplacedItems)
+}
+
+// AddRoot adds an additional root node to the tree.  It is useful to
+// call .AddRoot() to re-attach part of the tree that has been broken
+// off.
+func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalAddr) {
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
+	tree.Roots.Insert(rootNode)
+
+	var stats rootStats
+	stats.Leafs.D = len(tree.leafToRoots)
+	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
+		stats.Leafs.N = i
+		progressWriter.Set(stats)
+		if tree.Leafs.Has(leaf) || !tree.leafToRoots[leaf].Has(rootNode) {
+			continue
+		}
+		tree.Leafs.Insert(leaf)
+		for j, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
+			newPtr := keyio.ItemPtr{
+				Node: leaf,
+				Idx:  j,
+			}
+			if oldPtr, exists := tree.Items.Load(itemKey); !exists {
+				tree.Items.Store(itemKey, newPtr)
+				stats.AddedItems++
+			} else if tree.shouldReplace(oldPtr.Node, newPtr.Node) {
+				tree.Items.Store(itemKey, newPtr)
+				stats.ReplacedItems++
+			}
+			tree.forrest.cbAddedItem(ctx, tree.ID, itemKey)
+			progressWriter.Set(stats)
+		}
+	}
+	stats.Leafs.N = len(tree.leafToRoots)
+	progressWriter.Set(stats)
+	progressWriter.Done()
+}
+
+func (tree *RebuiltTree) shouldReplace(oldNode, newNode btrfsvol.LogicalAddr) bool {
+	oldDist, _ := tree.COWDistance(tree.forrest.graph.Nodes[oldNode].Owner)
+	newDist, _ := tree.COWDistance(tree.forrest.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 := tree.forrest.graph.Nodes[oldNode].Generation
+		newGen := tree.forrest.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, tree.forrest.graph.Nodes[oldNode],
+				newNode, tree.forrest.graph.Nodes[newNode]))
+		}
+	}
+}
+
+// main public API /////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// 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++
+	}
+}
+
+// Resolve a key to a keyio.ItemPtr.
+func (tree *RebuiltTree) Resolve(key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
+	return tree.Items.Load(key)
+}
+
+// Load reads an item from a tree.
+func (tree *RebuiltTree) Load(ctx context.Context, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
+	ptr, ok := tree.Resolve(key)
+	if !ok {
+		return nil, false
+	}
+	return tree.forrest.keyIO.ReadItem(ctx, ptr)
+}
+
+// Search searches for an item from a tree.
+func (tree *RebuiltTree) Search(fn func(btrfsprim.Key) int) (key btrfsprim.Key, ok bool) {
+	k, _, ok := tree.Items.Search(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
+		return fn(k)
+	})
+	return k, ok
+}
+
+// Search searches for a range of items from a tree.
+func (tree *RebuiltTree) SearchAll(fn func(btrfsprim.Key) int) []btrfsprim.Key {
+	kvs := tree.Items.SearchAll(func(k btrfsprim.Key, _ keyio.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
+}
+
+// LeafToRoots returns the list of potential roots (to pass to
+// .AddRoot) that include a given leaf-node.
+func (tree *RebuiltTree) LeafToRoots(leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
+	if tree.forrest.graph.Nodes[leaf].Level != 0 {
+		panic(fmt.Errorf("should not happen: (tree=%v).LeafToRoots(leaf=%v): not a leaf",
+			tree.ID, leaf))
+	}
+	ret := make(containers.Set[btrfsvol.LogicalAddr])
+	for root := range tree.leafToRoots[leaf] {
+		if tree.Roots.Has(root) {
+			panic(fmt.Errorf("should not happen: (tree=%v).LeafToRoots(leaf=%v): tree contains root=%v but not leaf",
+				tree.ID, leaf, root))
+		}
+		ret.Insert(root)
+	}
+	if len(ret) == 0 {
+		return nil
+	}
+	return ret
+}
+
+// Keys returns a map of all keys in node that would be valid in this tree.
+//
+// Do not mutate the returned map; it is a pointer to the
+// RebuiltTree's internal map!
+func (tree *RebuiltTree) Keys() *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
+	return &tree.keys
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 26f2a44..b4feacf 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -48,7 +48,7 @@ type rebuilder struct {
 	graph graph.Graph
 	keyIO *keyio.Handle
 
-	rebuilt *btrees.RebuiltTrees
+	rebuilt *btrees.RebuiltForrest
 
 	curKey       keyAndTree
 	treeQueue    containers.Set[btrfsprim.ObjID]
@@ -73,7 +73,7 @@ func NewRebuilder(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		graph: nodeGraph,
 		keyIO: keyIO,
 	}
-	o.rebuilt = btrees.NewRebuiltTrees(sb, nodeGraph, keyIO,
+	o.rebuilt = btrees.NewRebuiltForrest(sb, nodeGraph, keyIO,
 		o.cbAddedItem, o.cbLookupRoot, o.cbLookupUUID)
 	return o, nil
 }
@@ -116,7 +116,7 @@ func (o *rebuilder) Rebuild(_ctx context.Context) error {
 			if err := _ctx.Err(); err != nil {
 				return err
 			}
-			o.rebuilt.AddTree(stepCtx, treeID)
+			o.rebuilt.Tree(stepCtx, treeID)
 		}
 
 		// Handle items in the queue (drain o.itemQueue, fill o.augmentQueue and o.treeQueue)
@@ -139,7 +139,7 @@ func (o *rebuilder) Rebuild(_ctx context.Context) error {
 				return err
 			}
 			o.curKey = key
-			itemBody, ok := o.rebuilt.Load(itemCtx, key.TreeID, key.Key)
+			itemBody, ok := o.rebuilt.Tree(itemCtx, key.TreeID).Load(itemCtx, key.Key)
 			if !ok {
 				o.ioErr(itemCtx, fmt.Errorf("could not read previously read item: %v", key))
 			}
@@ -180,7 +180,7 @@ func (o *rebuilder) Rebuild(_ctx context.Context) error {
 					return err
 				}
 				progressWriter.Set(progress)
-				o.rebuilt.AddRoot(treeCtx, treeID, nodeAddr)
+				o.rebuilt.Tree(treeCtx, treeID).AddRoot(treeCtx, nodeAddr)
 				progress.N++
 			}
 		}
@@ -206,7 +206,7 @@ func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (off
 		o.itemQueue.Insert(o.curKey)
 		return 0, btrfsitem.Root{}, false
 	}
-	itemBody, ok := o.rebuilt.Load(ctx, btrfsprim.ROOT_TREE_OBJECTID, key)
+	itemBody, ok := o.rebuilt.Tree(ctx, btrfsprim.ROOT_TREE_OBJECTID).Load(ctx, key)
 	if !ok {
 		o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
 	}
@@ -231,7 +231,7 @@ func (o *rebuilder) cbLookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id b
 		o.itemQueue.Insert(o.curKey)
 		return 0, false
 	}
-	itemBody, ok := o.rebuilt.Load(ctx, btrfsprim.UUID_TREE_OBJECTID, key)
+	itemBody, ok := o.rebuilt.Tree(ctx, btrfsprim.UUID_TREE_OBJECTID).Load(ctx, key)
 	if !ok {
 		o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
 	}
@@ -376,7 +376,7 @@ func (o *rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, cho
 	}
 	choicesWithDist := make(map[btrfsvol.LogicalAddr]int, len(choices))
 	for choice := range choices {
-		dist, ok := o.rebuilt.COWDistance(ctx, treeID, o.graph.Nodes[choice].Owner)
+		dist, ok := o.rebuilt.Tree(ctx, treeID).COWDistance(o.graph.Nodes[choice].Owner)
 		if !ok {
 			panic(fmt.Errorf("should not happen: .wantAugment called for tree=%v with invalid choice=%v", treeID, choice))
 		}
@@ -402,7 +402,7 @@ func (o *rebuilder) want(ctx context.Context, reason string, treeID btrfsprim.Ob
 }
 
 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) {
+	if o.rebuilt.Tree(ctx, treeID) == nil {
 		o.itemQueue.Insert(o.curKey)
 		return btrfsprim.Key{}, false
 	}
@@ -413,7 +413,7 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btr
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	if key, ok := o.rebuilt.Search(ctx, treeID, func(key btrfsprim.Key) int {
+	if key, ok := o.rebuilt.Tree(ctx, treeID).Search(func(key btrfsprim.Key) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
 	}); ok {
@@ -423,10 +423,10 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btr
 	// OK, we need to insert it
 
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.Keys(treeID).Subrange(
+	o.rebuilt.Tree(ctx, treeID).Keys().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))
+			wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(v.Node))
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
@@ -446,24 +446,24 @@ func (o *rebuilder) wantOff(ctx context.Context, reason string, treeID btrfsprim
 }
 
 func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, tgt btrfsprim.Key) (ok bool) {
-	if !o.rebuilt.AddTree(ctx, treeID) {
+	if o.rebuilt.Tree(ctx, treeID) == nil {
 		o.itemQueue.Insert(o.curKey)
 		return false
 	}
 
 	// check if we already have it
 
-	if _, ok := o.rebuilt.Search(ctx, treeID, tgt.Cmp); ok {
+	if _, ok := o.rebuilt.Tree(ctx, treeID).Search(tgt.Cmp); ok {
 		return true
 	}
 
 	// OK, we need to insert it
 
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.Keys(treeID).Subrange(
+	o.rebuilt.Tree(ctx, treeID).Keys().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))
+			wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(v.Node))
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
@@ -471,7 +471,7 @@ func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, tgt bt
 }
 
 func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(keyio.ItemPtr) bool) {
-	if !o.rebuilt.AddTree(ctx, treeID) {
+	if o.rebuilt.Tree(ctx, treeID) == nil {
 		o.itemQueue.Insert(o.curKey)
 		return
 	}
@@ -482,12 +482,12 @@ func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	keys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
+	keys := o.rebuilt.Tree(ctx, treeID).SearchAll(func(key btrfsprim.Key) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
 	})
 	for _, itemKey := range keys {
-		itemPtr, ok := o.rebuilt.Resolve(ctx, treeID, itemKey)
+		itemPtr, ok := o.rebuilt.Tree(ctx, treeID).Resolve(itemKey)
 		if !ok {
 			o.ioErr(ctx, fmt.Errorf("could not resolve previously read item: %v", itemKey))
 		}
@@ -499,11 +499,11 @@ func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 	// OK, we need to insert it
 
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.Keys(treeID).Subrange(
+	o.rebuilt.Tree(ctx, treeID).Keys().Subrange(
 		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Cmp(k) },
 		func(k btrfsprim.Key, v keyio.ItemPtr) bool {
 			if fn(v) {
-				wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
+				wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(v.Node))
 			}
 			return true
 		})
@@ -530,13 +530,13 @@ func (o *rebuilder) _wantRange(
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key",
 		fmt.Sprintf("tree=%v key={%v %v ?}", treeID, objID, typ))
 
-	if !o.rebuilt.AddTree(ctx, treeID) {
+	if o.rebuilt.Tree(ctx, treeID) == nil {
 		o.itemQueue.Insert(o.curKey)
 		return
 	}
 
 	sizeFn := func(key btrfsprim.Key) (uint64, error) {
-		ptr, ok := o.rebuilt.Keys(treeID).Load(key)
+		ptr, ok := o.rebuilt.Tree(ctx, treeID).Keys().Load(key)
 		if !ok {
 			panic(fmt.Errorf("should not happen: could not load key: %v", key))
 		}
@@ -558,7 +558,7 @@ func (o *rebuilder) _wantRange(
 		ItemType: typ,
 		Offset:   end - 1,
 	}
-	runKeys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
+	runKeys := o.rebuilt.Tree(ctx, treeID).SearchAll(func(key btrfsprim.Key) int {
 		switch {
 		case runMin.Cmp(key) < 0:
 			return 1
@@ -645,7 +645,7 @@ func (o *rebuilder) _wantRange(
 			ItemType: typ,
 			Offset:   gap.End - 1,
 		}
-		o.rebuilt.Keys(treeID).Subrange(
+		o.rebuilt.Tree(ctx, treeID).Keys().Subrange(
 			func(key btrfsprim.Key, _ keyio.ItemPtr) int {
 				switch {
 				case runMin.Cmp(key) < 0:
@@ -680,7 +680,7 @@ func (o *rebuilder) _wantRange(
 				}
 				wantCtx := dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key",
 					fmt.Sprintf("tree=%v key={%v %v %v-%v}", treeID, objID, typ, gap.Beg, gap.End))
-				o.wantAugment(wantCtx, treeID, o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
+				o.wantAugment(wantCtx, treeID, o.rebuilt.Tree(wantCtx, treeID).LeafToRoots(v.Node))
 				last = runEnd
 
 				return true