Let btree search have access to the item size

4 files changed, 59 insertions, 50 deletions

diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index 553bf41..ca528d7 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"io"
 	iofs "io/fs"
+	"math"
 	"strings"
 
 	"github.com/datawire/dlib/derror"
@@ -45,14 +46,14 @@ type Trees interface {
 	TreeWalk(ctx context.Context, treeID ObjID, errHandle func(*TreeError), cbs TreeWalkHandler)
 
 	TreeLookup(treeID ObjID, key Key) (Item, error)
-	TreeSearch(treeID ObjID, fn func(Key) int) (Item, error)
+	TreeSearch(treeID ObjID, fn func(key Key, size uint32) int) (Item, error) // size is math.MaxUint32 for key-pointers
 
 	// If some items are able to be read, but there is an error reading the
 	// full set, then it might return *both* a list of items and an error.
 	//
 	// If no such item is found, an error that is io/fs.ErrNotExist is
 	// returned.
-	TreeSearchAll(treeID ObjID, fn func(Key) int) ([]Item, error)
+	TreeSearchAll(treeID ObjID, fn func(key Key, size uint32) int) ([]Item, error) // size is math.MaxUint32 for key-pointers
 
 	// For bootstrapping purposes.
 	Superblock() (*Superblock, error)
@@ -231,7 +232,7 @@ func LookupTreeRoot(fs Trees, treeID ObjID) (*TreeRoot, error) {
 			Generation: sb.BlockGroupRootGeneration,
 		}, nil
 	default:
-		rootItem, err := fs.TreeSearch(ROOT_TREE_OBJECTID, func(key Key) int {
+		rootItem, err := fs.TreeSearch(ROOT_TREE_OBJECTID, func(key Key, _ uint32) int {
 			if key.ObjectID == treeID && key.ItemType == btrfsitem.ROOT_ITEM_KEY {
 				return 0
 			}
@@ -402,7 +403,7 @@ func (fs *FS) treeWalk(ctx context.Context, path TreePath, errHandle func(*TreeE
 	}
 }
 
-func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key) int) (TreePath, *diskio.Ref[btrfsvol.LogicalAddr, Node], error) {
+func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key, uint32) int) (TreePath, *diskio.Ref[btrfsvol.LogicalAddr, Node], error) {
 	path := TreePath{
 		TreeID: treeRoot.TreeID,
 		Nodes: []TreePathElem{
@@ -437,7 +438,7 @@ func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key) int) (TreePath, *diskio
 			firstBad := len(node.Data.BodyInternal)
 			for firstBad > lastGood+1 {
 				midpoint := (lastGood + firstBad) / 2
-				direction := fn(node.Data.BodyInternal[midpoint].Key)
+				direction := fn(node.Data.BodyInternal[midpoint].Key, math.MaxUint32)
 				if direction < 0 {
 					firstBad = midpoint
 				} else {
@@ -469,7 +470,9 @@ func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key) int) (TreePath, *diskio
 			end := len(node.Data.BodyLeaf)
 			for beg < end {
 				midpoint := (beg + end) / 2
-				direction := fn(node.Data.BodyLeaf[midpoint].Key)
+				direction := fn(
+					node.Data.BodyLeaf[midpoint].Key,
+					node.Data.BodyLeaf[midpoint].BodySize)
 				switch {
 				case direction < 0:
 					end = midpoint
@@ -606,7 +609,7 @@ func (fs *FS) next(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 	return path, node, nil
 }
 
-func (fs *FS) TreeSearch(treeID ObjID, fn func(Key) int) (Item, error) {
+func (fs *FS) TreeSearch(treeID ObjID, fn func(Key, uint32) int) (Item, error) {
 	rootInfo, err := LookupTreeRoot(fs, treeID)
 	if err != nil {
 		return Item{}, err
@@ -618,15 +621,21 @@ func (fs *FS) TreeSearch(treeID ObjID, fn func(Key) int) (Item, error) {
 	return node.Data.BodyLeaf[path.Node(-1).ItemIdx], nil
 }
 
+func KeySearch(fn func(Key) int) func(Key, uint32) int {
+	return func(key Key, _ uint32) int {
+		return fn(key)
+	}
+}
+
 func (fs *FS) TreeLookup(treeID ObjID, key Key) (Item, error) {
-	item, err := fs.TreeSearch(treeID, key.Cmp)
+	item, err := fs.TreeSearch(treeID, KeySearch(key.Cmp))
 	if err != nil {
 		err = fmt.Errorf("item with key=%v: %w", key, err)
 	}
 	return item, err
 }
 
-func (fs *FS) TreeSearchAll(treeID ObjID, fn func(Key) int) ([]Item, error) {
+func (fs *FS) TreeSearchAll(treeID ObjID, fn func(Key, uint32) int) ([]Item, error) {
 	rootInfo, err := LookupTreeRoot(fs, treeID)
 	if err != nil {
 		return nil, err
@@ -649,7 +658,7 @@ func (fs *FS) TreeSearchAll(treeID ObjID, fn func(Key) int) ([]Item, error) {
 			break
 		}
 		prevItem := prevNode.Data.BodyLeaf[prevPath.Node(-1).ItemIdx]
-		if fn(prevItem.Key) != 0 {
+		if fn(prevItem.Key, prevItem.BodySize) != 0 {
 			break
 		}
 		ret = append(ret, prevItem)
@@ -665,7 +674,7 @@ func (fs *FS) TreeSearchAll(treeID ObjID, fn func(Key) int) ([]Item, error) {
 			break
 		}
 		nextItem := nextNode.Data.BodyLeaf[nextPath.Node(-1).ItemIdx]
-		if fn(nextItem.Key) != 0 {
+		if fn(nextItem.Key, nextItem.BodySize) != 0 {
 			break
 		}
 		ret = append(ret, nextItem)
diff --git a/lib/btrfs/io4_fs.go b/lib/btrfs/io4_fs.go
index 0ef922b..d28e36f 100644
--- a/lib/btrfs/io4_fs.go
+++ b/lib/btrfs/io4_fs.go
@@ -134,7 +134,7 @@ func (sv *Subvolume) LoadFullInode(inode ObjID) (*FullInode, error) {
 				Inode: inode,
 			},
 		}
-		items, err := sv.FS.TreeSearchAll(sv.TreeID, func(key Key) int {
+		items, err := sv.FS.TreeSearchAll(sv.TreeID, func(key Key, _ uint32) int {
 			return containers.NativeCmp(inode, key.ObjectID)
 		})
 		if err != nil {
diff --git a/lib/btrfs/types_node.go b/lib/btrfs/types_node.go
index 781b2fa..840236d 100644
--- a/lib/btrfs/types_node.go
+++ b/lib/btrfs/types_node.go
@@ -247,8 +247,9 @@ func (node *Node) marshalInternalTo(bodyBuf []byte) error {
 // Node: "leaf" ////////////////////////////////////////////////////////////////////////////////////
 
 type Item struct {
-	Key  Key
-	Body btrfsitem.Item
+	Key      Key
+	BodySize uint32 // [ignored-when-writing]
+	Body     btrfsitem.Item
 }
 
 type ItemHeader struct {
@@ -287,8 +288,9 @@ func (node *Node) unmarshalLeaf(bodyBuf []byte) (int, error) {
 		dataBuf := bodyBuf[dataOff : dataOff+dataSize]
 
 		node.BodyLeaf = append(node.BodyLeaf, Item{
-			Key:  itemHead.Key,
-			Body: btrfsitem.UnmarshalItem(itemHead.Key, node.ChecksumType, dataBuf),
+			Key:      itemHead.Key,
+			BodySize: itemHead.DataSize,
+			Body:     btrfsitem.UnmarshalItem(itemHead.Key, node.ChecksumType, dataBuf),
 		})
 	}
 
diff --git a/lib/btrfsprogs/btrfsutil/broken_btree.go b/lib/btrfsprogs/btrfsutil/broken_btree.go
index 93bb53d..7fc812d 100644
--- a/lib/btrfsprogs/btrfsutil/broken_btree.go
+++ b/lib/btrfsprogs/btrfsutil/broken_btree.go
@@ -21,13 +21,10 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
-type indexValue[T any] struct {
-	Key btrfs.Key
-	Val T
-}
-
-func (v indexValue[T]) keyFn() btrfs.Key {
-	return v.Key
+type indexValue struct {
+	Key      btrfs.Key
+	ItemSize uint32
+	Path     btrfs.TreePath
 }
 
 var maxKey = btrfs.Key{
@@ -88,7 +85,7 @@ type spanError struct {
 
 type cachedIndex struct {
 	TreeRootErr error
-	Items       *containers.RBTree[btrfs.Key, indexValue[btrfs.TreePath]]
+	Items       *containers.RBTree[btrfs.Key, indexValue]
 	Errors      map[int]map[btrfs.Key][]spanError
 }
 
@@ -157,8 +154,8 @@ func (bt *brokenTrees) treeIndex(treeID btrfs.ObjID) cachedIndex {
 	if err != nil {
 		cacheEntry.TreeRootErr = err
 	} else {
-		cacheEntry.Items = &containers.RBTree[btrfs.Key, indexValue[btrfs.TreePath]]{
-			KeyFn: indexValue[btrfs.TreePath].keyFn,
+		cacheEntry.Items = &containers.RBTree[btrfs.Key, indexValue]{
+			KeyFn: func(iv indexValue) btrfs.Key { return iv.Key },
 		}
 		dlog.Infof(bt.ctx, "indexing tree %v...", treeID)
 		bt.inner.RawTreeWalk(
@@ -190,9 +187,10 @@ func (bt *brokenTrees) treeIndex(treeID btrfs.ObjID) cachedIndex {
 						// and force me to figure out how to handle it.
 						panic(fmt.Errorf("dup key=%v in tree=%v", item.Key, treeID))
 					}
-					cacheEntry.Items.Insert(indexValue[btrfs.TreePath]{
-						Key: item.Key,
-						Val: path.DeepCopy(),
+					cacheEntry.Items.Insert(indexValue{
+						Key:      item.Key,
+						ItemSize: item.BodySize,
+						Path:     path.DeepCopy(),
 					})
 					return nil
 				},
@@ -209,42 +207,42 @@ func (bt *brokenTrees) treeIndex(treeID btrfs.ObjID) cachedIndex {
 }
 
 func (bt *brokenTrees) TreeLookup(treeID btrfs.ObjID, key btrfs.Key) (btrfs.Item, error) {
-	item, err := bt.TreeSearch(treeID, key.Cmp)
+	item, err := bt.TreeSearch(treeID, btrfs.KeySearch(key.Cmp))
 	if err != nil {
 		err = fmt.Errorf("item with key=%v: %w", key, err)
 	}
 	return item, err
 }
 
-func (bt *brokenTrees) TreeSearch(treeID btrfs.ObjID, fn func(btrfs.Key) int) (btrfs.Item, error) {
+func (bt *brokenTrees) TreeSearch(treeID btrfs.ObjID, fn func(btrfs.Key, uint32) int) (btrfs.Item, error) {
 	index := bt.treeIndex(treeID)
 	if index.TreeRootErr != nil {
 		return btrfs.Item{}, index.TreeRootErr
 	}
-	indexItem := index.Items.Search(func(indexItem indexValue[btrfs.TreePath]) int {
-		return fn(indexItem.Key)
+	indexItem := index.Items.Search(func(indexItem indexValue) int {
+		return fn(indexItem.Key, indexItem.ItemSize)
 	})
 	if indexItem == nil {
 		return btrfs.Item{}, iofs.ErrNotExist
 	}
 
-	node, err := bt.inner.ReadNode(indexItem.Value.Val.Node(-2).NodeAddr)
+	node, err := bt.inner.ReadNode(indexItem.Value.Path.Node(-2).NodeAddr)
 	if err != nil {
 		return btrfs.Item{}, err
 	}
 
-	item := node.Data.BodyLeaf[indexItem.Value.Val.Node(-1).ItemIdx]
+	item := node.Data.BodyLeaf[indexItem.Value.Path.Node(-1).ItemIdx]
 
 	return item, nil
 }
 
-func (bt *brokenTrees) TreeSearchAll(treeID btrfs.ObjID, fn func(btrfs.Key) int) ([]btrfs.Item, error) {
+func (bt *brokenTrees) TreeSearchAll(treeID btrfs.ObjID, fn func(btrfs.Key, uint32) int) ([]btrfs.Item, error) {
 	index := bt.treeIndex(treeID)
 	if index.TreeRootErr != nil {
 		return nil, index.TreeRootErr
 	}
-	indexItems := index.Items.SearchRange(func(indexItem indexValue[btrfs.TreePath]) int {
-		return fn(indexItem.Key)
+	indexItems := index.Items.SearchRange(func(indexItem indexValue) int {
+		return fn(indexItem.Key, indexItem.ItemSize)
 	})
 	if len(indexItems) == 0 {
 		return nil, iofs.ErrNotExist
@@ -253,26 +251,26 @@ func (bt *brokenTrees) TreeSearchAll(treeID btrfs.ObjID, fn func(btrfs.Key) int)
 	ret := make([]btrfs.Item, len(indexItems))
 	var node *diskio.Ref[btrfsvol.LogicalAddr, btrfs.Node]
 	for i := range indexItems {
-		if node == nil || node.Addr != indexItems[i].Val.Node(-2).NodeAddr {
+		if node == nil || node.Addr != indexItems[i].Path.Node(-2).NodeAddr {
 			var err error
-			node, err = bt.inner.ReadNode(indexItems[i].Val.Node(-2).NodeAddr)
+			node, err = bt.inner.ReadNode(indexItems[i].Path.Node(-2).NodeAddr)
 			if err != nil {
 				return nil, err
 			}
 		}
-		ret[i] = node.Data.BodyLeaf[indexItems[i].Val.Node(-1).ItemIdx]
+		ret[i] = node.Data.BodyLeaf[indexItems[i].Path.Node(-1).ItemIdx]
 	}
 
 	var errs derror.MultiError
 	for _, invLvl := range maps.SortedKeys(index.Errors) {
 		for _, beg := range maps.Keys(index.Errors[invLvl]) {
-			if fn(beg) < 0 {
+			if fn(beg, math.MaxUint32) < 0 {
 				continue
 			}
 			for _, spanErr := range index.Errors[invLvl][beg] {
 				end := spanErr.End
 				err := spanErr.Err
-				if fn(end) > 0 {
+				if fn(end, math.MaxUint32) > 0 {
 					continue
 				}
 				errs = append(errs, err)
@@ -298,7 +296,7 @@ func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfs.ObjID, errHand
 		return
 	}
 	var node *diskio.Ref[btrfsvol.LogicalAddr, btrfs.Node]
-	_ = index.Items.Walk(func(indexItem *containers.RBNode[indexValue[btrfs.TreePath]]) error {
+	_ = index.Items.Walk(func(indexItem *containers.RBNode[indexValue]) error {
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
@@ -306,17 +304,17 @@ func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfs.ObjID, errHand
 			return bt.ctx.Err()
 		}
 		if cbs.Item != nil {
-			if node == nil || node.Addr != indexItem.Value.Val.Node(-2).NodeAddr {
+			if node == nil || node.Addr != indexItem.Value.Path.Node(-2).NodeAddr {
 				var err error
-				node, err = bt.inner.ReadNode(indexItem.Value.Val.Node(-2).NodeAddr)
+				node, err = bt.inner.ReadNode(indexItem.Value.Path.Node(-2).NodeAddr)
 				if err != nil {
-					errHandle(&btrfs.TreeError{Path: indexItem.Value.Val, Err: err})
+					errHandle(&btrfs.TreeError{Path: indexItem.Value.Path, Err: err})
 					return nil
 				}
 			}
-			item := node.Data.BodyLeaf[indexItem.Value.Val.Node(-1).ItemIdx]
-			if err := cbs.Item(indexItem.Value.Val, item); err != nil {
-				errHandle(&btrfs.TreeError{Path: indexItem.Value.Val, Err: err})
+			item := node.Data.BodyLeaf[indexItem.Value.Path.Node(-1).ItemIdx]
+			if err := cbs.Item(indexItem.Value.Path, item); err != nil {
+				errHandle(&btrfs.TreeError{Path: indexItem.Value.Path, Err: err})
 			}
 		}
 		return nil