Add an AddrDelta type for safer arithmentic

7 files changed, 64 insertions, 48 deletions

diff --git a/cmd/btrfs-fsck/pass1.go b/cmd/btrfs-fsck/pass1.go
index f82ec03..4507313 100644
--- a/cmd/btrfs-fsck/pass1.go
+++ b/cmd/btrfs-fsck/pass1.go
@@ -41,7 +41,7 @@ func pass1(fs *btrfs.FS, superblock *util.Ref[btrfs.PhysicalAddr, btrfs.Superblo
 
 	fsFoundNodes := make(map[btrfs.LogicalAddr]struct{})
 	fsReconstructedChunks := make(map[btrfs.LogicalAddr]struct {
-		Size    uint64
+		Size    btrfs.AddrDelta
 		Stripes []btrfsitem.ChunkStripe
 	})
 	for _, dev := range fs.Devices {
@@ -221,7 +221,7 @@ func pass1ReconstructChunksOneDev(
 	foundNodes map[btrfs.LogicalAddr][]btrfs.PhysicalAddr,
 ) (
 	chunks map[btrfs.LogicalAddr]struct {
-		Size    uint64
+		Size    btrfs.AddrDelta
 		Stripes []btrfsitem.ChunkStripe
 	},
 ) {
@@ -254,7 +254,7 @@ func pass1ReconstructChunksOneDev(
 	type stripe struct {
 		PAddr btrfs.PhysicalAddr
 		LAddr btrfs.LogicalAddr
-		Size  uint64
+		Size  btrfs.AddrDelta
 	}
 	var stripes []*stripe
 	for _, paddr := range sortedPAddrs {
@@ -265,14 +265,14 @@ func pass1ReconstructChunksOneDev(
 			lastStripe = stripes[len(stripes)-1]
 		}
 		if lastStripe != nil &&
-			paddr == lastStripe.PAddr+btrfs.PhysicalAddr(lastStripe.Size) &&
-			laddr == lastStripe.LAddr+btrfs.LogicalAddr(lastStripe.Size) {
-			lastStripe.Size += uint64(superblock.Data.NodeSize)
+			paddr == lastStripe.PAddr.Add(lastStripe.Size) &&
+			laddr == lastStripe.LAddr.Add(lastStripe.Size) {
+			lastStripe.Size += btrfs.AddrDelta(superblock.Data.NodeSize)
 		} else {
 			stripes = append(stripes, &stripe{
 				PAddr: paddr,
 				LAddr: laddr,
-				Size:  uint64(superblock.Data.NodeSize),
+				Size:  btrfs.AddrDelta(superblock.Data.NodeSize),
 			})
 		}
 	}
@@ -280,7 +280,7 @@ func pass1ReconstructChunksOneDev(
 
 	// organize those stripes in to chunks
 	chunks = make(map[btrfs.LogicalAddr]struct {
-		Size    uint64
+		Size    btrfs.AddrDelta
 		Stripes []btrfsitem.ChunkStripe
 	})
 	for _, stripe := range stripes {
@@ -303,7 +303,7 @@ func pass1ReconstructChunksOneDev(
 }
 
 func pass1PrintChunks(chunks map[btrfs.LogicalAddr]struct {
-	Size    uint64
+	Size    btrfs.AddrDelta
 	Stripes []btrfsitem.ChunkStripe
 }) {
 	laddrs := make([]btrfs.LogicalAddr, 0, len(chunks))
@@ -313,8 +313,8 @@ func pass1PrintChunks(chunks map[btrfs.LogicalAddr]struct {
 	sort.Slice(laddrs, func(i, j int) bool {
 		return laddrs[i] < laddrs[j]
 	})
-	lprev := btrfs.LogicalAddr(0)
-	pprev := btrfs.PhysicalAddr(0)
+	var lprev btrfs.LogicalAddr
+	var pprev btrfs.PhysicalAddr
 	for _, laddr := range laddrs {
 		ldelta := laddr - lprev
 		chunk := chunks[laddr]
@@ -327,25 +327,25 @@ func pass1PrintChunks(chunks map[btrfs.LogicalAddr]struct {
 			fmt.Printf("chunkstripe: laddr=%v (+%v) => paddr=%v (+%v) ; size=%v (%s)\n",
 				laddr, ldelta,
 				stripe.Offset, pdelta,
-				btrfs.PhysicalAddr(chunk.Size),
+				chunk.Size,
 				adj)
-			pprev = stripe.Offset + btrfs.PhysicalAddr(chunk.Size)
+			pprev = stripe.Offset.Add(chunk.Size)
 		}
-		lprev = laddr + btrfs.LogicalAddr(chunk.Size)
+		lprev = laddr.Add(chunk.Size)
 	}
 }
 
 func pass1ProcessBlockGroups(blockgroups []sysBlockGroup) {
 	// organize in to a more manageable datastructure
 	type groupAttrs struct {
-		Size  btrfs.LogicalAddr
+		Size  btrfs.AddrDelta
 		Flags btrfsitem.BlockGroupFlags
 	}
 	groups := make(map[btrfs.LogicalAddr]groupAttrs)
 	for _, bg := range blockgroups {
 		laddr := btrfs.LogicalAddr(bg.Key.ObjectID)
 		attrs := groupAttrs{
-			Size:  btrfs.LogicalAddr(bg.Key.Offset),
+			Size:  btrfs.AddrDelta(bg.Key.Offset),
 			Flags: bg.BG.Flags,
 		}
 		// If there's a conflict, but they both say the same thing (existing == attrs),
@@ -369,7 +369,7 @@ func pass1ProcessBlockGroups(blockgroups []sysBlockGroup) {
 	// cluster
 	type cluster struct {
 		LAddr btrfs.LogicalAddr
-		Size  btrfs.LogicalAddr
+		Size  btrfs.AddrDelta
 		Flags btrfsitem.BlockGroupFlags
 	}
 	var clusters []*cluster
@@ -380,7 +380,7 @@ func pass1ProcessBlockGroups(blockgroups []sysBlockGroup) {
 		if len(clusters) > 0 {
 			lastCluster = clusters[len(clusters)-1]
 		}
-		if lastCluster != nil && laddr == lastCluster.LAddr+lastCluster.Size && attrs.Flags == lastCluster.Flags {
+		if lastCluster != nil && laddr == lastCluster.LAddr.Add(lastCluster.Size) && attrs.Flags == lastCluster.Flags {
 			lastCluster.Size += attrs.Size
 		} else {
 			clusters = append(clusters, &cluster{
@@ -392,12 +392,12 @@ func pass1ProcessBlockGroups(blockgroups []sysBlockGroup) {
 	}
 
 	// print
-	prev := btrfs.LogicalAddr(0)
+	var prev btrfs.LogicalAddr
 	for _, cluster := range clusters {
 		delta := cluster.LAddr - prev
 		fmt.Printf("blockgroup cluster: laddr=%v (+%v); size=%v ; flags=%v\n",
 			cluster.LAddr, delta, cluster.Size, cluster.Flags)
-		prev = cluster.LAddr + cluster.Size
+		prev = cluster.LAddr.Add(cluster.Size)
 	}
 }
 
@@ -405,7 +405,7 @@ func pass1ProcessDevExtents(devextents []sysDevExtent) {
 	// organize in to a more manageable datastructure
 	type extAttrs struct {
 		LAddr btrfs.LogicalAddr
-		Size  uint64
+		Size  btrfs.AddrDelta
 	}
 	exts := make(map[btrfs.PhysicalAddr]extAttrs)
 	for _, de := range devextents {
@@ -436,7 +436,7 @@ func pass1ProcessDevExtents(devextents []sysDevExtent) {
 	type stripe struct {
 		PAddr btrfs.PhysicalAddr
 		LAddr btrfs.LogicalAddr
-		Size  uint64
+		Size  btrfs.AddrDelta
 	}
 	var stripes []*stripe
 	for _, paddr := range sortedPAddrs {
@@ -447,8 +447,8 @@ func pass1ProcessDevExtents(devextents []sysDevExtent) {
 			lastStripe = stripes[len(stripes)-1]
 		}
 		if lastStripe != nil &&
-			paddr == lastStripe.PAddr+btrfs.PhysicalAddr(lastStripe.Size) &&
-			attrs.LAddr == lastStripe.LAddr+btrfs.LogicalAddr(lastStripe.Size) {
+			paddr == lastStripe.PAddr.Add(lastStripe.Size) &&
+			attrs.LAddr == lastStripe.LAddr.Add(lastStripe.Size) {
 			lastStripe.Size += attrs.Size
 		} else {
 			stripes = append(stripes, &stripe{
@@ -460,17 +460,17 @@ func pass1ProcessDevExtents(devextents []sysDevExtent) {
 	}
 
 	// print
-	lprev := btrfs.LogicalAddr(0)
-	pprev := btrfs.PhysicalAddr(0)
+	var lprev btrfs.LogicalAddr
+	var pprev btrfs.PhysicalAddr
 	for _, stripe := range stripes {
 		pdelta := stripe.PAddr - pprev
 		ldelta := stripe.LAddr - lprev
 		fmt.Printf("devextent cluster: paddr=%v (+%v) => laddr=%v (+%v) ; size=%v\n",
 			stripe.PAddr, pdelta,
 			stripe.LAddr, ldelta,
-			btrfs.PhysicalAddr(stripe.Size))
-		pprev = stripe.PAddr + btrfs.PhysicalAddr(stripe.Size)
-		lprev = stripe.LAddr + btrfs.LogicalAddr(stripe.Size)
+			stripe.Size)
+		pprev = stripe.PAddr.Add(stripe.Size)
+		lprev = stripe.LAddr.Add(stripe.Size)
 	}
 }
 
@@ -478,7 +478,7 @@ func pass1WriteReconstructedChunks(
 	fs *btrfs.FS,
 	superblock btrfs.Superblock,
 	fsReconstructedChunks map[btrfs.LogicalAddr]struct {
-		Size    uint64
+		Size    btrfs.AddrDelta
 		Stripes []btrfsitem.ChunkStripe
 	},
 ) {
@@ -530,7 +530,7 @@ func pass1WriteReconstructedChunks(
 		chunk := fsReconstructedChunks[laddr]
 		for j, stripe := range chunk.Stripes {
 			fmt.Printf("Pass 1: chunks[%v].stripes[%v] = { laddr=%v => { dev_id=%v, paddr=%v }, size=%v }\n",
-				i, j, laddr, stripe.DeviceID, stripe.Offset, btrfs.LogicalAddr(chunk.Size))
+				i, j, laddr, stripe.DeviceID, stripe.Offset, chunk.Size)
 		}
 		reconstructedNode.Data.BodyLeaf = append(reconstructedNode.Data.BodyLeaf, btrfs.Item{
 			Head: btrfs.ItemHeader{
diff --git a/pkg/btrfs/btrfsitem/item_blockgroup.go b/pkg/btrfs/btrfsitem/item_blockgroup.go
index a28d4bf..737cfa0 100644
--- a/pkg/btrfs/btrfsitem/item_blockgroup.go
+++ b/pkg/btrfs/btrfsitem/item_blockgroup.go
@@ -6,6 +6,8 @@ import (
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
+// key.objectid = logical_addr
+// key.offset = size of chunk
 type BlockGroup struct { // BLOCK_GROUP_ITEM=192
 	Used          int64           `bin:"off=0, siz=8"`
 	ChunkObjectID internal.ObjID  `bin:"off=8, siz=8"`
diff --git a/pkg/btrfs/btrfsitem/item_chunk.go b/pkg/btrfs/btrfsitem/item_chunk.go
index eb3a16d..8f98ba1 100644
--- a/pkg/btrfs/btrfsitem/item_chunk.go
+++ b/pkg/btrfs/btrfsitem/item_chunk.go
@@ -8,21 +8,24 @@ import (
 )
 
 // Maps logical address to physical.
+//
+// key.objectid = BTRFS_FIRST_CHUNK_TREE_OBJECTID
+// key.offset = logical_addr
 type Chunk struct { // CHUNK_ITEM=228
 	Head    ChunkHeader
 	Stripes []ChunkStripe
 }
 
 type ChunkHeader struct {
-	Size           uint64          `bin:"off=0x0,  siz=0x8"`
-	Owner          internal.ObjID  `bin:"off=0x8,  siz=0x8"` // root referencing this chunk (always EXTENT_TREE_OBJECTID=2)
-	StripeLen      uint64          `bin:"off=0x10, siz=0x8"` // ???
-	Type           BlockGroupFlags `bin:"off=0x18, siz=0x8"`
-	IOOptimalAlign uint32          `bin:"off=0x20, siz=0x4"`
-	IOOptimalWidth uint32          `bin:"off=0x24, siz=0x4"`
-	IOMinSize      uint32          `bin:"off=0x28, siz=0x4"` // sector size
-	NumStripes     uint16          `bin:"off=0x2c, siz=0x2"` // [ignored-when-writing]
-	SubStripes     uint16          `bin:"off=0x2e, siz=0x2"` // ???
+	Size           internal.AddrDelta `bin:"off=0x0,  siz=0x8"`
+	Owner          internal.ObjID     `bin:"off=0x8,  siz=0x8"` // root referencing this chunk (always EXTENT_TREE_OBJECTID=2)
+	StripeLen      uint64             `bin:"off=0x10, siz=0x8"` // ???
+	Type           BlockGroupFlags    `bin:"off=0x18, siz=0x8"`
+	IOOptimalAlign uint32             `bin:"off=0x20, siz=0x4"`
+	IOOptimalWidth uint32             `bin:"off=0x24, siz=0x4"`
+	IOMinSize      uint32             `bin:"off=0x28, siz=0x4"` // sector size
+	NumStripes     uint16             `bin:"off=0x2c, siz=0x2"` // [ignored-when-writing]
+	SubStripes     uint16             `bin:"off=0x2e, siz=0x2"` // ???
 	binstruct.End  `bin:"off=0x30"`
 }
 
diff --git a/pkg/btrfs/btrfsitem/item_devextent.go b/pkg/btrfs/btrfsitem/item_devextent.go
index dfa6d03..7f08b5e 100644
--- a/pkg/btrfs/btrfsitem/item_devextent.go
+++ b/pkg/btrfs/btrfsitem/item_devextent.go
@@ -5,11 +5,13 @@ import (
 	"lukeshu.com/btrfs-tools/pkg/btrfs/internal"
 )
 
+// key.objectid = device_id
+// key.offset = physical_addr
 type DevExtent struct { // DEV_EXTENT=204
 	ChunkTree     int64                `bin:"off=0, siz=8"`
 	ChunkObjectID internal.ObjID       `bin:"off=8, siz=8"`
 	ChunkOffset   internal.LogicalAddr `bin:"off=16, siz=8"`
-	Length        uint64               `bin:"off=24, siz=8"`
+	Length        internal.AddrDelta   `bin:"off=24, siz=8"`
 	ChunkTreeUUID internal.UUID        `bin:"off=32, siz=16"`
 	binstruct.End `bin:"off=48"`
 }
diff --git a/pkg/btrfs/internal.go b/pkg/btrfs/internal.go
index 5c0fc8b..e00e946 100644
--- a/pkg/btrfs/internal.go
+++ b/pkg/btrfs/internal.go
@@ -11,6 +11,7 @@ type (
 	ObjID        = internal.ObjID
 	LogicalAddr  = internal.LogicalAddr
 	PhysicalAddr = internal.PhysicalAddr
+	AddrDelta    = internal.AddrDelta
 
 	// complex types
 
diff --git a/pkg/btrfs/internal/addr.go b/pkg/btrfs/internal/addr.go
index 7067982..a3f41ed 100644
--- a/pkg/btrfs/internal/addr.go
+++ b/pkg/btrfs/internal/addr.go
@@ -9,6 +9,7 @@ import (
 type (
 	PhysicalAddr int64
 	LogicalAddr  int64
+	AddrDelta    int64
 )
 
 func formatAddr(addr int64, f fmt.State, verb rune) {
@@ -23,3 +24,10 @@ func formatAddr(addr int64, f fmt.State, verb rune) {
 
 func (a PhysicalAddr) Format(f fmt.State, verb rune) { formatAddr(int64(a), f, verb) }
 func (a LogicalAddr) Format(f fmt.State, verb rune)  { formatAddr(int64(a), f, verb) }
+func (d AddrDelta) Format(f fmt.State, verb rune)    { formatAddr(int64(d), f, verb) }
+
+func (a PhysicalAddr) Sub(b PhysicalAddr) AddrDelta { return AddrDelta(a - b) }
+func (a LogicalAddr) Sub(b LogicalAddr) AddrDelta   { return AddrDelta(a - b) }
+
+func (a PhysicalAddr) Add(b AddrDelta) PhysicalAddr { return a + PhysicalAddr(b) }
+func (a LogicalAddr) Add(b AddrDelta) LogicalAddr   { return a + LogicalAddr(b) }
diff --git a/pkg/btrfs/io2_fs.go b/pkg/btrfs/io2_fs.go
index 588a60a..8b76e2d 100644
--- a/pkg/btrfs/io2_fs.go
+++ b/pkg/btrfs/io2_fs.go
@@ -150,20 +150,20 @@ type QualifiedPhysicalAddr struct {
 	Addr PhysicalAddr
 }
 
-func (fs *FS) Resolve(laddr LogicalAddr) (paddrs map[QualifiedPhysicalAddr]struct{}, maxlen uint64) {
+func (fs *FS) Resolve(laddr LogicalAddr) (paddrs map[QualifiedPhysicalAddr]struct{}, maxlen AddrDelta) {
 	paddrs = make(map[QualifiedPhysicalAddr]struct{})
-	maxlen = math.MaxUint64
+	maxlen = math.MaxInt64
 
 	for _, chunk := range fs.chunks {
 		low := LogicalAddr(chunk.Key.Offset)
-		high := low + LogicalAddr(chunk.Chunk.Head.Size)
+		high := low.Add(chunk.Chunk.Head.Size)
 		if low <= laddr && laddr < high {
-			offsetWithinChunk := uint64(laddr) - chunk.Key.Offset
+			offsetWithinChunk := laddr.Sub(low)
 			maxlen = util.Min(maxlen, chunk.Chunk.Head.Size-offsetWithinChunk)
 			for _, stripe := range chunk.Chunk.Stripes {
 				paddrs[QualifiedPhysicalAddr{
 					Dev:  stripe.DeviceUUID,
-					Addr: stripe.Offset + PhysicalAddr(offsetWithinChunk),
+					Addr: stripe.Offset.Add(offsetWithinChunk),
 				}] = struct{}{}
 			}
 		}
@@ -189,7 +189,7 @@ func (fs *FS) maybeShortReadAt(dat []byte, laddr LogicalAddr) (int, error) {
 	if len(paddrs) == 0 {
 		return 0, fmt.Errorf("read: could not map logical address %v", laddr)
 	}
-	if uint64(len(dat)) > maxlen {
+	if AddrDelta(len(dat)) > maxlen {
 		dat = dat[:maxlen]
 	}
 
@@ -231,7 +231,7 @@ func (fs *FS) maybeShortWriteAt(dat []byte, laddr LogicalAddr) (int, error) {
 	if len(paddrs) == 0 {
 		return 0, fmt.Errorf("write: could not map logical address %v", laddr)
 	}
-	if uint64(len(dat)) > maxlen {
+	if AddrDelta(len(dat)) > maxlen {
 		dat = dat[:maxlen]
 	}