rethink the SumRun interface to fit diskio.Sequence

4 files changed, 186 insertions, 106 deletions

diff --git a/lib/btrfs/btrfsvol/addr.go b/lib/btrfs/btrfsvol/addr.go
index 76b531d..94320ef 100644
--- a/lib/btrfs/btrfsvol/addr.go
+++ b/lib/btrfs/btrfsvol/addr.go
@@ -56,3 +56,13 @@ func (a QualifiedPhysicalAddr) Cmp(b QualifiedPhysicalAddr) int {
 	}
 	return int(a.Addr - b.Addr)
 }
+
+type _IntAddr[T any] interface {
+	~int64
+	Sub(T) AddrDelta
+	Add(AddrDelta) T
+}
+
+// IntAddr is a type constraint that you can use to accept LogicalAddr
+// or PhysicalAddr.
+type IntAddr[T _IntAddr[T]] _IntAddr[T]
diff --git a/lib/btrfsprogs/btrfsinspect/scanforextents/csums.go b/lib/btrfsprogs/btrfsinspect/scanforextents/csums.go
index 6a42946..10b70aa 100644
--- a/lib/btrfsprogs/btrfsinspect/scanforextents/csums.go
+++ b/lib/btrfsprogs/btrfsinspect/scanforextents/csums.go
@@ -7,7 +7,6 @@ package scanforextents
 import (
 	"context"
 	"encoding/gob"
-	"fmt"
 	"io"
 	"os"
 	"runtime"
@@ -26,34 +25,66 @@ import (
 
 const csumBlockSize = 4 * 1024
 
+// ShortSum //////////////////////////////////////////////////////////
+
 type ShortSum string
 
-type LogicalSumRun struct {
-	Addr btrfsvol.LogicalAddr
+// SumRun ////////////////////////////////////////////////////////////
+
+type SumRun[Addr btrfsvol.IntAddr[Addr]] struct {
+	// How big a ShortSum is in this Run.
+	ChecksumSize int
+	// Base address where this run starts.
+	Addr Addr
+	// All of the ShortSums in this run, concatenated together.
+	//
+	// This is a 'string' rather than a 'ShortSum' to make it hard
+	// to accidentally use it as a single sum.
 	Sums string
 }
 
-type AllSums struct {
-	ChecksumSize int
-	Logical      []LogicalSumRun
-	Physical     map[btrfsvol.DeviceID]string
+func (run SumRun[Addr]) NumSums() int {
+	return len(run.Sums) / run.ChecksumSize
 }
 
-func ReadAllSums(filename string) (AllSums, error) {
-	fh, err := os.Open(filename)
-	if err != nil {
-		return AllSums{}, err
+func (run SumRun[Addr]) Size() btrfsvol.AddrDelta {
+	return btrfsvol.AddrDelta(run.NumSums()) * csumBlockSize
+}
+
+// Get implements diskio.Sequence[int, ShortSum]
+func (run SumRun[Addr]) Get(sumIdx int64) (ShortSum, error) {
+	if sumIdx < 0 || int(sumIdx) > run.NumSums() {
+		return "", io.EOF
 	}
-	defer fh.Close()
-	var val AllSums
-	if err := gob.NewDecoder(fh).Decode(&val); err != nil {
-		return AllSums{}, err
+	off := int(sumIdx) * run.ChecksumSize
+	return ShortSum(run.Sums[off : off+run.ChecksumSize]), nil
+}
+
+func (run SumRun[Addr]) SumForAddr(addr Addr) (ShortSum, bool) {
+	if addr < run.Addr || addr >= run.Addr.Add(run.Size()) {
+		return "", false
 	}
-	return val, nil
+	off := int((addr-run.Addr)/csumBlockSize) * run.ChecksumSize
+	return ShortSum(run.Sums[off : off+run.ChecksumSize]), true
 }
 
-func WriteAllSums(w io.Writer, sums AllSums) error {
-	return gob.NewEncoder(w).Encode(sums)
+func (run SumRun[Addr]) Walk(ctx context.Context, fn func(Addr, ShortSum) error) error {
+	for addr, off := run.Addr, 0; off < len(run.Sums); addr, off = addr+csumBlockSize, off+run.ChecksumSize {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		if err := fn(addr, ShortSum(run.Sums[off:off+run.ChecksumSize])); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// AllSums ///////////////////////////////////////////////////////////
+
+type AllSums struct {
+	Logical  []SumRun[btrfsvol.LogicalAddr]
+	Physical map[btrfsvol.DeviceID]SumRun[btrfsvol.PhysicalAddr]
 }
 
 func (as AllSums) SumForPAddr(paddr btrfsvol.QualifiedPhysicalAddr) (ShortSum, bool) {
@@ -61,51 +92,71 @@ func (as AllSums) SumForPAddr(paddr btrfsvol.QualifiedPhysicalAddr) (ShortSum, b
 	if !ok {
 		return "", false
 	}
-	off := int(paddr.Addr/csumBlockSize) * as.ChecksumSize
-	if off+as.ChecksumSize > len(run) {
-		return "", false
-	}
-	return ShortSum(run[off : off+as.ChecksumSize]), true
+	return run.SumForAddr(paddr.Addr)
 }
 
-func (as AllSums) SumForLAddr(laddr btrfsvol.LogicalAddr) (ShortSum, bool) {
+func (as AllSums) RunForLAddr(laddr btrfsvol.LogicalAddr) (SumRun[btrfsvol.LogicalAddr], bool) {
 	for _, run := range as.Logical {
-		size := btrfsvol.AddrDelta(len(run.Sums)/as.ChecksumSize) * csumBlockSize
 		if run.Addr > laddr {
-			return "", false
+			return SumRun[btrfsvol.LogicalAddr]{}, false
 		}
-		if run.Addr.Add(size) <= laddr {
+		if run.Addr.Add(run.Size()) <= laddr {
 			continue
 		}
-		off := int(laddr.Sub(run.Addr)/csumBlockSize) * as.ChecksumSize
-		return ShortSum(run.Sums[off : off+as.ChecksumSize]), true
+		return run, true
+	}
+	return SumRun[btrfsvol.LogicalAddr]{}, false
+}
+
+func (as AllSums) SumForLAddr(laddr btrfsvol.LogicalAddr) (ShortSum, bool) {
+	run, ok := as.RunForLAddr(laddr)
+	if !ok {
+		return "", false
 	}
-	return "", false
+	return run.SumForAddr(laddr)
 }
 
-func (as AllSums) WalkLogical(fn func(btrfsvol.LogicalAddr, ShortSum) error) error {
+func (as AllSums) WalkLogical(ctx context.Context, fn func(btrfsvol.LogicalAddr, ShortSum) error) error {
 	for _, run := range as.Logical {
-		for laddr, off := run.Addr, 0; off < len(run.Sums); laddr, off = laddr+csumBlockSize, off+as.ChecksumSize {
-			if err := fn(laddr, ShortSum(run.Sums[off:off+as.ChecksumSize])); err != nil {
-				return err
-			}
+		if err := run.Walk(ctx, fn); err != nil {
+			return err
 		}
 	}
 	return nil
 }
 
+// Read/Write AllSums ////////////////////////////////////////////////
+
+func ReadAllSums(filename string) (AllSums, error) {
+	fh, err := os.Open(filename)
+	if err != nil {
+		return AllSums{}, err
+	}
+	defer fh.Close()
+	var val AllSums
+	if err := gob.NewDecoder(fh).Decode(&val); err != nil {
+		return AllSums{}, err
+	}
+	return val, nil
+}
+
+func WriteAllSums(w io.Writer, sums AllSums) error {
+	return gob.NewEncoder(w).Encode(sums)
+}
+
 func SumEverything(ctx context.Context, fs *btrfs.FS) (AllSums, error) {
 	var ret AllSums
 
 	// ChecksumSize
 	var alg btrfssum.CSumType
+	var csumSize int
 	if err := func() error {
 		sb, err := fs.Superblock()
 		if err != nil {
 			return err
 		}
 		alg = sb.ChecksumType
-		ret.ChecksumSize = alg.Size()
+		csumSize = alg.Size()
 		return nil
 	}(); err != nil {
 		return ret, err
@@ -131,24 +182,26 @@ func SumEverything(ctx context.Context, fs *btrfs.FS) (AllSums, error) {
 						laddr := btrfsvol.LogicalAddr(item.Key.Offset) + (btrfsvol.LogicalAddr(i) * csumBlockSize)
 						if laddr != curAddr {
 							if curSums.Len() > 0 {
-								ret.Logical = append(ret.Logical, LogicalSumRun{
-									Addr: curAddr,
-									Sums: curSums.String(),
+								ret.Logical = append(ret.Logical, SumRun[btrfsvol.LogicalAddr]{
+									ChecksumSize: csumSize,
+									Addr:         curAddr,
+									Sums:         curSums.String(),
 								})
 							}
 							curAddr = laddr
 							curSums.Reset()
 						}
-						curSums.Write(sum[:ret.ChecksumSize])
+						curSums.Write(sum[:csumSize])
 					}
 					return nil
 				},
 			},
 		)
 		if curSums.Len() > 0 {
-			ret.Logical = append(ret.Logical, LogicalSumRun{
-				Addr: curAddr,
-				Sums: curSums.String(),
+			ret.Logical = append(ret.Logical, SumRun[btrfsvol.LogicalAddr]{
+				ChecksumSize: csumSize,
+				Addr:         curAddr,
+				Sums:         curSums.String(),
 			})
 		}
 	}()
@@ -165,7 +218,7 @@ func SumEverything(ctx context.Context, fs *btrfs.FS) (AllSums, error) {
 		devs := fs.LV.PhysicalVolumes()
 
 		var mu sync.Mutex
-		ret.Physical = make(map[btrfsvol.DeviceID]string, len(devs))
+		ret.Physical = make(map[btrfsvol.DeviceID]SumRun[btrfsvol.PhysicalAddr], len(devs))
 
 		grp := dgroup.NewGroup(ctx, dgroup.GroupConfig{})
 		for devID, dev := range devs {
@@ -173,7 +226,7 @@ func SumEverything(ctx context.Context, fs *btrfs.FS) (AllSums, error) {
 			grp.Go(dev.Name(), func(ctx context.Context) error {
 				devSize := dev.Size()
 				numSums := int(devSize / csumBlockSize)
-				sums := make([]byte, numSums*ret.ChecksumSize)
+				sums := make([]byte, numSums*csumSize)
 				lastPct := -1
 				progress := func(curSum int) {
 					pct := int(100 * float64(curSum) / float64(numSums))
@@ -192,12 +245,16 @@ func SumEverything(ctx context.Context, fs *btrfs.FS) (AllSums, error) {
 					if err != nil {
 						return err
 					}
-					copy(sums[i*ret.ChecksumSize:], sum[:ret.ChecksumSize])
+					copy(sums[i*csumSize:], sum[:csumSize])
 				}
 				progress(numSums)
 				sumsStr := string(sums)
 				mu.Lock()
-				ret.Physical[devID] = sumsStr
+				ret.Physical[devID] = SumRun[btrfsvol.PhysicalAddr]{
+					ChecksumSize: csumSize,
+					Addr:         0,
+					Sums:         sumsStr,
+				}
 				mu.Unlock()
 				return nil
 			})
@@ -213,55 +270,3 @@ func SumEverything(ctx context.Context, fs *btrfs.FS) (AllSums, error) {
 	// Return
 	return ret, nil
 }
-
-func ChecksumLogical(fs btrfs.Trees, alg btrfssum.CSumType, laddr btrfsvol.LogicalAddr) (btrfssum.CSum, error) {
-	var dat [csumBlockSize]byte
-	if _, err := fs.ReadAt(dat[:], laddr); err != nil {
-		return btrfssum.CSum{}, err
-	}
-	return alg.Sum(dat[:])
-}
-
-func ChecksumPhysical(dev *btrfs.Device, alg btrfssum.CSumType, paddr btrfsvol.PhysicalAddr) (btrfssum.CSum, error) {
-	var dat [csumBlockSize]byte
-	if _, err := dev.ReadAt(dat[:], paddr); err != nil {
-		return btrfssum.CSum{}, err
-	}
-	return alg.Sum(dat[:])
-}
-
-func ChecksumQualifiedPhysical(fs *btrfs.FS, alg btrfssum.CSumType, paddr btrfsvol.QualifiedPhysicalAddr) (btrfssum.CSum, error) {
-	dev := fs.LV.PhysicalVolumes()[paddr.Dev]
-	if dev == nil {
-		return btrfssum.CSum{}, fmt.Errorf("no such device_id=%v", paddr.Dev)
-	}
-	return ChecksumPhysical(dev, alg, paddr.Addr)
-}
-
-func LookupCSum(fs btrfs.Trees, alg btrfssum.CSumType, laddr btrfsvol.LogicalAddr) (map[btrfsvol.LogicalAddr]btrfssum.CSum, error) {
-	item, err := fs.TreeSearch(btrfs.CSUM_TREE_OBJECTID, func(key btrfs.Key, size uint32) int {
-		itemBeg := btrfsvol.LogicalAddr(key.ObjectID)
-		numSums := int64(size) / int64(alg.Size())
-		itemEnd := itemBeg + btrfsvol.LogicalAddr(numSums*csumBlockSize)
-		switch {
-		case itemEnd <= laddr:
-			return 1
-		case laddr < itemBeg:
-			return -1
-		default:
-			return 0
-		}
-	})
-	if err != nil {
-		return nil, err
-	}
-	body, ok := item.Body.(btrfsitem.ExtentCSum)
-	if !ok {
-		return nil, fmt.Errorf("item body is %T not ExtentCSum", item.Body)
-	}
-	ret := make(map[btrfsvol.LogicalAddr]btrfssum.CSum, len(body.Sums))
-	for i, sum := range body.Sums {
-		ret[btrfsvol.LogicalAddr(item.Key.ObjectID)+(btrfsvol.LogicalAddr(i)*csumBlockSize)] = sum
-	}
-	return ret, nil
-}
diff --git a/lib/btrfsprogs/btrfsinspect/scanforextents/csums_raw.go b/lib/btrfsprogs/btrfsinspect/scanforextents/csums_raw.go
new file mode 100644
index 0000000..14be110
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/scanforextents/csums_raw.go
@@ -0,0 +1,66 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package scanforextents
+
+import (
+	"fmt"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfssum"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+)
+
+func ChecksumLogical(fs btrfs.Trees, alg btrfssum.CSumType, laddr btrfsvol.LogicalAddr) (btrfssum.CSum, error) {
+	var dat [csumBlockSize]byte
+	if _, err := fs.ReadAt(dat[:], laddr); err != nil {
+		return btrfssum.CSum{}, err
+	}
+	return alg.Sum(dat[:])
+}
+
+func ChecksumPhysical(dev *btrfs.Device, alg btrfssum.CSumType, paddr btrfsvol.PhysicalAddr) (btrfssum.CSum, error) {
+	var dat [csumBlockSize]byte
+	if _, err := dev.ReadAt(dat[:], paddr); err != nil {
+		return btrfssum.CSum{}, err
+	}
+	return alg.Sum(dat[:])
+}
+
+func ChecksumQualifiedPhysical(fs *btrfs.FS, alg btrfssum.CSumType, paddr btrfsvol.QualifiedPhysicalAddr) (btrfssum.CSum, error) {
+	dev := fs.LV.PhysicalVolumes()[paddr.Dev]
+	if dev == nil {
+		return btrfssum.CSum{}, fmt.Errorf("no such device_id=%v", paddr.Dev)
+	}
+	return ChecksumPhysical(dev, alg, paddr.Addr)
+}
+
+func LookupCSum(fs btrfs.Trees, alg btrfssum.CSumType, laddr btrfsvol.LogicalAddr) (map[btrfsvol.LogicalAddr]btrfssum.CSum, error) {
+	item, err := fs.TreeSearch(btrfs.CSUM_TREE_OBJECTID, func(key btrfs.Key, size uint32) int {
+		itemBeg := btrfsvol.LogicalAddr(key.ObjectID)
+		numSums := int64(size) / int64(alg.Size())
+		itemEnd := itemBeg + btrfsvol.LogicalAddr(numSums*csumBlockSize)
+		switch {
+		case itemEnd <= laddr:
+			return 1
+		case laddr < itemBeg:
+			return -1
+		default:
+			return 0
+		}
+	})
+	if err != nil {
+		return nil, err
+	}
+	body, ok := item.Body.(btrfsitem.ExtentCSum)
+	if !ok {
+		return nil, fmt.Errorf("item body is %T not ExtentCSum", item.Body)
+	}
+	ret := make(map[btrfsvol.LogicalAddr]btrfssum.CSum, len(body.Sums))
+	for i, sum := range body.Sums {
+		ret[btrfsvol.LogicalAddr(item.Key.ObjectID)+(btrfsvol.LogicalAddr(i)*csumBlockSize)] = sum
+	}
+	return ret, nil
+}
diff --git a/lib/btrfsprogs/btrfsinspect/scanforextents/scan.go b/lib/btrfsprogs/btrfsinspect/scanforextents/scan.go
index 8920dfa..77a4ed9 100644
--- a/lib/btrfsprogs/btrfsinspect/scanforextents/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/scanforextents/scan.go
@@ -25,10 +25,12 @@ func ScanForExtents(ctx context.Context, fs *btrfs.FS, blockgroups *BlockGroupTr
 
 	dlog.Info(ctx, "Reverse-indexing and validating logical sums...")
 	var totalSums int
-	_ = sums.WalkLogical(func(btrfsvol.LogicalAddr, ShortSum) error {
+	if err := sums.WalkLogical(ctx, func(btrfsvol.LogicalAddr, ShortSum) error {
 		totalSums++
 		return nil
-	})
+	}); err != nil {
+		return err
+	}
 	sum2laddrs := make(map[ShortSum][]btrfsvol.LogicalAddr)
 	var curSum int
 	lastPct := -1
@@ -39,10 +41,7 @@ func ScanForExtents(ctx context.Context, fs *btrfs.FS, blockgroups *BlockGroupTr
 			lastPct = pct
 		}
 	}
-	if err := sums.WalkLogical(func(laddr btrfsvol.LogicalAddr, expShortSum ShortSum) error {
-		if err := ctx.Err(); err != nil {
-			return err
-		}
+	if err := sums.WalkLogical(ctx, func(laddr btrfsvol.LogicalAddr, expShortSum ShortSum) error {
 		progress(curSum)
 		curSum++
 		readSum, err := ChecksumLogical(fs, sb.ChecksumType, laddr)
@@ -53,7 +52,7 @@ func ScanForExtents(ctx context.Context, fs *btrfs.FS, blockgroups *BlockGroupTr
 			}
 			return err
 		}
-		readShortSum := ShortSum(readSum[:sums.ChecksumSize])
+		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))