unwrap: Add ArrayAndSessionIterator(), close #3272

It can be used to work with the results of
CreateCallAndPrefetchIteratorScript() execution. The first item must be
an array and the optional second item must be an iterator, containing
remaining elements.

Signed-off-by: Evgenii Stratonikov <fyfyrchik@runbox.com>

pkg/rpcclient/unwrap/unwrap.go

@@ -167,12 +167,9 @@ func SessionIterator(r *result.Invoke, err error) (uuid.UUID, result.Iterator, e
 	if err != nil {
 		return uuid.UUID{}, result.Iterator{}, err
 	}
-	if t := itm.Type(); t != stackitem.InteropT {
-		return uuid.UUID{}, result.Iterator{}, fmt.Errorf("expected InteropInterface, got %s", t)
-	}
-	iter, ok := itm.Value().(result.Iterator)
-	if !ok {
-		return uuid.UUID{}, result.Iterator{}, errors.New("the item is InteropInterface, but not an Iterator")
+	iter, err := itemToSessionIterator(itm)
+	if err != nil {
+		return uuid.UUID{}, result.Iterator{}, err
 	}
 	if (r.Session == uuid.UUID{}) && iter.ID != nil {
 		return uuid.UUID{}, result.Iterator{}, ErrNoSessionID
@@ -180,6 +177,54 @@ func SessionIterator(r *result.Invoke, err error) (uuid.UUID, result.Iterator, e
 	return r.Session, iter, nil
 }
 
+// ArrayAndSessionIterator expects correct execution (HALT state) with one or two stack
+// items returned. If there is 1 item, it must be an array. If there is a second item,
+// it must be an iterator. This is exactly the result of smartcontract.CreateCallAndPrefetchIteratorScript.
+// Sessions must be enabled on the RPC server for this to function correctly.
+func ArrayAndSessionIterator(r *result.Invoke, err error) ([]stackitem.Item, uuid.UUID, result.Iterator, error) {
+	if err := checkResOK(r, err); err != nil {
+		return nil, uuid.UUID{}, result.Iterator{}, err
+	}
+	if len(r.Stack) == 0 {
+		return nil, uuid.UUID{}, result.Iterator{}, errors.New("result stack is empty")
+	}
+	if len(r.Stack) != 1 && len(r.Stack) != 2 {
+		return nil, uuid.UUID{}, result.Iterator{}, fmt.Errorf("expected 1 or 2 result items, got %d", len(r.Stack))
+	}
+
+	// Unwrap array.
+	itm := r.Stack[0]
+	arr, ok := itm.Value().([]stackitem.Item)
+	if !ok {
+		return nil, uuid.UUID{}, result.Iterator{}, errors.New("not an array")
+	}
+
+	// Check whether iterator exists and unwrap it.
+	if len(r.Stack) == 1 {
+		return arr, uuid.UUID{}, result.Iterator{}, nil
+	}
+
+	iter, err := itemToSessionIterator(r.Stack[1])
+	if err != nil {
+		return nil, uuid.UUID{}, result.Iterator{}, err
+	}
+	if (r.Session == uuid.UUID{}) {
+		return nil, uuid.UUID{}, result.Iterator{}, ErrNoSessionID
+	}
+	return arr, r.Session, iter, nil
+}
+
+func itemToSessionIterator(itm stackitem.Item) (result.Iterator, error) {
+	if t := itm.Type(); t != stackitem.InteropT {
+		return result.Iterator{}, fmt.Errorf("expected InteropInterface, got %s", t)
+	}
+	iter, ok := itm.Value().(result.Iterator)
+	if !ok {
+		return result.Iterator{}, errors.New("the item is InteropInterface, but not an Iterator")
+	}
+	return iter, nil
+}
+
 // Array expects correct execution (HALT state) with a single array stack item
 // returned. This item is returned to the caller. Notice that this function can
 // be used for structures as well since they're also represented as slices of

pkg/rpcclient/unwrap/unwrap_test.go

@@ -50,6 +50,10 @@ func TestStdErrors(t *testing.T) {
 			_, _, err = SessionIterator(r, err)
 			return nil, err
 		},
+		func(r *result.Invoke, err error) (any, error) {
+			_, _, _, err = ArrayAndSessionIterator(r, err)
+			return nil, err
+		},
 		func(r *result.Invoke, err error) (any, error) {
 			return Array(r, err)
 		},
@@ -256,6 +260,39 @@ func TestSessionIterator(t *testing.T) {
 	require.Equal(t, iter, ri)
 }
 
+func TestArraySessionIterator(t *testing.T) {
+	_, _, _, err := ArrayAndSessionIterator(&result.Invoke{State: "HALT", Stack: []stackitem.Item{stackitem.Make(42)}}, nil)
+	require.Error(t, err)
+
+	_, _, _, err = ArrayAndSessionIterator(&result.Invoke{State: "HALT", Stack: []stackitem.Item{stackitem.NewInterop(42)}}, nil)
+	require.Error(t, err)
+
+	arr := stackitem.NewArray([]stackitem.Item{stackitem.Make(42)})
+	ra, rs, ri, err := ArrayAndSessionIterator(&result.Invoke{State: "HALT", Stack: []stackitem.Item{arr}}, nil)
+	require.NoError(t, err)
+	require.Equal(t, arr.Value(), ra)
+	require.Empty(t, rs)
+	require.Empty(t, ri)
+
+	_, _, _, err = ArrayAndSessionIterator(&result.Invoke{State: "HALT", Stack: []stackitem.Item{arr, stackitem.NewInterop(42)}}, nil)
+	require.Error(t, err)
+
+	iid := uuid.New()
+	iter := result.Iterator{ID: &iid}
+	_, _, _, err = ArrayAndSessionIterator(&result.Invoke{State: "HALT", Stack: []stackitem.Item{arr, stackitem.NewInterop(iter)}}, nil)
+	require.ErrorIs(t, err, ErrNoSessionID)
+
+	sid := uuid.New()
+	_, rs, ri, err = ArrayAndSessionIterator(&result.Invoke{Session: sid, State: "HALT", Stack: []stackitem.Item{arr, stackitem.NewInterop(iter)}}, nil)
+	require.NoError(t, err)
+	require.Equal(t, arr.Value(), ra)
+	require.Equal(t, sid, rs)
+	require.Equal(t, iter, ri)
+
+	_, _, _, err = ArrayAndSessionIterator(&result.Invoke{Session: sid, State: "HALT", Stack: []stackitem.Item{arr, stackitem.NewInterop(iter), stackitem.Make(42)}}, nil)
+	require.Error(t, err)
+}
+
 func TestArray(t *testing.T) {
 	_, err := Array(&result.Invoke{State: "HALT", Stack: []stackitem.Item{stackitem.Make(42)}}, nil)
 	require.Error(t, err)