From 5304e68b9a169994ec36f531e6ae70423a886c35 Mon Sep 17 00:00:00 2001
From: Alex Vanin <alexey@nspcc.ru>
Date: Wed, 1 Jun 2022 13:57:23 +0300
Subject: [PATCH] [#477] Add getTreeNodes method

Find multiple treeNodes from
subtree in single request.

Signed-off-by: Alex Vanin <alexey@nspcc.ru>
---
 internal/neofs/tree.go | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/internal/neofs/tree.go b/internal/neofs/tree.go
index a3e24cc..0fadb3a 100644
--- a/internal/neofs/tree.go
+++ b/internal/neofs/tree.go
@@ -476,25 +476,40 @@ func (c *TreeClient) DeleteBucketTagging(ctx context.Context, cnrID *cid.ID) err
 }
 
 func (c *TreeClient) getTreeNode(ctx context.Context, cnrID *cid.ID, nodeID uint64, key string) (*TreeNode, error) {
+	nodes, err := c.getTreeNodes(ctx, cnrID, nodeID, key)
+	if err != nil {
+		return nil, err
+	}
+	// if there will be many allocations, consider having separate
+	// implementations of 'getTreeNode' and 'getTreeNodes'
+	return nodes[key], nil
+}
+
+func (c *TreeClient) getTreeNodes(ctx context.Context, cnrID *cid.ID, nodeID uint64, keys ...string) (map[string]*TreeNode, error) {
 	subtree, err := c.getSubTree(ctx, cnrID, versionTree, nodeID, 1)
 	if err != nil {
 		return nil, err
 	}
 
-	var treeNode *TreeNode
+	treeNodes := make(map[string]*TreeNode, len(keys))
 
 	for _, s := range subtree {
 		node, err := newTreeNode(s)
 		if err != nil {
 			return nil, err
 		}
-		if _, ok := node.Get(key); ok {
-			treeNode = node
+		for _, key := range keys {
+			if _, ok := node.Get(key); ok {
+				treeNodes[key] = node
+				break
+			}
+		}
+		if len(treeNodes) == len(keys) {
 			break
 		}
 	}
 
-	return treeNode, nil
+	return treeNodes, nil
 }
 
 func (c *TreeClient) GetVersions(ctx context.Context, cnrID *cid.ID, filepath string) ([]*data.NodeVersion, error) {