github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/vm/iterator_test.go (about)

     1  package vm
     2  
     3  import (
     4  	"fmt"
     5  	"math/big"
     6  	"testing"
     7  
     8  	"github.com/nspcc-dev/neo-go/internal/random"
     9  	"github.com/nspcc-dev/neo-go/pkg/core/interop/interopnames"
    10  	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
    11  	"github.com/nspcc-dev/neo-go/pkg/smartcontract/callflag"
    12  	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  type arrayIterator struct {
    17  	index  int
    18  	values []stackitem.Item
    19  }
    20  
    21  func TestCreateCallAndUnwrapIteratorScript(t *testing.T) {
    22  	ctrHash := random.Uint160()
    23  	ctrMethod := "mymethod"
    24  	param := stackitem.NewBigInteger(big.NewInt(42))
    25  
    26  	const totalItems = 8
    27  	values := make([]stackitem.Item, totalItems)
    28  	for i := range values {
    29  		values[i] = stackitem.NewBigInteger(big.NewInt(int64(i)))
    30  	}
    31  
    32  	checkStack := func(t *testing.T, script []byte, index int, prefetch bool) {
    33  		v := load(script)
    34  		it := &arrayIterator{index: -1, values: values}
    35  		v.SyscallHandler = func(v *VM, id uint32) error {
    36  			switch id {
    37  			case interopnames.ToID([]byte(interopnames.SystemContractCall)):
    38  				require.Equal(t, ctrHash.BytesBE(), v.Estack().Pop().Value())
    39  				require.Equal(t, []byte(ctrMethod), v.Estack().Pop().Value())
    40  				require.Equal(t, big.NewInt(int64(callflag.All)), v.Estack().Pop().Value())
    41  				require.Equal(t, []stackitem.Item{param}, v.Estack().Pop().Value())
    42  				v.Estack().PushItem(stackitem.NewInterop(it))
    43  			case interopnames.ToID([]byte(interopnames.SystemIteratorNext)):
    44  				require.Equal(t, it, v.Estack().Pop().Value())
    45  				it.index++
    46  				v.Estack().PushVal(it.index < len(it.values))
    47  			case interopnames.ToID([]byte(interopnames.SystemIteratorValue)):
    48  				require.Equal(t, it, v.Estack().Pop().Value())
    49  				v.Estack().PushVal(it.values[it.index])
    50  			default:
    51  				return fmt.Errorf("unexpected syscall: %d", id)
    52  			}
    53  			return nil
    54  		}
    55  		require.NoError(t, v.Run())
    56  
    57  		if prefetch && index <= len(values) {
    58  			require.Equal(t, 2, v.Estack().Len())
    59  
    60  			it, ok := v.Estack().Pop().Interop().Value().(*arrayIterator)
    61  			require.True(t, ok)
    62  			require.Equal(t, index-1, it.index)
    63  			require.Equal(t, values[:index], v.Estack().Pop().Array())
    64  			return
    65  		}
    66  		if len(values) < index {
    67  			index = len(values)
    68  		}
    69  		require.Equal(t, 1, v.Estack().Len())
    70  		require.Equal(t, values[:index], v.Estack().Pop().Array())
    71  	}
    72  
    73  	t.Run("truncate", func(t *testing.T) {
    74  		t.Run("zero", func(t *testing.T) {
    75  			const index = 0
    76  			script, err := smartcontract.CreateCallAndUnwrapIteratorScript(ctrHash, ctrMethod, index, param)
    77  			require.NoError(t, err)
    78  
    79  			// The behaviour is a bit unexpected, but not a problem (why would anyone fetch 0 items).
    80  			// Let's have test, to make it obvious.
    81  			checkStack(t, script, index+1, false)
    82  		})
    83  		t.Run("all", func(t *testing.T) {
    84  			const index = totalItems + 1
    85  			script, err := smartcontract.CreateCallAndUnwrapIteratorScript(ctrHash, ctrMethod, index, param)
    86  			require.NoError(t, err)
    87  
    88  			checkStack(t, script, index, false)
    89  		})
    90  		t.Run("partial", func(t *testing.T) {
    91  			const index = totalItems / 2
    92  			script, err := smartcontract.CreateCallAndUnwrapIteratorScript(ctrHash, ctrMethod, index, param)
    93  			require.NoError(t, err)
    94  
    95  			checkStack(t, script, index, false)
    96  		})
    97  	})
    98  	t.Run("prefetch", func(t *testing.T) {
    99  		t.Run("zero", func(t *testing.T) {
   100  			const index = 0
   101  			script, err := smartcontract.CreateCallAndPrefetchIteratorScript(ctrHash, ctrMethod, index, param)
   102  			require.NoError(t, err)
   103  
   104  			checkStack(t, script, index+1, true)
   105  		})
   106  		t.Run("all", func(t *testing.T) {
   107  			const index = totalItems + 1 // +1 to test with iterator dropped
   108  			script, err := smartcontract.CreateCallAndPrefetchIteratorScript(ctrHash, ctrMethod, index, param)
   109  			require.NoError(t, err)
   110  
   111  			checkStack(t, script, index, true)
   112  		})
   113  		t.Run("partial", func(t *testing.T) {
   114  			const index = totalItems / 2
   115  			script, err := smartcontract.CreateCallAndPrefetchIteratorScript(ctrHash, ctrMethod, index, param)
   116  			require.NoError(t, err)
   117  
   118  			checkStack(t, script, index, true)
   119  		})
   120  	})
   121  }