btrfstree: Have errors include context of what was being searched for

4 files changed, 232 insertions, 38 deletions

diff --git a/lib/btrfs/btrfstree/btree.go b/lib/btrfs/btrfstree/btree.go
index 7b3721b..4c10ffa 100644
--- a/lib/btrfs/btrfstree/btree.go
+++ b/lib/btrfs/btrfstree/btree.go
@@ -13,6 +13,15 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 )
 
+type TreeSearcher interface {
+	// How the search should be described in the event of an
+	// error.
+	fmt.Stringer
+
+	// size is math.MaxUint32 for key-pointers
+	Search(key btrfsprim.Key, size uint32) int
+}
+
 // TreeOperator is an interface for performing basic btree operations.
 type TreeOperator interface {
 	// TreeWalk walks a tree, triggering callbacks for every node,
@@ -40,7 +49,7 @@ type TreeOperator interface {
 	TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, errHandle func(*TreeError), cbs TreeWalkHandler)
 
 	TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key) (Item, error)
-	TreeSearch(treeID btrfsprim.ObjID, fn func(key btrfsprim.Key, size uint32) int) (Item, error) // size is math.MaxUint32 for key-pointers
+	TreeSearch(treeID btrfsprim.ObjID, search TreeSearcher) (Item, error)
 
 	// 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.
