From 66501f9ef92648f43c18d27be34a498518c0b905 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 12 Sep 2019 11:32:09 +0300
Subject: [PATCH] vm: make REMOVE consume array from stack

---
 pkg/vm/vm.go      | 23 +++++++++++++++++++----
 pkg/vm/vm_test.go | 12 +++---------
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index b9c667c69..0587a6929 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -728,10 +728,25 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		}
 	case REMOVE:
 		key := int(v.estack.Pop().BigInt().Int64())
-		elem := v.estack.Peek(0)
-		a := elem.Array()
-		a = append(a[:key], a[key+1:]...)
-		elem.value = makeStackItem(a)
+		elem := v.estack.Pop()
+		switch t := elem.value.(type) {
+		case *ArrayItem:
+			a := t.value
+			if key < 0 || key >= len(a) {
+				panic("REMOVE: invalid index")
+			}
+			a = append(a[:key], a[key+1:]...)
+			t.value = a
+		case *StructItem:
+			a := t.value
+			if key < 0 || key >= len(a) {
+				panic("REMOVE: invalid index")
+			}
+			a = append(a[:key], a[key+1:]...)
+			t.value = a
+		default:
+			panic("REMOVE: invalid type")
+		}
 
 	case ARRAYSIZE:
 		elem := v.estack.Pop()
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index d813ebc89..2f6857c4e 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -1314,23 +1314,17 @@ func TestREMOVEBadIndex(t *testing.T) {
 }
 
 func TestREMOVEGood(t *testing.T) {
-	prog := makeProgram(REMOVE)
+	prog := makeProgram(DUP, PUSH2, REMOVE)
 	elements := []int{22, 34, 42, 55, 81}
 	reselements := []int{22, 34, 55, 81}
 	vm := load(prog)
 	vm.estack.PushVal(1)
 	vm.estack.PushVal(elements)
-	vm.estack.PushVal(2)
 	vm.Run()
 	assert.Equal(t, false, vm.state.HasFlag(faultState))
 	assert.Equal(t, 2, vm.estack.Len())
-	a := vm.estack.Peek(0).Array()
-	assert.Equal(t, len(reselements), len(a))
-	for k, v := range reselements {
-		e := a[k].Value().(*big.Int)
-		assert.Equal(t, int64(v), e.Int64())
-	}
-	assert.Equal(t, int64(1), vm.estack.Peek(1).BigInt().Int64())
+	assert.Equal(t, makeStackItem(reselements), vm.estack.Pop().value)
+	assert.Equal(t, makeStackItem(1), vm.estack.Pop().value)
 }
 
 func makeProgram(opcodes ...Instruction) []byte {