Merge pull request #566 from nspcc-dev/feature/varindex

compiler: implement assignment to a variable index

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

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)
+}
+