rebuild root items

this was sitting here uncommited from Wednesday

5 files changed, 137 insertions, 38 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
index 9cef716..1f8b617 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
@@ -29,11 +29,13 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		uuidMap: uuidMap,
 	}
 
-	orphanedNodes, badNodes, err := classifyNodes(ctx, nfs, nodeScanResults)
+	orphanedNodes, badNodes, treeAncestors, err := classifyNodes(ctx, nfs, nodeScanResults)
 	if err != nil {
 		return nil, err
 	}
 
+	uuidMap.considerAncestors(ctx, treeAncestors)
+
 	rebuiltNodes, err := reInitBrokenNodes(ctx, nfs, badNodes)
 	if err != nil {
 		return nil, err
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go
index 9daf97f..da956de 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go
@@ -61,25 +61,7 @@ func (fs *RebuiltTrees) WriteAt(p []byte, off btrfsvol.LogicalAddr) (int, error)
 // btrfstree.NodeFile
 
 func (fs *RebuiltTrees) 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 := fs.uuidMap.TreeParent[tree]
-	if !ok {
-		// could not look up parent info
-		return 0, false
-	}
-	if parentUUID == (btrfsprim.UUID{}) {
-		// no parent
-		return 0, true
-	}
-	parentObjID, ok := fs.uuidMap.UUID2ObjID[parentUUID]
-	if !ok {
-		// could not look up parent info
-		return 0, false
-	}
-	return parentObjID, true
+	return fs.uuidMap.ParentTree(tree)
 }
 
 // btrfstree.NodeSource
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
index 38946b0..9860355 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
@@ -42,6 +42,103 @@ func (m uuidMap) missingRootItems() map[btrfsprim.ObjID]struct{} {
 	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
+}
+
+func (m uuidMap) considerAncestors(ctx context.Context, treeAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]) {
+	if missing := m.missingRootItems(); len(missing) == 0 {
+		return
+	} else {
+		dlog.Infof(ctx, "Rebuilding %d root items...", len(missing))
+	}
+
+	var fullAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]
+	var getFullAncestors func(child btrfsprim.ObjID) containers.Set[btrfsprim.ObjID]
+	getFullAncestors = func(child btrfsprim.ObjID) containers.Set[btrfsprim.ObjID] {
+		if memoized := fullAncestors[child]; memoized != nil {
+			return memoized
+		}
+		ret := make(containers.Set[btrfsprim.ObjID])
+		fullAncestors[child] = ret
+
+		// Try to use 'm' if possible.
+		knownAncestors := make(containers.Set[btrfsprim.ObjID])
+		if parent, ok := m.ParentTree(child); ok {
+			if parent == 0 {
+				return ret
+			}
+			knownAncestors.Insert(parent)
+		} else {
+			knownAncestors.InsertFrom(treeAncestors[child])
+		}
+
+		for _, ancestor := range maps.SortedKeys(knownAncestors) {
+			ret.Insert(ancestor)
+			ret.InsertFrom(getFullAncestors(ancestor))
+		}
+		return ret
+	}
+
+	for pass := 1; true; pass++ {
+		dlog.Infof(ctx, "... pass %v", pass)
+		added := 0
+		fullAncestors = make(map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID])
+		for _, missingRoot := range maps.SortedKeys(m.missingRootItems()) {
+			if _, ok := m.TreeParent[missingRoot]; ok {
+				// This one is incomplete because it doesn't have a UUID, not
+				// because it doesn't have a parent.
+				continue
+			}
+			dlog.Infof(ctx, "... rebuilding %v", missingRoot)
+			potentialParents := make(containers.Set[btrfsprim.ObjID])
+			potentialParents.InsertFrom(getFullAncestors(missingRoot))
+			for ancestor := range getFullAncestors(missingRoot) {
+				potentialParents.DeleteFrom(getFullAncestors(ancestor))
+			}
+			if len(potentialParents) == 1 {
+				parent := potentialParents.TakeOne()
+				dlog.Infof(ctx, "... ... the parent of %v is %v", missingRoot, parent)
+				parentUUID, ok := m.ObjID2UUID[parent]
+				if !ok {
+					dlog.Errorf(ctx, "... ... but can't synthesize a root item because UUID of %v is unknown", parent)
+					continue
+				}
+				m.TreeParent[missingRoot] = parentUUID
+				added++
+			}
+		}
+		if added == 0 {
+			break
+		}
+	}
+
+	if missing := m.missingRootItems(); len(missing) > 0 {
+		dlog.Errorf(ctx, "... could not rebuild root items for %d trees: %v", len(missing), maps.SortedKeys(missing))
+	} else {
+		dlog.Info(ctx, "... rebuilt all root items")
+	}
+}
+
 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)
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
index e171e49..9274a89 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
@@ -11,10 +11,12 @@ 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"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
@@ -28,9 +30,11 @@ type badNode struct {
 //
 //  1. the set of nodes don't have another node claiming it as a child, and
 //  2. the list of bad nodes (in no particular order)
+//  3. tree ancestors thing
 func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDevicesResult) (
 	orphanedNodes map[btrfsvol.LogicalAddr]struct{},
 	badNodes []badNode,
+	treeAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID],
 	err error,
 ) {
 	dlog.Info(ctx, "Walking trees to identify orphan and broken nodes...")
@@ -48,6 +52,8 @@ func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDev
 		}
 	}
 
