better algo

3 files changed, 139 insertions, 103 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/scanforextents/blockgroups.go b/lib/btrfsprogs/btrfsinspect/scanforextents/blockgroups.go
index 99009e6..64ec8c3 100644
--- a/lib/btrfsprogs/btrfsinspect/scanforextents/blockgroups.go
+++ b/lib/btrfsprogs/btrfsinspect/scanforextents/blockgroups.go
@@ -8,16 +8,23 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"sort"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"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/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
 type NodeScanResults = map[btrfsvol.DeviceID]btrfsinspect.ScanOneDevResult
 
-func ReadNodeScanResults(fs *btrfs.FS, filename string) (*BlockGroupTree, error) {
+type BlockGroup struct {
+	LAddr btrfsvol.LogicalAddr
+	Size  btrfsvol.AddrDelta
+	Flags btrfsvol.BlockGroupFlags
+}
+
+func ReadNodeScanResults(fs *btrfs.FS, filename string) (map[btrfsvol.LogicalAddr]BlockGroup, error) {
 	scanResultsBytes, err := os.ReadFile(filename)
 	if err != nil {
 		return nil, err
@@ -36,43 +43,8 @@ func ReadNodeScanResults(fs *btrfs.FS, filename string) (*BlockGroupTree, error)
 	return bgTree, nil
 }
 
-type BlockGroup struct {
-	LAddr btrfsvol.LogicalAddr
-	Size  btrfsvol.AddrDelta
-	Flags btrfsvol.BlockGroupFlags
-}
-
-type BlockGroupTree = containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], BlockGroup]
-
-func LookupBlockGroup(tree *BlockGroupTree, laddr btrfsvol.LogicalAddr, size btrfsvol.AddrDelta) *BlockGroup {
-	a := struct {
-		LAddr btrfsvol.LogicalAddr
-		Size  btrfsvol.AddrDelta
-	}{
-		LAddr: laddr,
-		Size:  size,
-	}
-	node := tree.Search(func(b BlockGroup) int {
-		switch {
-		case a.LAddr.Add(a.Size) <= b.LAddr:
-			// 'a' is wholly to the left of 'b'.
-			return -1
-		case b.LAddr.Add(b.Size) <= a.LAddr:
-			// 'a' is wholly to the right of 'b'.
-			return 1
-		default:
-			// There is some overlap.
-			return 0
-		}
-	})
-	if node == nil {
-		return nil
-	}
-	bg := node.Value
-	return &bg
-}
-
-func ReduceScanResults(fs *btrfs.FS, scanResults NodeScanResults) (*BlockGroupTree, error) {
+func ReduceScanResults(fs *btrfs.FS, scanResults NodeScanResults) (map[btrfsvol.LogicalAddr]BlockGroup, error) {
+	// Reduce
 	bgSet := make(map[BlockGroup]struct{})
 	for _, found := range scanResults {
 		for _, bg := range found.FoundBlockGroups {
@@ -83,19 +55,24 @@ func ReduceScanResults(fs *btrfs.FS, scanResults NodeScanResults) (*BlockGroupTr
 			}] = struct{}{}
 		}
 	}
-	bgTree := &BlockGroupTree{
-		KeyFn: func(bg BlockGroup) containers.NativeOrdered[btrfsvol.LogicalAddr] {
-			return containers.NativeOrdered[btrfsvol.LogicalAddr]{Val: bg.LAddr}
-		},
-	}
-	for bg := range bgSet {
-		if laddr, _ := fs.LV.ResolveAny(bg.LAddr, bg.Size); laddr >= 0 {
-			continue
-		}
-		if LookupBlockGroup(bgTree, bg.LAddr, bg.Size) != nil {
+
+	// Sanity check
+	bgList := maps.Keys(bgSet)
+	sort.Slice(bgList, func(i, j int) bool {
+		return bgList[i].LAddr < bgList[j].LAddr
+	})
+	var pos btrfsvol.LogicalAddr
+	for _, bg := range bgList {
+		if bg.LAddr < pos || bg.Size < 0 {
 			return nil, fmt.Errorf("found block groups are inconsistent")
 		}
-		bgTree.Insert(bg)
+		pos = bg.LAddr.Add(bg.Size)
 	}
-	return bgTree, nil
+
+	// Return
+	bgMap := make(map[btrfsvol.LogicalAddr]BlockGroup, len(bgSet))
+	for bg := range bgSet {
+		bgMap[bg.LAddr] = bg
+	}
+	return bgMap, nil
 }
diff --git a/lib/btrfsprogs/btrfsinspect/scanforextents/gaps.go b/lib/btrfsprogs/btrfsinspect/scanforextents/gaps.go
index f67b498..5963e53 100644
--- a/lib/btrfsprogs/btrfsinspect/scanforextents/gaps.go
+++ b/lib/btrfsprogs/btrfsinspect/scanforextents/gaps.go
@@ -12,6 +12,7 @@ import (
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
 type PhysicalGap struct {
@@ -53,30 +54,25 @@ func roundUp[T constraints.Integer](x, multiple T) T {
 }
 
 func WalkGaps(ctx context.Context,
-	gaps []PhysicalGap, blockSize btrfsvol.PhysicalAddr,
-	progress func(cur, total int64),
-	main func(btrfsvol.PhysicalAddr) error,
+	sums AllSums, gaps map[btrfsvol.DeviceID][]PhysicalGap,
+	fn func(btrfsvol.DeviceID, SumRun[btrfsvol.PhysicalAddr]) error,
 ) error {
-	var totalBlocks int64
-	for _, gap := range gaps {
-		for paddr := roundUp(gap.Beg, blockSize); paddr+blockSize <= gap.End; paddr += blockSize {
-			totalBlocks++
-		}
-	}
-
-	var curBlock int64
-	for _, gap := range gaps {
-		for paddr := roundUp(gap.Beg, blockSize); paddr+blockSize <= gap.End; paddr += blockSize {
+	for _, devID := range maps.SortedKeys(gaps) {
+		for _, gap := range gaps[devID] {
 			if err := ctx.Err(); err != nil {
 				return err
 			}
-			progress(curBlock, totalBlocks)
-			curBlock++
-			if err := main(paddr); err != nil {
+			begAddr := roundUp(gap.Beg, csumBlockSize)
+			begOff := int(begAddr/csumBlockSize) * sums.Physical[devID].ChecksumSize
+			endOff := int(gap.End/csumBlockSize) * sums.Physical[devID].ChecksumSize
+			if err := fn(devID, SumRun[btrfsvol.PhysicalAddr]{
+				ChecksumSize: sums.Physical[devID].ChecksumSize,
+				Addr:         begAddr,
+				Sums:         sums.Physical[devID].Sums[begOff:endOff],
+			}); err != nil {
 				return err
 			}
 		}
 	}
-	progress(curBlock, totalBlocks)
 	return nil
 }
diff --git a/lib/btrfsprogs/btrfsinspect/scanforextents/scan.go b/lib/btrfsprogs/btrfsinspect/scanforextents/scan.go
index 77a4ed9..12ade04 100644
--- a/lib/btrfsprogs/btrfsinspect/scanforextents/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/scanforextents/scan.go
@@ -7,71 +7,128 @@ package scanforextents
 import (
 	"context"
 	"errors"
-	"fmt"
 
 	"github.com/datawire/dlib/dlog"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"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"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
-func ScanForExtents(ctx context.Context, fs *btrfs.FS, blockgroups *BlockGroupTree, sums AllSums) error {
-	sb, err := fs.Superblock()
-	if err != nil {
-		return err
+func ScanForExtents(ctx context.Context, fs *btrfs.FS, blockgroups map[btrfsvol.LogicalAddr]BlockGroup, sums AllSums) error {
+	dlog.Info(ctx, "Pairing up blockgroups and sums...")
+	bgSums := make(map[btrfsvol.LogicalAddr][]SumRun[btrfsvol.LogicalAddr])
+	for _, blockgroup := range blockgroups {
+		for laddr := blockgroup.LAddr; laddr < blockgroup.LAddr.Add(blockgroup.Size); {
+			run, ok := sums.RunForLAddr(laddr)
+			if !ok {
+				laddr += csumBlockSize
+				continue
+			}
+			off := int((laddr-run.Addr)/csumBlockSize) * run.ChecksumSize
+			deltaAddr := slices.Min[btrfsvol.AddrDelta](
+				blockgroup.Size-laddr.Sub(blockgroup.LAddr),
+				btrfsvol.AddrDelta((len(run.Sums)-off)/run.ChecksumSize)*csumBlockSize)
+			deltaOff := int(deltaAddr/csumBlockSize) * run.ChecksumSize
+			bgSums[blockgroup.LAddr] = append(bgSums[blockgroup.LAddr], SumRun[btrfsvol.LogicalAddr]{
+				ChecksumSize: run.ChecksumSize,
+				Addr:         laddr,
+				Sums:         run.Sums[off : off+deltaOff],
+			})
+			laddr = laddr.Add(deltaAddr)
+		}
+		return nil
 	}
+	dlog.Info(ctx, "... done pairing")
 
-	dlog.Info(ctx, "Reverse-indexing and validating logical sums...")
-	var totalSums int
-	if err := sums.WalkLogical(ctx, func(btrfsvol.LogicalAddr, ShortSum) error {
-		totalSums++
-		return nil
-	}); err != nil {
-		return err
+	dlog.Info(ctx, "Searching for unmapped blockgroups in unmapped regions...")
+	gaps := ListPhysicalGaps(fs)
+	bgMatches := make(map[btrfsvol.LogicalAddr][]btrfsvol.QualifiedPhysicalAddr)
+	for i, bgLAddr := range maps.SortedKeys(blockgroups) {
+		bgRuns := bgSums[bgLAddr]
+		if len(bgRuns) != 1 {
+			// TODO(lukeshu): We aught to handle this rather than erroring and skipping
+			// it.
+			dlog.Errorf(ctx, "blockgroup laddr=%v has holes (%v runs)", bgLAddr, len(bgRuns))
+			continue
+		}
+		bgRun := bgRuns[0]
+
+		if err := WalkGaps(ctx, sums, gaps, func(devID btrfsvol.DeviceID, gap SumRun[btrfsvol.PhysicalAddr]) error {
+			matches, err := diskio.IndexAll[int64, ShortSum](gap, bgRun)
+			if err != nil {
+				return err
+			}
+			for _, match := range matches {
+				bgMatches[bgLAddr] = append(bgMatches[bgLAddr], btrfsvol.QualifiedPhysicalAddr{
+					Dev:  devID,
+					Addr: gap.Addr + (btrfsvol.PhysicalAddr(match) * csumBlockSize),
+				})
+			}
+			return nil
+		}); err != nil {
+			return err
+		}
+		dlog.Infof(ctx, "... (%v/%v) blockgroup[laddr=%v] has %v matches",
+			i+1, len(bgSums), bgLAddr, len(bgMatches[bgLAddr]))
 	}
-	sum2laddrs := make(map[ShortSum][]btrfsvol.LogicalAddr)
-	var curSum int
-	lastPct := -1
-	progress := func(curSum int) {
-		pct := int(100 * float64(curSum) / float64(totalSums))
-		if pct != lastPct || curSum == totalSums {
-			dlog.Infof(ctx, "... reversed+validated %v%%", pct)
-			lastPct = pct
+	dlog.Info(ctx, "... done searching")
+
+	dlog.Info(ctx, "Applying those mappings...")
+	for _, bgLAddr := range maps.SortedKeys(bgMatches) {
+		matches := bgMatches[bgLAddr]
+		if len(matches) != 1 {
+			continue
+		}
+		blockgroup := blockgroups[bgLAddr]
+		mapping := btrfsvol.Mapping{
+			LAddr:      blockgroup.LAddr,
+			PAddr:      matches[0],
+			Size:       blockgroup.Size,
+			SizeLocked: true,
+			Flags: containers.Optional[btrfsvol.BlockGroupFlags]{
+				OK:  true,
+				Val: blockgroup.Flags,
+			},
+		}
+		if err := fs.LV.AddMapping(mapping); err != nil {
+			dlog.Error(ctx, err)
 		}
 	}
-	if err := sums.WalkLogical(ctx, func(laddr btrfsvol.LogicalAddr, expShortSum ShortSum) error {
-		progress(curSum)
-		curSum++
-		readSum, err := ChecksumLogical(fs, sb.ChecksumType, laddr)
-		if err != nil {
+	dlog.Info(ctx, "... done applying")
+
+	dlog.Info(ctx, "Reverse-indexing remaining unmapped logical sums...")
+	sum2laddrs := make(map[ShortSum][]btrfsvol.LogicalAddr)
+	var numUnmappedBlocks int64
+	if err := sums.WalkLogical(ctx, func(laddr btrfsvol.LogicalAddr, sum ShortSum) error {
+		var dat [csumBlockSize]byte
+		if _, err := fs.ReadAt(dat[:], laddr); err != nil {
 			if errors.Is(err, btrfsvol.ErrCouldNotMap) {
-				sum2laddrs[expShortSum] = append(sum2laddrs[expShortSum], laddr)
+				sum2laddrs[sum] = append(sum2laddrs[sum], laddr)
+				numUnmappedBlocks++
 				return nil
 			}
 			return err
 		}
-		readShortSum := ShortSum(readSum[:len(expShortSum)])
-		if readShortSum != expShortSum {
-			return fmt.Errorf("checksum mismatch at laddr=%v: CSUM_TREE=%x != read=%x",
-				laddr, []byte(expShortSum), []byte(readShortSum))
-		}
 		return nil
 	}); err != nil {
 		return err
 	}
-	progress(totalSums)
-	dlog.Info(ctx, "... done reverse-indexing and validating")
+	dlog.Infof(ctx, "... done reverse-indexing; %v still unmapped logical sums",
+		numUnmappedBlocks)
+
+	/* TODO
 
-	dlog.Info(ctx, "Cross-referencing sums (and blockgroups) to re-construct mappings...")
+	dlog.Info(ctx, "Cross-referencing sums to re-construct mappings...")
 	newMappings := &ExtentMappings{
 		InLV:          &fs.LV,
 		InBlockGroups: blockgroups,
 		InSums:        sums,
 		InReverseSums: sum2laddrs,
 	}
-	devs := fs.LV.PhysicalVolumes()
 	gaps := ListPhysicalGaps(fs)
 	for _, devID := range maps.SortedKeys(gaps) {
 		if err := newMappings.ScanOneDev(ctx,
@@ -98,9 +155,13 @@ func ScanForExtents(ctx context.Context, fs *btrfs.FS, blockgroups *BlockGroupTr
 	}
 	dlog.Info(ctx, "... done applying")
 
+	*/
+
 	return nil
 }
 
+/*
+
 type ExtentMappings struct {
 	// input
 	InLV          *btrfsvol.LogicalVolume[*btrfs.Device]
@@ -233,3 +294,5 @@ func (em *ExtentMappings) ScanOneDev(
 		},
 	)
 }
+
+*/