wip better reattach

2 files changed, 92 insertions, 44 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s4_reattach.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s4_reattach.go
index 28bd834..45fcc23 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s4_reattach.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s4_reattach.go
@@ -15,9 +15,22 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
-	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
+func (a RebuiltNode) ContainsWholeRegion(min, max btrfsprim.Key) int {
+	switch {
+	case min.Cmp(a.MinKey) < 0:
+		// 'a' is too far right
+		return -1
+	case max.Cmp(a.MaxKey) > 0:
+		// 'a' is too far left
+		return 1
+	default:
+		// just right
+		return 0
+	}
+}
+
 func reAttachNodes(ctx context.Context, fs _FS, orphanedNodes map[btrfsvol.LogicalAddr]struct{}, rebuiltNodes map[btrfsvol.LogicalAddr]*RebuiltNode) error {
 	dlog.Info(ctx, "Attaching orphaned nodes to rebuilt nodes...")
 
@@ -28,23 +41,17 @@ func reAttachNodes(ctx context.Context, fs _FS, orphanedNodes map[btrfsvol.Logic
 
 	// Index 'rebuiltNodes' for fast lookups.
 	dlog.Info(ctx, "... indexing rebuilt nodes...")
-	gaps := make(map[btrfsprim.ObjID]map[uint8][]*RebuiltNode)
-	maxLevel := make(map[btrfsprim.ObjID]uint8)
+	var byLevel [][]*RebuiltNode
 	for _, node := range rebuiltNodes {
-		for treeID := range node.InTrees {
-			maxLevel[treeID] = slices.Max(maxLevel[treeID], node.Head.Level)
-			if gaps[treeID] == nil {
-				gaps[treeID] = make(map[uint8][]*RebuiltNode)
-			}
-			gaps[treeID][node.Head.Level] = append(gaps[treeID][node.Head.Level], node)
+		for int(node.Head.Level) >= len(byLevel) {
+			byLevel = append(byLevel, []*RebuiltNode(nil))
 		}
+		byLevel[node.Head.Level] = append(byLevel[node.Head.Level], node)
 	}
-	for _, byTreeID := range gaps {
-		for _, slice := range byTreeID {
-			sort.Slice(slice, func(i, j int) bool {
-				return slice[i].MinKey.Cmp(slice[j].MinKey) < 0
-			})
-		}
+	for _, slice := range byLevel {
+		sort.Slice(slice, func(i, j int) bool {
+			return slice[i].MinKey.Cmp(slice[j].MinKey) < 0
+		})
 	}
 	dlog.Info(ctx, "... done indexing")
 
@@ -76,40 +83,69 @@ func reAttachNodes(ctx context.Context, fs _FS, orphanedNodes map[btrfsvol.Logic
 		if !ok {
 			continue
 		}
-		treeGaps := gaps[foundRef.Data.Head.Owner]
-		var attached bool
-		for level := foundRef.Data.Head.Level + 1; treeGaps != nil && level <= maxLevel[foundRef.Data.Head.Owner] && !attached; level++ {
-			parentGen, ok := treeGaps[level]
+
+		// `trees` is the set of trees that the node may be
+		// placed in; '0' is a wildcard that means "any tree".
+		// We still keep track of the others, in order to try
+		// to avoid using the wildcard.
+		trees := make(containers.Set[btrfsprim.ObjID])
+		tree := foundRef.Data.Head.Owner
+		for {
+			trees[tree] = struct{}{}
+			var ok bool
+			tree, ok = fs.ParentTree(tree)
 			if !ok {
-				continue
+				// error; accept anything
+				trees[0] = struct{}{}
+				break
+			}
+			if tree == 0 {
+				// end of the line
+				break
+			}
+		}
+		attached := make(containers.Set[btrfsprim.ObjID])
+		for level := int(foundRef.Data.Head.Level) + 1; level < len(byLevel) && len(attached) == 0; level++ {
+			for _, parent := range byLevel[level] {
+				if parent.ContainsWholeRegion(foundMinKey, foundMaxKey) != 0 {
+					continue
+				}
+				if parent.Node.Head.Generation < foundRef.Data.Head.Generation {
+					continue
+				}
+				if !trees.HasIntersection(parent.InTrees) {
+					continue
+				}
+				parent.BodyInternal = append(parent.BodyInternal, btrfstree.KeyPointer{
+					Key:        foundMinKey,
+					BlockPtr:   foundLAddr,
+					Generation: foundRef.Data.Head.Generation,
+				})
+				attached.InsertFrom(parent.InTrees)
 			}
-			parentIdx, ok := slices.Search(parentGen, func(parent *RebuiltNode) int {
-				switch {
-				case foundMinKey.Cmp(parent.MinKey) < 0:
-					// 'parent' is too far right
-					return -1
-				case foundMaxKey.Cmp(parent.MaxKey) > 0:
-					// 'parent' is too far left
-					return 1
-				default:
-					// just right
-					return 0
+		}
+		if _, wildcard := trees[0]; wildcard && len(attached) == 0 {
+			for level := int(foundRef.Data.Head.Level) + 1; level < len(byLevel) && len(attached) == 0; level++ {
+				for _, parent := range byLevel[level] {
+					if parent.ContainsWholeRegion(foundMinKey, foundMaxKey) != 0 {
+						continue
+					}
+					if parent.Node.Head.Generation < foundRef.Data.Head.Generation {
+						continue
+					}
+					parent.BodyInternal = append(parent.BodyInternal, btrfstree.KeyPointer{
+						Key:        foundMinKey,
+						BlockPtr:   foundLAddr,
+						Generation: foundRef.Data.Head.Generation,
+					})
+					attached.InsertFrom(parent.InTrees)
 				}
-			})
-			if !ok {
-				continue
 			}
-			parent := parentGen[parentIdx]
-			parent.BodyInternal = append(parent.BodyInternal, btrfstree.KeyPointer{
-				Key:        foundMinKey,
-				BlockPtr:   foundLAddr,
-				Generation: foundRef.Data.Head.Generation,
-			})
-			parent.Head.Generation = slices.Max(parent.Head.Generation, foundRef.Data.Head.Generation)
-			attached = true
-			numAttached++
 		}
-		if !attached {
+
+		if len(attached) > 0 {
+			numAttached++
+		} else {
 			dlog.Errorf(ctx, "could not find a broken node to attach node to reattach node@%v to",
 				foundRef.Addr)
 		}
diff --git a/lib/containers/set.go b/lib/containers/set.go
index 4e40894..8cd700a 100644
--- a/lib/containers/set.go
+++ b/lib/containers/set.go
@@ -67,3 +67,15 @@ func (o *Set[T]) Delete(v T) {
 	}
 	delete(*o, v)
 }
+
+func (small Set[T]) HasIntersection(big Set[T]) bool {
+	if len(big) < len(small) {
+		small, big = big, small
+	}
+	for v := range small {
+		if _, ok := big[v]; ok {
+			return true
+		}
+	}
+	return false
+}