implement TreeSearchAll

2 files changed, 185 insertions, 22 deletions

diff --git a/pkg/btrfs/btree.go b/pkg/btrfs/btree.go
index 0553c46..fb3dd65 100644
--- a/pkg/btrfs/btree.go
+++ b/pkg/btrfs/btree.go
@@ -12,6 +12,16 @@ import (
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
+// - The first element will always have an ItemIdx of -1.
+//
+// - For .Item() callbacks, the last element will always have a
+//   NodeAddr of 0.
+//
+// ex:
+//
+//     {-1, 0x01, 3}→{9, 0x02, 2}→{9, 0x03, 1}→{9, 0x04, 0}→{9, 0, 0}
+type TreePath []TreePathElem
+
 // A TreePathElem essentially represents a KeyPointer.
 type TreePathElem struct {
 	// ItemIdx is the index of this KeyPointer in the parent Node;
@@ -34,12 +44,6 @@ func (elem TreePathElem) writeNodeTo(w io.Writer) {
 	}
 }
 
-// - The first element will always have an ItemIdx of -1.
-//
-// - For .Item() callbacks, the last element will always have a
-//   NodeAddr of 0.
-type TreePath []TreePathElem
-
 func (path TreePath) String() string {
 	if len(path) == 0 {
 		return "(empty-path)"
@@ -105,12 +109,8 @@ func (fs *FS) treeWalk(path TreePath, cbs TreeWalkHandler) error {
 			return err
 		}
 	}
-	node, err := fs.ReadNode(path[len(path)-1].NodeAddr)
+	node, err := fs.readNodeAtLevel(path[len(path)-1].NodeAddr, path[len(path)-1].NodeLevel)
 	if node != nil {
-		if exp := path[len(path)-1].NodeLevel; exp != math.MaxUint8 && node.Data.Head.Level != exp && err == nil {
-			err = fmt.Errorf("btrfs.FS.TreeWalk: node@%v: expected level %v but has level %v",
-				node.Addr, exp, node.Data.Head.Level)
-		}
 		path[len(path)-1].NodeLevel = node.Data.Head.Level
 	}
 	if cbs.Node != nil {
@@ -174,17 +174,24 @@ func (fs *FS) treeWalk(path TreePath, cbs TreeWalkHandler) error {
 	return nil
 }
 
-func (fs *FS) TreeSearch(treeRoot LogicalAddr, fn func(Key) int) (Key, btrfsitem.Item, error) {
-	nodeAddr := treeRoot
-
+func (fs *FS) treeSearch(treeRoot LogicalAddr, fn func(Key) int) (TreePath, *util.Ref[LogicalAddr, Node], error) {
+	path := TreePath{
+		TreePathElem{
+			ItemIdx:   -1,
+			NodeAddr:  treeRoot,
+			NodeLevel: math.MaxUint8,
+		},
+	}
 	for {
-		if nodeAddr == 0 {
-			return Key{}, nil, iofs.ErrNotExist
+		if path[len(path)-1].NodeAddr == 0 {
+			return nil, nil, iofs.ErrNotExist
 		}
-		node, err := fs.ReadNode(nodeAddr)
+		node, err := fs.readNodeAtLevel(path[len(path)-1].NodeAddr, path[len(path)-1].NodeLevel)
 		if err != nil {
-			return Key{}, nil, err
+			return nil, nil, err
 		}
+		path[len(path)-1].NodeLevel = node.Data.Head.Level
+
 		if node.Data.Head.Level > 0 {
 			// internal node
 
@@ -208,9 +215,13 @@ func (fs *FS) TreeSearch(treeRoot LogicalAddr, fn func(Key) int) (Key, btrfsitem
 				}
 			}
 			if lastGood < 0 {
-				return Key{}, nil, iofs.ErrNotExist
+				return nil, nil, iofs.ErrNotExist
 			}
-			nodeAddr = node.Data.BodyInternal[lastGood].BlockPtr
+			path = append(path, TreePathElem{
+				ItemIdx:   lastGood,
+				NodeAddr:  node.Data.BodyInternal[lastGood].BlockPtr,
+				NodeLevel: node.Data.Head.Level - 1,
+			})
 		} else {
 			// leaf node
 
@@ -234,14 +245,153 @@ func (fs *FS) TreeSearch(treeRoot LogicalAddr, fn func(Key) int) (Key, btrfsitem
 				case direction > 0:
 					items = items[midpoint+1:]
 				case direction == 0:
-					return items[midpoint].Head.Key, items[midpoint].Body, nil
+					path = append(path, TreePathElem{
+						ItemIdx: midpoint,
+					})
+					return path, node, nil
 				}
 			}
-			return Key{}, nil, iofs.ErrNotExist
+			return nil, nil, iofs.ErrNotExist
 		}
 	}
 }
 
+func (fs *FS) prev(path TreePath, node *util.Ref[LogicalAddr, Node]) (TreePath, *util.Ref[LogicalAddr, Node], error) {
+	var err error
+	path = append(TreePath(nil), path...)
+
+	// go up
+	for path[len(path)-1].ItemIdx == 0 {
+		path = path[:len(path)-1]
+		if len(path) == 0 {
+			return nil, nil, nil
+		}
+	}
+	// go left
+	path[len(path)-1].ItemIdx--
+	// go down
+	for path[len(path)-1].NodeAddr != 0 {
+		if node.Addr != path[len(path)-1].NodeAddr {
+			node, err = fs.readNodeAtLevel(path[len(path)-1].NodeAddr, path[len(path)-1].NodeLevel)
+			if err != nil {
+				return nil, nil, err
+			}
+			path[len(path)-1].NodeLevel = node.Data.Head.Level
+		}
+		if node.Data.Head.Level > 0 {
+			path = append(path, TreePathElem{
+				ItemIdx:   len(node.Data.BodyInternal) - 1,
+				NodeAddr:  node.Data.BodyInternal[len(node.Data.BodyInternal)-1].BlockPtr,
+				NodeLevel: node.Data.Head.Level - 1,
+			})
+		} else {
+			path = append(path, TreePathElem{
+				ItemIdx: len(node.Data.BodyLeaf) - 1,
+			})
+		}
+	}
+	// return
+	if node.Addr != path[len(path)-2].NodeAddr {
+		node, err = fs.readNodeAtLevel(path[len(path)-2].NodeAddr, path[len(path)-2].NodeLevel)
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+	return path, node, nil
+}
+
+func (fs *FS) next(path TreePath, node *util.Ref[LogicalAddr, Node]) (TreePath, *util.Ref[LogicalAddr, Node], error) {
+	var err error
+	path = append(TreePath(nil), path...)
+
+	// go up
+	if node.Addr != path[len(path)-2].NodeAddr {
+		node, err = fs.readNodeAtLevel(path[len(path)-2].NodeAddr, path[len(path)-2].NodeLevel)
+		if err != nil {
+			return nil, nil, err
+		}
+		path[len(path)-2].NodeLevel = node.Data.Head.Level
+	}
+	for path[len(path)-1].ItemIdx+1 >= int(node.Data.Head.NumItems) {
+		path = path[:len(path)-1]
+		if len(path) == 0 {
+			return nil, nil, nil
+		}
+		if node.Addr != path[len(path)-2].NodeAddr {
+			node, err = fs.readNodeAtLevel(path[len(path)-2].NodeAddr, path[len(path)-2].NodeLevel)
+			if err != nil {
+				return nil, nil, err
+			}
+			path[len(path)-2].NodeLevel = node.Data.Head.Level
+		}
+	}
+	// go left
+	path[len(path)-1].ItemIdx++
+	// go down
+	for path[len(path)-1].NodeAddr != 0 {
+		if node.Addr != path[len(path)-1].NodeAddr {
+			node, err = fs.readNodeAtLevel(path[len(path)-1].NodeAddr, path[len(path)-1].NodeLevel)
+			if err != nil {
+				return nil, nil, err
+			}
+			path[len(path)-1].NodeLevel = node.Data.Head.Level
+		}
+		if node.Data.Head.Level > 0 {
+			path = append(path, TreePathElem{
+				ItemIdx:   0,
+				NodeAddr:  node.Data.BodyInternal[len(node.Data.BodyInternal)-1].BlockPtr,
+				NodeLevel: node.Data.Head.Level - 1,
+			})
+		} else {
+			path = append(path, TreePathElem{
+				ItemIdx: 0,
+			})
+		}
+	}
+	// return
+	if node.Addr != path[len(path)-2].NodeAddr {
+		node, err = fs.readNodeAtLevel(path[len(path)-2].NodeAddr, path[len(path)-2].NodeLevel)
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+	return path, node, nil
+}
+
+func (fs *FS) TreeSearch(treeRoot LogicalAddr, fn func(Key) int) (Key, btrfsitem.Item, error) {
+	path, node, err := fs.treeSearch(treeRoot, fn)
+	if err != nil {
+		return Key{}, nil, err
+	}
+	item := node.Data.BodyLeaf[path[len(path)-1].ItemIdx]
+	return item.Head.Key, item.Body, nil
+}
+
 func (fs *FS) TreeLookup(treeRoot LogicalAddr, key Key) (Key, btrfsitem.Item, error) {
 	return fs.TreeSearch(treeRoot, key.Cmp)
 }
+
+func (fs *FS) TreeSearchAll(treeRoot LogicalAddr, fn func(Key) int) ([]Item, error) {
+	middlePath, middleNode, err := fs.treeSearch(treeRoot, fn)
+	if err != nil {
+		return nil, err
+	}
+
+	var ret = []Item{middleNode.Data.BodyLeaf[middlePath[len(middlePath)-1].ItemIdx]}
+	for prevPath, prevNode := middlePath, middleNode; prevPath != nil; {
+		prevPath, prevNode, err = fs.prev(prevPath, prevNode)
+		if err != nil {
+			return nil, err
+		}
+		ret = append(ret, prevNode.Data.BodyLeaf[prevPath[len(prevPath)-1].ItemIdx])
+	}
+	util.ReverseSlice(ret)
+	for nextPath, nextNode := middlePath, middleNode; nextPath != nil; {
+		nextPath, nextNode, err = fs.next(nextPath, nextNode)
+		if err != nil {
+			return nil, err
+		}
+		ret = append(ret, nextNode.Data.BodyLeaf[nextPath[len(nextPath)-1].ItemIdx])
+	}
+	return ret, nil
+}
diff --git a/pkg/btrfs/types_btree.go b/pkg/btrfs/types_btree.go
index 4c77437..58a08e0 100644
--- a/pkg/btrfs/types_btree.go
+++ b/pkg/btrfs/types_btree.go
@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"math"
 
 	"lukeshu.com/btrfs-tools/pkg/binstruct"
 	"lukeshu.com/btrfs-tools/pkg/btrfs/btrfsitem"
@@ -388,3 +389,15 @@ func (fs *FS) ReadNode(addr LogicalAddr) (*util.Ref[LogicalAddr, Node], error) {
 		return nil
 	})
 }
+
+func (fs *FS) readNodeAtLevel(addr LogicalAddr, expLevel uint8) (*util.Ref[LogicalAddr, Node], error) {
+	node, err := fs.ReadNode(addr)
+	if err != nil {
+		return node, err
+	}
+	if expLevel != math.MaxUint8 && node.Data.Head.Level != expLevel {
+		return node, fmt.Errorf("btrfs.FS.ReadNode: node@%v: expected level %v but has level %v",
+			node.Addr, expLevel, node.Data.Head.Level)
+	}
+	return node, nil
+}