Merge branch 'lukeshu/tree-api-pt0-prep'

1 files changed, 75 insertions, 7 deletions

diff --git a/lib/btrfs/btrfstree/btree_forrest.go b/lib/btrfs/btrfstree/btree_forrest.go
index 0f46d42..8f8e2de 100644
--- a/lib/btrfs/btrfstree/btree_forrest.go
+++ b/lib/btrfs/btrfstree/btree_forrest.go
@@ -5,6 +5,7 @@
 package btrfstree
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
@@ -16,7 +17,7 @@ import (
 // A TreeRoot is more-or-less a btrfsitem.Root, but simpler; returned by
 // LookupTreeRoot.
 type TreeRoot struct {
-	TreeID     btrfsprim.ObjID
+	ID         btrfsprim.ObjID
 	RootNode   btrfsvol.LogicalAddr
 	Level      uint8
 	Generation btrfsprim.Generation
@@ -24,32 +25,32 @@ type TreeRoot struct {
 
 // LookupTreeRoot is a utility function to help with implementing the
 // 'TreeOperator' interface.
-func LookupTreeRoot(fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*TreeRoot, error) {
+func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*TreeRoot, error) {
 	switch treeID {
 	case btrfsprim.ROOT_TREE_OBJECTID:
 		return &TreeRoot{
-			TreeID:     treeID,
+			ID:         treeID,
 			RootNode:   sb.RootTree,
 			Level:      sb.RootLevel,
 			Generation: sb.Generation, // XXX: same generation as LOG_TREE?
 		}, nil
 	case btrfsprim.CHUNK_TREE_OBJECTID:
 		return &TreeRoot{
-			TreeID:     treeID,
+			ID:         treeID,
 			RootNode:   sb.ChunkTree,
 			Level:      sb.ChunkLevel,
 			Generation: sb.ChunkRootGeneration,
 		}, nil
 	case btrfsprim.TREE_LOG_OBJECTID:
 		return &TreeRoot{
-			TreeID:     treeID,
+			ID:         treeID,
 			RootNode:   sb.LogTree,
 			Level:      sb.LogLevel,
 			Generation: sb.Generation, // XXX: same generation as ROOT_TREE?
 		}, nil
 	case btrfsprim.BLOCK_GROUP_TREE_OBJECTID:
 		return &TreeRoot{
-			TreeID:     treeID,
+			ID:         treeID,
 			RootNode:   sb.BlockGroupRoot,
 			Level:      sb.BlockGroupRootLevel,
 			Generation: sb.BlockGroupRootGeneration,
@@ -65,7 +66,7 @@ func LookupTreeRoot(fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*Tr
 		switch rootItemBody := rootItem.Body.(type) {
 		case *btrfsitem.Root:
 			return &TreeRoot{
-				TreeID:     treeID,
+				ID:         treeID,
 				RootNode:   rootItemBody.ByteNr,
 				Level:      rootItemBody.Level,
 				Generation: rootItemBody.Generation,
@@ -77,3 +78,70 @@ func LookupTreeRoot(fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*Tr
 		}
 	}
 }
+
+type TreeOperatorImpl struct {
+	NodeSource
+}
+
+func (fs TreeOperatorImpl) RawTree(ctx context.Context, treeID btrfsprim.ObjID) (*RawTree, error) {
+	sb, err := fs.Superblock()
+	if err != nil {
+		return nil, err
+	}
+	rootInfo, err := LookupTreeRoot(ctx, fs, *sb, treeID)
+	if err != nil {
+		return nil, err
+	}
+	return &RawTree{
+		Forrest:  fs,
+		TreeRoot: *rootInfo,
+	}, nil
+}
+
+// TreeWalk implements the 'TreeOperator' interface.
+func (fs TreeOperatorImpl) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, errHandle func(*TreeError), cbs TreeWalkHandler) {
+	tree, err := fs.RawTree(ctx, treeID)
+	if err != nil {
+		errHandle(&TreeError{Path: Path{{FromTree: treeID, ToMaxKey: btrfsprim.MaxKey}}, Err: err})
+		return
+	}
+	tree.TreeWalk(ctx, cbs)
+}
+
+// TreeLookup implements the 'TreeOperator' interface.
+func (fs TreeOperatorImpl) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key) (Item, error) {
+	ctx := context.TODO()
+	tree, err := fs.RawTree(ctx, treeID)
+	if err != nil {
+		return Item{}, err
+	}
+	return tree.TreeLookup(ctx, key)
+}
+
+// TreeSearch implements the 'TreeOperator' interface.
+func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, searcher TreeSearcher) (Item, error) {
+	ctx := context.TODO()
+	tree, err := fs.RawTree(ctx, treeID)
+	if err != nil {
+		return Item{}, err
+	}
+	return tree.TreeSearch(ctx, searcher)
+}
+
+// TreeSearchAll implements the 'TreeOperator' interface.
+func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, searcher TreeSearcher) ([]Item, error) {
+	ctx := context.TODO()
+	tree, err := fs.RawTree(ctx, treeID)
+	if err != nil {
+		return nil, err
+	}
+
+	var ret []Item
+	err = tree.TreeSubrange(ctx, 1, searcher, func(item Item) bool {
+		item.Body = item.Body.CloneItem()
+		ret = append(ret, item)
+		return true
+	})
+
+	return ret, err
+}