implement btree lookup

3 files changed, 95 insertions, 0 deletions

diff --git a/pkg/btrfs/btree.go b/pkg/btrfs/btree.go
index 99c909f..72da1b9 100644
--- a/pkg/btrfs/btree.go
+++ b/pkg/btrfs/btree.go
@@ -8,6 +8,7 @@ import (
 	"math"
 	"strings"
 
+	"lukeshu.com/btrfs-tools/pkg/btrfs/btrfsitem"
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
@@ -172,3 +173,75 @@ func (fs *FS) treeWalk(path TreeWalkPath, cbs TreeWalkHandler) error {
 	}
 	return nil
 }
+
+func (fs *FS) TreeSearch(treeRoot LogicalAddr, fn func(Key) int) (Key, btrfsitem.Item, error) {
+	nodeAddr := treeRoot
+
+	for {
+		if nodeAddr == 0 {
+			return Key{}, nil, iofs.ErrNotExist
+		}
+		node, err := fs.ReadNode(nodeAddr)
+		if err != nil {
+			return Key{}, nil, err
+		}
+		if node.Data.Head.Level > 0 {
+			// internal node
+
+			// Search for the right-most node.Data.BodyInternal item for which
+			// `fn(item.Key) >= 0`.
+			//
+			//    + + + + 0 - - - -
+			//
+			// There may or may not be a value that returns '0'.
+			//
+			// Implement this search as a binary search.
+			lastGood := -1
+			firstBad := len(node.Data.BodyInternal)
+			for firstBad > lastGood+1 {
+				midpoint := (lastGood + firstBad) / 2
+				direction := fn(node.Data.BodyInternal[midpoint].Key)
+				if direction < 0 {
+					firstBad = midpoint
+				} else {
+					lastGood = midpoint
+				}
+			}
+			if lastGood < 0 {
+				return Key{}, nil, iofs.ErrNotExist
+			}
+			nodeAddr = node.Data.BodyInternal[lastGood].BlockPtr
+		} else {
+			// leaf node
+
+			// Search for a member of node.Data.BodyLeaf for which
+			// `fn(item.Head.Key) == 0`.
+			//
+			//    + + + + 0 - - - -
+			//
+			// Such an item might not exist; in this case, return nil/ErrNotExist.
+			// Multiple such items might exist; in this case, it does not matter which
+			// is returned.
+			//
+			// Implement this search as a binary search.
+			items := node.Data.BodyLeaf
+			for len(items) > 0 {
+				midpoint := len(items) / 2
+				direction := fn(items[midpoint].Head.Key)
+				switch {
+				case direction < 0:
+					items = items[:midpoint]
+				case direction > 0:
+					items = items[midpoint+1:]
+				case direction == 0:
+					return items[midpoint].Head.Key, items[midpoint].Body, nil
+				}
+			}
+			return Key{}, nil, iofs.ErrNotExist
+		}
+	}
+}
+
+func (fs *FS) TreeLookup(treeRoot LogicalAddr, key Key) (Key, btrfsitem.Item, error) {
+	return fs.TreeSearch(treeRoot, key.Cmp)
+}
diff --git a/pkg/btrfs/internal/misc.go b/pkg/btrfs/internal/misc.go
index 2a98464..3dc77d6 100644
--- a/pkg/btrfs/internal/misc.go
+++ b/pkg/btrfs/internal/misc.go
@@ -4,6 +4,7 @@ import (
 	"time"
 
 	"lukeshu.com/btrfs-tools/pkg/binstruct"
+	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
 type Generation uint64
@@ -15,6 +16,16 @@ type Key struct {
 	binstruct.End `bin:"off=0x11"`
 }
 
+func (a Key) Cmp(b Key) int {
+	if d := util.CmpUint(a.ObjectID, b.ObjectID); d != 0 {
+		return d
+	}
+	if d := util.CmpUint(a.ItemType, b.ItemType); d != 0 {
+		return d
+	}
+	return util.CmpUint(a.Offset, b.Offset)
+}
+
 type Time struct {
 	Sec           int64  `bin:"off=0x0, siz=0x8"` // Number of seconds since 1970-01-01T00:00:00Z.
 	NSec          uint32 `bin:"off=0x8, siz=0x4"` // Number of nanoseconds since the beginning of the second.
diff --git a/pkg/util/generic.go b/pkg/util/generic.go
index 70fed5b..dbe077f 100644
--- a/pkg/util/generic.go
+++ b/pkg/util/generic.go
@@ -77,3 +77,14 @@ func SortedMapKeys[K constraints.Ordered, V any](m map[K]V) []K {
 	SortSlice(ret)
 	return ret
 }
+
+func CmpUint[T constraints.Unsigned](a, b T) int {
+	switch {
+	case a < b:
+		return -1
+	case a == b:
+		return 0
+	default:
+		return 1
+	}
+}