1 files changed, 162 insertions, 0 deletions
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go b/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go
new file mode 100644
index 0000000..3430502
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildmappings/logicalsums.go
@@ -0,0 +1,162 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package rebuildmappings
+
+import (
+	"context"
+	"strings"
+
+	"github.com/datawire/dlib/dlog"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfssum"
+	"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/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
+)
+
+func ExtractLogicalSums(ctx context.Context, scanResults btrfsinspect.ScanDevicesResult) btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr] {
+	dlog.Info(ctx, "... Mapping the logical address space...")
+	type record struct {
+		Gen btrfs.Generation
+		Sum btrfssum.ShortSum
+	}
+	addrspace := make(map[btrfsvol.LogicalAddr]record)
+	var sumSize int
+	for _, devResults := range scanResults {
+		sumSize = devResults.Checksums.ChecksumSize
+		for _, sumItem := range devResults.FoundExtentCSums {
+			_ = sumItem.Sums.Walk(ctx, func(pos btrfsvol.LogicalAddr, sum btrfssum.ShortSum) error {
+				new := record{
+					Gen: sumItem.Generation,
+					Sum: sum,
+				}
+				if old, ok := addrspace[pos]; ok {
+					switch {
+					case old.Gen > new.Gen:
+						// do nothing
+					case old.Gen < new.Gen:
+						addrspace[pos] = new
+					case old.Gen == new.Gen:
+						if old != new {
+							dlog.Errorf(ctx, "mismatch of laddr=%v sum: %v != %v", pos, old, new)
+						}
+					}
+				} else {
+					addrspace[pos] = new
+				}
+				return nil
+			})
+		}
+	}
+	dlog.Info(ctx, "... ... done mapping")
+
+	var flattened btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr]
+	if len(addrspace) == 0 {
+		return flattened
+	}
+
+	dlog.Info(ctx, "... Flattening the map ...")
+	var curAddr btrfsvol.LogicalAddr
+	var curSums strings.Builder
+	for _, laddr := range maps.SortedKeys(addrspace) {
+		if laddr != curAddr+(btrfsvol.LogicalAddr(curSums.Len()/sumSize)*btrfssum.BlockSize) {
+			if curSums.Len() > 0 {
+				flattened.Runs = append(flattened.Runs, btrfssum.SumRun[btrfsvol.LogicalAddr]{
+					ChecksumSize: sumSize,
+					Addr:         curAddr,
+					Sums:         btrfssum.ShortSum(curSums.String()),
+				})
+			}
+			curAddr = laddr
+			curSums.Reset()
+		}
+		curSums.WriteString(string(addrspace[laddr].Sum))
+	}
+	if curSums.Len() > 0 {
+		flattened.Runs = append(flattened.Runs, btrfssum.SumRun[btrfsvol.LogicalAddr]{
+			ChecksumSize: sumSize,
+			Addr:         curAddr,
+			Sums:         btrfssum.ShortSum(curSums.String()),
+		})
+	}
+	flattened.Addr = flattened.Runs[0].Addr
+	last := flattened.Runs[len(flattened.Runs)-1]
+	end := last.Addr.Add(last.Size())
+	flattened.Size = end.Sub(flattened.Addr)
+	dlog.Info(ctx, "... ... done flattening")
+
+	return flattened
+}
+
+func ListUnmappedLogicalRegions(fs *btrfs.FS, logicalSums btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr]) []btrfssum.SumRun[btrfsvol.LogicalAddr] {
+	// There are a lot of ways this algorithm could be made
+	// faster.
+	var ret []btrfssum.SumRun[btrfsvol.LogicalAddr]
+	var cur struct {
+		Addr btrfsvol.LogicalAddr
+		Size btrfsvol.AddrDelta
+	}
+	for _, run := range logicalSums.Runs {
+		for addr := run.Addr; addr < run.Addr.Add(run.Size()); addr += btrfssum.BlockSize {
+			if _, maxlen := fs.LV.Resolve(addr); maxlen < btrfssum.BlockSize {
+				if cur.Size == 0 {
+					cur.Addr = addr
+					cur.Size = 0
+				}
+				cur.Size += btrfssum.BlockSize
+			} else if cur.Size > 0 {
+				begIdx := int(cur.Addr.Sub(run.Addr)/btrfssum.BlockSize) * run.ChecksumSize
+				lenIdx := (int(cur.Size) / btrfssum.BlockSize) * run.ChecksumSize
+				endIdx := begIdx + lenIdx
+				ret = append(ret, btrfssum.SumRun[btrfsvol.LogicalAddr]{
+					ChecksumSize: run.ChecksumSize,
+					Addr:         cur.Addr,
+					Sums:         run.Sums[begIdx:endIdx],
+				})
+				cur.Size = 0
+			}
+		}
+		if cur.Size > 0 {
+			begIdx := int(cur.Addr.Sub(run.Addr)/btrfssum.BlockSize) * run.ChecksumSize
+			lenIdx := (int(cur.Size) / btrfssum.BlockSize) * run.ChecksumSize
+			endIdx := begIdx + lenIdx
+			ret = append(ret, btrfssum.SumRun[btrfsvol.LogicalAddr]{
+				ChecksumSize: run.ChecksumSize,
+				Addr:         cur.Addr,
+				Sums:         run.Sums[begIdx:endIdx],
+			})
+			cur.Size = 0
+		}
+	}
+	return ret
+}
+
+func SumsForLogicalRegion(sums btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr], beg btrfsvol.LogicalAddr, size btrfsvol.AddrDelta) btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr] {
+	runs := btrfssum.SumRunWithGaps[btrfsvol.LogicalAddr]{
+		Addr: beg,
+		Size: size,
+	}
+	for laddr := beg; laddr < beg.Add(size); {
+		run, next, ok := sums.RunForAddr(laddr)
+		if !ok {
+			laddr = next
+			continue
+		}
+		off := int((laddr-run.Addr)/btrfssum.BlockSize) * run.ChecksumSize
+		deltaAddr := slices.Min[btrfsvol.AddrDelta](
+			size-laddr.Sub(beg),
+			btrfsvol.AddrDelta((len(run.Sums)-off)/run.ChecksumSize)*btrfssum.BlockSize)
+		deltaOff := int(deltaAddr/btrfssum.BlockSize) * run.ChecksumSize
+		runs.Runs = append(runs.Runs, btrfssum.SumRun[btrfsvol.LogicalAddr]{
+			ChecksumSize: run.ChecksumSize,
+			Addr:         laddr,
+			Sums:         run.Sums[off : off+deltaOff],
+		})
+		laddr = laddr.Add(deltaAddr)
+	}
+	return runs
+}