frostfs-contract/commonclient/iterator.go

package commonclient

import (
	"fmt"

	"github.com/google/uuid"
	"github.com/nspcc-dev/neo-go/pkg/rpcclient/unwrap"
	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
	"github.com/nspcc-dev/neo-go/pkg/util"
	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
)

// ReadIteratorItems calls method that returns iterator and traverses the iterator until all items are read into array.
func ReadIteratorItems(inv Invoker, batchSize int, contract util.Uint160, method string, params ...any) ([]stackitem.Item, error) {
	if batchSize <= 0 {
		panic("batch size must be positive")
	}

	script, err := smartcontract.CreateCallAndPrefetchIteratorScript(contract, method, batchSize, params...)
	if err != nil {
		return nil, fmt.Errorf("couldn't create unwrap script: %w", err)
	}

	arr, sessionID, iter, err := unwrap.ArrayAndSessionIterator(inv.Run(script))
	if err != nil {
		return nil, fmt.Errorf("unwrap session iterator: %w", err)
	}
	if (sessionID == uuid.UUID{}) {
		return arr, nil
	}

	defer func() {
		_ = inv.TerminateSession(sessionID)
	}()

	var shouldStop bool
	res := arr
	for !shouldStop {
		items, err := inv.TraverseIterator(sessionID, &iter, batchSize)
		if err != nil {
			return nil, err
		}

		res = append(res, items...)
		shouldStop = len(items) < batchSize
	}

	return res, nil
}