From 8b6bddca3c3f7ffc8f4005922bbd96022ec5bc5a Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 24 Sep 2019 15:58:31 +0300
Subject: [PATCH] vm: support Map in PICKITEM and SETITEM

---
 pkg/vm/vm.go      | 31 ++++++++++++++++++++-----------
 pkg/vm/vm_test.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 46e60a5f7..584b7c14e 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -705,11 +705,11 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		v.estack.PushVal(l)
 
 	case PICKITEM:
-		var (
-			key   = v.estack.Pop()
-			obj   = v.estack.Pop()
-			index = int(key.BigInt().Int64())
-		)
+		key := v.estack.Pop()
+		validateMapKey(key)
+
+		obj := v.estack.Pop()
+		index := int(key.BigInt().Int64())
 
 		switch t := obj.value.(type) {
 		// Struct and Array items have their underlying value as []StackItem.
@@ -720,6 +720,12 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 			}
 			item := arr[index]
 			v.estack.PushVal(item)
+		case *MapItem:
+			if !t.Has(key.value) {
+				panic("invalid key")
+			}
+			k := toMapKey(key.value)
+			v.estack.Push(&Element{value: t.value[k]})
 		default:
 			arr := obj.Bytes()
 			if index < 0 || index >= len(arr) {
@@ -730,21 +736,24 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		}
 
 	case SETITEM:
-		var (
-			item  = v.estack.Pop().value
-			key   = v.estack.Pop()
-			obj   = v.estack.Pop()
-			index = int(key.BigInt().Int64())
-		)
+		item := v.estack.Pop().value
+		key := v.estack.Pop()
+		validateMapKey(key)
+
+		obj := v.estack.Pop()
 
 		switch t := obj.value.(type) {
 		// Struct and Array items have their underlying value as []StackItem.
 		case *ArrayItem, *StructItem:
 			arr := t.Value().([]StackItem)
+			index := int(key.BigInt().Int64())
 			if index < 0 || index >= len(arr) {
 				panic("SETITEM: invalid index")
 			}
 			arr[index] = item
+		case *MapItem:
+			t.Add(key.value, item)
+
 		default:
 			panic(fmt.Sprintf("SETITEM: invalid item type %s", t))
 		}
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index 2c0842bae..ad5aa475d 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -646,6 +646,39 @@ func TestPICKITEMByteArray(t *testing.T) {
 	assert.Equal(t, makeStackItem(2), vm.estack.Pop().value)
 }
 
+func TestPICKITEMMap(t *testing.T) {
+	prog := makeProgram(PICKITEM)
+	vm := load(prog)
+
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(3))
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(makeStackItem(5))
+
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(3), vm.estack.Pop().value)
+}
+
+func TestSETITEMMap(t *testing.T) {
+	prog := makeProgram(SETITEM, PICKITEM)
+	vm := load(prog)
+
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(3))
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(5)
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(5)
+	vm.estack.PushVal([]byte{0, 1})
+
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem([]byte{0, 1}), vm.estack.Pop().value)
+}
+
 func TestSIZENoArgument(t *testing.T) {
 	prog := makeProgram(SIZE)
 	vm := load(prog)