@@ -50,7 +59,7 @@ type TreeOperator interface {
 	//
 	// If no such item is found, an error that is ErrNoItem is
 	// returned.
-	TreeSearchAll(treeID btrfsprim.ObjID, fn func(key btrfsprim.Key, size uint32) int) ([]Item, error) // size is math.MaxUint32 for key-pointers
+	TreeSearchAll(treeID btrfsprim.ObjID, search TreeSearcher) ([]Item, error)
 }
 
 type TreeWalkHandler struct {
diff --git a/lib/btrfs/btrfstree/btree_forrest.go b/lib/btrfs/btrfstree/btree_forrest.go
index 625fb30..0f46d42 100644
--- a/lib/btrfs/btrfstree/btree_forrest.go
+++ b/lib/btrfs/btrfstree/btree_forrest.go
@@ -22,19 +22,6 @@ type TreeRoot struct {
 	Generation btrfsprim.Generation
 }
 
-func RootItemSearchFn(treeID btrfsprim.ObjID) func(btrfsprim.Key, uint32) int {
-	return func(key btrfsprim.Key, _ uint32) int {
-		if key.ObjectID == treeID && key.ItemType == btrfsitem.ROOT_ITEM_KEY {
-			return 0
-		}
-		return btrfsprim.Key{
-			ObjectID: treeID,
-			ItemType: btrfsitem.ROOT_ITEM_KEY,
-			Offset:   0,
-		}.Compare(key)
-	}
-}
-
 // LookupTreeRoot is a utility function to help with implementing the
 // 'TreeOperator' interface.
 func LookupTreeRoot(fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*TreeRoot, error) {
@@ -68,12 +55,12 @@ func LookupTreeRoot(fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*Tr
 			Generation: sb.BlockGroupRootGeneration,
 		}, nil
 	default:
-		rootItem, err := fs.TreeSearch(btrfsprim.ROOT_TREE_OBJECTID, RootItemSearchFn(treeID))
+		rootItem, err := fs.TreeSearch(btrfsprim.ROOT_TREE_OBJECTID, SearchRootItem(treeID))
 		if err != nil {
 			if errors.Is(err, ErrNoItem) {
 				err = ErrNoTree
 			}
-			return nil, err
+			return nil, fmt.Errorf("tree %s: %w", treeID.Format(btrfsprim.ROOT_TREE_OBJECTID), err)
 		}
 		switch rootItemBody := rootItem.Body.(type) {
 		case *btrfsitem.Root:
diff --git a/lib/btrfs/btrfstree/btree_searchers.go b/lib/btrfs/btrfstree/btree_searchers.go
new file mode 100644
index 0000000..6972e43
--- /dev/null
+++ b/lib/btrfs/btrfstree/btree_searchers.go
@@ -0,0 +1,209 @@
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfstree
+
+import (
+	"fmt"
+	"strings"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"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/containers"
+)
+
+type SearchItemType int8
+
+const (
+	ItemTypeAny SearchItemType = iota
+	ItemTypeExact
+)
+
+type SearchOffset int8
+
+const (
+	OffsetAny SearchOffset = iota
+	OffsetExact
+	OffsetRange // .Search behaves same as OffsetAny (TODO?)
+	OffsetName  // .Search behaves same as OffsetAny
+)
+
+// Search is a fairly generic and reusable implementation of
+// TreeSearcher.
+type Search struct {
+	ObjectID btrfsprim.ObjID
+
+	ItemTypeMatching SearchItemType
+	ItemType         btrfsprim.ItemType
+
+	// Offset is totally ignored if .ItemTypeMatching=ItemTypeany.
+	OffsetMatching SearchOffset
+	OffsetLow      uint64 // only for .OffsetMatching==OffsetExact or .OffsetMatching==OffsetRange
+	OffsetHigh     uint64 // only for .OffsetMatching==OffsetRange
+	OffsetName     string // only for .OffsetMatching==OffsetName
+}
+
+var (
+	_ containers.Ordered[Search] = Search{}
+	_ TreeSearcher               = Search{}
+)
+
+// Compare implements containers.Ordered.
+func (a Search) Compare(b Search) int {
+	if d := containers.NativeCompare(a.ObjectID, b.ObjectID); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.ItemType, b.ItemType); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.OffsetMatching, b.OffsetMatching); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.OffsetLow, b.OffsetLow); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.OffsetHigh, b.OffsetHigh); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.OffsetName, b.OffsetName); d != 0 {
+		return d
+	}
+	return 0
+}
+
+// String implements fmt.Stringer (and TreeSearcher).
+func (o Search) String() string {
+	var buf strings.Builder
+
+	fmt.Fprintf(&buf, "(%v ", o.ObjectID)
+
+	switch o.ItemTypeMatching {
+	case ItemTypeAny:
+		buf.WriteString("? ?)")
+		return buf.String()
+	case ItemTypeExact:
+		fmt.Fprintf(&buf, "%v", o.ItemType)
+	default:
+		panic(fmt.Errorf("should not happen: ItemTypeMatching=%#v", o.ItemTypeMatching))
+	}
+
+	buf.WriteString(" ")
+
+	switch o.OffsetMatching {
+	case OffsetAny:
+		buf.WriteString("?")
+	case OffsetExact:
+		fmt.Fprintf(&buf, "%v", o.OffsetLow)
+	case OffsetRange:
+		fmt.Fprintf(&buf, "%v-%v", o.OffsetLow, o.OffsetHigh)
+	case OffsetName:
+		fmt.Fprintf(&buf, "name=%q", o.OffsetName)
+	default:
+		panic(fmt.Errorf("should not happen: OffsetMatching=%#v", o.OffsetMatching))
+	}
+
+	buf.WriteString(")")
+
+	return buf.String()
+}
+
+// Search implements TreeSearcher.
+func (o Search) Search(k btrfsprim.Key, _ uint32) int {
+	if d := containers.NativeCompare(o.ObjectID, k.ObjectID); d != 0 {
+		return d
+	}
+
+	switch o.ItemTypeMatching {
+	case ItemTypeAny:
+		return 0
+	case ItemTypeExact:
+		if d := containers.NativeCompare(o.ItemType, k.ItemType); d != 0 {
+			return d
+		}
+	default:
+		panic(fmt.Errorf("should not happen: ItemTypeMatching=%#v", o.ItemTypeMatching))
+	}
+
+	switch o.OffsetMatching {
+	case OffsetAny, OffsetRange, OffsetName:
+		return 0
+	case OffsetExact:
+		return containers.NativeCompare(o.OffsetLow, k.Offset)
+	default:
+		panic(fmt.Errorf("should not happen: OffsetMatching=%#v", o.OffsetMatching))
+	}
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+// SearchObject returns a Search that searches all items belonging to
+// a given object.
+func SearchObject(objID btrfsprim.ObjID) Search {
+	return Search{
+		ObjectID:         objID,
+		ItemTypeMatching: ItemTypeAny,
+	}
+}
+
+// SearchExactKey returns a Search that searches for the exact key.
+func SearchExactKey(k btrfsprim.Key) Search {
+	return Search{
+		ObjectID: k.ObjectID,
+
+		ItemTypeMatching: ItemTypeExact,
+		ItemType:         k.ItemType,
+
+		OffsetMatching: OffsetExact,
+		OffsetLow:      k.Offset,
+	}
+}
+
+// SearchRootItem returns a Search that searches for the root item for
+// the given tree.
+func SearchRootItem(treeID btrfsprim.ObjID) Search {
+	return Search{
+		ObjectID: treeID,
+
+		ItemTypeMatching: ItemTypeExact,
+		ItemType:         btrfsprim.ROOT_ITEM_KEY,
+
+		OffsetMatching: OffsetAny,
+	}
+}
+
+type csumSearcher struct {
+	laddr   btrfsvol.LogicalAddr
+	algSize int
+}
+
+func (s csumSearcher) String() string { return fmt.Sprintf("csum for laddr=%v", s.laddr) }
+func (s csumSearcher) Search(key btrfsprim.Key, size uint32) int {
+	if d := containers.NativeCompare(btrfsprim.EXTENT_CSUM_OBJECTID, key.ObjectID); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(btrfsprim.EXTENT_CSUM_KEY, key.ItemType); d != 0 {
+		return d
+	}
+	itemBeg := btrfsvol.LogicalAddr(key.Offset)
+	numSums := int64(size) / int64(s.algSize)
+	itemEnd := itemBeg + btrfsvol.LogicalAddr(numSums*btrfssum.BlockSize)
+	switch {
+	case itemEnd <= s.laddr:
+		return 1
+	case s.laddr < itemBeg:
+		return -1
+	default:
+		return 0
+	}
+}
+
+// SearchCSum returns a TreeSearcher that searches for a csum-run
+// containing the csum for a given LogicalAddress.
+func SearchCSum(laddr btrfsvol.LogicalAddr, algSize int) TreeSearcher {
+	return csumSearcher{
+		laddr:   laddr,
+		algSize: algSize,
+	}
+}
diff --git a/lib/btrfs/btrfstree/btree_tree.go b/lib/btrfs/btrfstree/btree_tree.go
index c51efa5..df58c0c 100644
--- a/lib/btrfs/btrfstree/btree_tree.go
+++ b/lib/btrfs/btrfstree/btree_tree.go
@@ -405,7 +405,7 @@ func (fs TreeOperatorImpl) next(path TreePath, node *diskio.Ref[btrfsvol.Logical
 }
 
 // TreeSearch implements the 'TreeOperator' interface.
-func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, fn func(btrfsprim.Key, uint32) int) (Item, error) {
+func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, searcher TreeSearcher) (Item, error) {
 	sb, err := fs.Superblock()
 	if err != nil {
 		return Item{}, err
@@ -414,9 +414,9 @@ func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, fn func(btrfsprim.
 	if err != nil {
 		return Item{}, err
 	}
-	path, node, err := fs.treeSearch(*rootInfo, fn)
+	path, node, err := fs.treeSearch(*rootInfo, searcher.Search)
 	if err != nil {
-		return Item{}, err
+		return Item{}, fmt.Errorf("item with %s: %w", searcher, err)
 	}
 	item := node.Data.BodyLeaf[path.Node(-1).FromItemSlot]
 	item.Body = item.Body.CloneItem()
@@ -424,24 +424,13 @@ func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, fn func(btrfsprim.
 	return item, nil
 }
 
-// KeySearch returns a comparator suitable to be passed to TreeSearch.
-func KeySearch(fn func(btrfsprim.Key) int) func(btrfsprim.Key, uint32) int {
-	return func(key btrfsprim.Key, _ uint32) int {
-		return fn(key)
-	}
-}
-
 // TreeLookup implements the 'TreeOperator' interface.
 func (fs TreeOperatorImpl) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key) (Item, error) {
-	item, err := fs.TreeSearch(treeID, KeySearch(key.Compare))
-	if err != nil {
-		err = fmt.Errorf("item with key=%v: %w", key, err)
-	}
-	return item, err
+	return fs.TreeSearch(treeID, SearchExactKey(key))
 }
 
 // TreeSearchAll implements the 'TreeOperator' interface.
-func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfsprim.Key, uint32) int) ([]Item, error) {
+func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, searcher TreeSearcher) ([]Item, error) {
 	sb, err := fs.Superblock()
 	if err != nil {
 		return nil, err
@@ -450,9 +439,9 @@ func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfspr
 	if err != nil {
 		return nil, err
 	}
-	middlePath, middleNode, err := fs.treeSearch(*rootInfo, fn)
+	middlePath, middleNode, err := fs.treeSearch(*rootInfo, searcher.Search)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("items with %s: %w", searcher, err)
 	}
 	middleItem := middleNode.Data.BodyLeaf[middlePath.Node(-1).FromItemSlot]
 
@@ -469,7 +458,7 @@ func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfspr
 			break
 		}
 		prevItem := prevNode.Data.BodyLeaf[prevPath.Node(-1).FromItemSlot]
-		if fn(prevItem.Key, prevItem.BodySize) != 0 {
+		if searcher.Search(prevItem.Key, prevItem.BodySize) != 0 {
 			break
 		}
 		item := prevItem
@@ -482,7 +471,7 @@ func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfspr
 		middleNode, err = fs.ReadNode(middlePath)
 		if err != nil {
 			FreeNodeRef(middleNode)
-			return nil, err
+			return nil, fmt.Errorf("items with %s: %w", searcher, err)
 		}
 	}
 	nextPath, nextNode := middlePath, middleNode
@@ -496,7 +485,7 @@ func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfspr
 			break
 		}
 		nextItem := nextNode.Data.BodyLeaf[nextPath.Node(-1).FromItemSlot]
-		if fn(nextItem.Key, nextItem.BodySize) != 0 {
+		if searcher.Search(nextItem.Key, nextItem.BodySize) != 0 {
 			break
 		}
 		item := nextItem