+	treeAncestors = make(map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID])
+
 	var potentialRoot btrfsvol.LogicalAddr // zero for non-lost nodes, non-zero for lost nodes
 	nodeHandler := func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
 		if err != nil && (errors.Is(err, btrfstree.ErrNotANode) || errors.As(err, new(*btrfstree.IOError))) {
@@ -64,11 +70,12 @@ func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDev
 			if addr != potentialRoot {
 				delete(orphanedNodes, addr)
 			}
-			// TODO: Compare `nodeRef.Data.Head.Owner` and `path.Node(-1).FromTree` to
-			// maybe reconstruct a missing root item.  This is a sort of catch-22; we
-			// trust this data less because lost nodes may be discarded and not just
-			// lost, but non-lost nodes will never have a missing root item, so lost
-			// nodes are all we have to work with on this.
+		}
+		if err == nil && nodeRef.Data.Head.Owner != path.Node(-1).FromTree {
+			if treeAncestors[path.Node(-1).FromTree] == nil {
+				treeAncestors[path.Node(-1).FromTree] = make(containers.Set[btrfsprim.ObjID])
+			}
+			treeAncestors[path.Node(-1).FromTree].Insert(nodeRef.Data.Head.Owner)
 		}
 		progress()
 		return nil
@@ -130,5 +137,5 @@ func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDev
 	}
 
 	dlog.Infof(ctx, "... identified %d orphaned nodes and %d bad nodes", len(orphanedNodes), len(badNodes))
-	return orphanedNodes, badNodes, nil
+	return orphanedNodes, badNodes, treeAncestors, nil
 }
diff --git a/lib/containers/set.go b/lib/containers/set.go
index 8cd700a..42e5ad2 100644
--- a/lib/containers/set.go
+++ b/lib/containers/set.go
@@ -45,27 +45,38 @@ func (o *Set[T]) DecodeJSON(r io.RuneScanner) error {
 	})
 }
 
-func (o *Set[T]) Insert(v T) {
-	if *o == nil {
-		*o = Set[T]{}
+func (o Set[T]) Insert(v T) {
+	o[v] = struct{}{}
+}
+
+func (o Set[T]) InsertFrom(p Set[T]) {
+	for v := range p {
+		o[v] = struct{}{}
 	}
-	(*o)[v] = struct{}{}
 }
 
-func (o *Set[T]) InsertFrom(p Set[T]) {
-	if *o == nil {
-		*o = Set[T]{}
+func (o Set[T]) Delete(v T) {
+	if o == nil {
+		return
+	}
+	delete(o, v)
+}
+
+func (o Set[T]) DeleteFrom(p Set[T]) {
+	if o == nil {
+		return
 	}
 	for v := range p {
-		(*o)[v] = struct{}{}
+		delete(o, v)
 	}
 }
 
-func (o *Set[T]) Delete(v T) {
-	if *o == nil {
-		return
+func (o Set[T]) TakeOne() T {
+	for v := range o {
+		return v
 	}
-	delete(*o, v)
+	var zero T
+	return zero
 }
 
 func (small Set[T]) HasIntersection(big Set[T]) bool {