From 891a878af173c01f758db6338e6b60f51e19bdf4 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 19 Dec 2019 10:48:06 +0300
Subject: [PATCH] compiler: implement assignment to a variable index

Fixes #564.
---
 pkg/compiler/codegen.go    | 22 +++++++++++-----
 pkg/vm/tests/slice_test.go | 53 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 7 deletions(-)
 create mode 100644 pkg/vm/tests/slice_test.go

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 42970458b..06dacd81b 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -269,15 +269,23 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 				ast.Walk(c, n.Rhs[i])
 				name := t.X.(*ast.Ident).Name
 				c.emitLoadLocal(name)
-				// For now storm only supports basic index operations. Hence we
-				// cast this to an *ast.BasicLit (1, 2 , 3)
-				indexStr := t.Index.(*ast.BasicLit).Value
-				index, err := strconv.Atoi(indexStr)
-				if err != nil {
-					c.prog.Err = fmt.Errorf("failed to convert slice index to integer")
+				switch ind := t.Index.(type) {
+				case *ast.BasicLit:
+					indexStr := ind.Value
+					index, err := strconv.Atoi(indexStr)
+					if err != nil {
+						c.prog.Err = fmt.Errorf("failed to convert slice index to integer")
+						return nil
+					}
+					c.emitStoreStructField(index)
+				case *ast.Ident:
+					c.emitLoadLocal(ind.Name)
+					emitOpcode(c.prog.BinWriter, opcode.ROT)
+					emitOpcode(c.prog.BinWriter, opcode.SETITEM)
+				default:
+					c.prog.Err = fmt.Errorf("unsupported index expression")
 					return nil
 				}
-				c.emitStoreStructField(index)
 			}
 		}
 		return nil
diff --git a/pkg/vm/tests/slice_test.go b/pkg/vm/tests/slice_test.go
new file mode 100644
index 000000000..8d0d67f9d
--- /dev/null
+++ b/pkg/vm/tests/slice_test.go
@@ -0,0 +1,53 @@
+package vm_test
+
+import (
+	"math/big"
+	"testing"
+)
+
+var sliceTestCases = []testCase{
+	{
+		"constant index",
+		`
+		package foo
+		func Main() int {
+			a := []int{0,0}
+			a[1] = 42
+			return a[1]+0
+		}
+		`,
+		big.NewInt(42),
+	},
+	{
+		"variable index",
+		`
+		package foo
+		func Main() int {
+			a := []int{0,0}
+			i := 1
+			a[i] = 42
+			return a[1]+0
+		}
+		`,
+		big.NewInt(42),
+	},
+	{
+		"complex test",
+		`
+		package foo
+		func Main() int {
+			a := []int{1,2,3}
+			x := a[0]
+			a[x] = a[x] + 4
+			a[x] = a[x] + a[2]
+			return a[1]
+		}
+		`,
+		big.NewInt(9),
+	},
+}
+
+func TestSliceOperations(t *testing.T) {
+	runTestCases(t, sliceTestCases)
+}
+