From bf7d80f44b69019903267525ffdbf439ab6d19b8 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <e.stratonikov@yadro.com>
Date: Fri, 24 Mar 2023 11:54:59 +0300
Subject: [PATCH] [#172] morph/client: Support iterators via neo-go session API

Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>
---
 pkg/morph/client/client.go | 40 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/pkg/morph/client/client.go b/pkg/morph/client/client.go
index 1f32fc54..80682886 100644
--- a/pkg/morph/client/client.go
+++ b/pkg/morph/client/client.go
@@ -19,6 +19,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/gas"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/nep17"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/rolemgmt"
+	"github.com/nspcc-dev/neo-go/pkg/rpcclient/unwrap"
 	sc "github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/trigger"
 	"github.com/nspcc-dev/neo-go/pkg/util"
@@ -195,6 +196,45 @@ func (c *Client) Invoke(contract util.Uint160, fee fixedn.Fixed8, method string,
 	return nil
 }
 
+// TestInvokeIterator invokes contract method returning an iterator and executes cb on each element.
+// If cb returns an error, the session is closed and this error is returned as-is.
+// If the remove neo-go node does not support sessions, `unwrap.ErrNoSessionID` is returned.
+func (c *Client) TestInvokeIterator(cb func(stackitem.Item) error, contract util.Uint160, method string, args ...interface{}) error {
+	c.switchLock.RLock()
+	defer c.switchLock.RUnlock()
+
+	if c.inactive {
+		return ErrConnectionLost
+	}
+
+	val, err := c.rpcActor.Call(contract, method, args...)
+	if err != nil {
+		return err
+	} else if val.State != HaltState {
+		return wrapFrostFSError(&notHaltStateError{state: val.State, exception: val.FaultException})
+	}
+
+	sid, r, err := unwrap.SessionIterator(val, err)
+	if err != nil {
+		return err
+	}
+
+	defer func() {
+		_ = c.rpcActor.TerminateSession(sid)
+	}()
+
+	items, err := c.rpcActor.TraverseIterator(sid, &r, 0)
+	for err == nil && len(items) != 0 {
+		for i := range items {
+			if err = cb(items[i]); err != nil {
+				return err
+			}
+		}
+		items, err = c.rpcActor.TraverseIterator(sid, &r, 0)
+	}
+	return err
+}
+
 // TestInvoke invokes contract method locally in neo-go node. This method should
 // be used to read data from smart-contract.
 func (c *Client) TestInvoke(contract util.Uint160, method string, args ...any) (res []stackitem.Item, err error) {