compiler: fix struct field increment

Close #3981.

Signed-off-by: Tural Devrishev <tural@nspcc.ru>

Author: Tural Devrishev

pkg/compiler/codegen.go

@@ -238,6 +238,29 @@ func (c *codegen) emitStoreStructField(i int) {
 	emit.Opcodes(c.prog.BinWriter, opcode.ROT, opcode.SETITEM)
 }
 
+func (c *codegen) emitStoreSelectorExpr(n *ast.SelectorExpr) {
+	typ := c.typeOf(n.X)
+	if c.isInvalidType(typ) {
+		// This is a global variable from a package.
+		c.emitStoreVar(n.X.(*ast.Ident).Name, n.Sel.Name)
+		return
+	}
+	strct, ok := c.getStruct(typ)
+	if !ok {
+		c.prog.Err = fmt.Errorf("nested selector assigns not supported yet")
+		return
+	}
+	ast.Walk(c, n.X)
+	i := indexOfStruct(strct, n.Sel.Name)
+	c.emitStoreStructField(i)
+}
+
+func (c *codegen) emitStoreIndexExpr(n *ast.IndexExpr) {
+	ast.Walk(c, n.X)
+	ast.Walk(c, n.Index)
+	emit.Opcodes(c.prog.BinWriter, opcode.ROT, opcode.SETITEM)
+}
+
 // getVarIndex returns variable type and position in the corresponding slot,
 // according to the current scope.
 func (c *codegen) getVarIndex(pkg string, name string) *varInfo {
@@ -711,30 +734,15 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 				if !isAssignOp && !isMapKeyCheck {
 					ast.Walk(c, n.Rhs[i])
 				}
-				typ := c.typeOf(t.X)
-				if c.isInvalidType(typ) {
-					// Store to other package global variable.
-					c.emitStoreVar(t.X.(*ast.Ident).Name, t.Sel.Name)
-					return nil
-				}
-				strct, ok := c.getStruct(typ)
-				if !ok {
-					c.prog.Err = fmt.Errorf("nested selector assigns not supported yet")
-					return nil
-				}
-				ast.Walk(c, t.X)                      // load the struct
-				i := indexOfStruct(strct, t.Sel.Name) // get the index of the field
-				c.emitStoreStructField(i)             // store the field
+				c.emitStoreSelectorExpr(t)
 
 			// Assignments to index expressions.
 			// slice[0] = 10
 			case *ast.IndexExpr:
 				if !isAssignOp && !isMapKeyCheck {
 					ast.Walk(c, n.Rhs[i])
 				}
-				ast.Walk(c, t.X)
-				ast.Walk(c, t.Index)
-				emit.Opcodes(c.prog.BinWriter, opcode.ROT, opcode.SETITEM)
+				c.emitStoreIndexExpr(t)
 			}
 		}
 		return nil
@@ -1207,11 +1215,17 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		ast.Walk(c, n.X)
 		c.emitToken(n.Tok, c.typeOf(n.X))
 
-		// For now, only identifiers are supported for (post) for stmts.
-		// for i := 0; i < 10; i++ {}
-		// Where the post stmt is ( i++ )
-		if ident, ok := n.X.(*ast.Ident); ok {
-			c.emitStoreVar("", ident.Name)
+		switch t := n.X.(type) {
+		case *ast.Ident:
+			c.emitStoreVar("", t.Name)
+
+		case *ast.SelectorExpr:
+			// myStruct.t--
+			c.emitStoreSelectorExpr(t)
+
+		case *ast.IndexExpr:
+			// slice[0]++
+			c.emitStoreIndexExpr(t)
 		}
 		return nil
 

pkg/compiler/slice_test.go

@@ -338,6 +338,60 @@ var sliceTestCases = []testCase{
 		`,
 		[]byte("str"),
 	},
+	{
+		"increment elem, alias",
+		`func F%d() int {
+			s := []int{0}
+			t := s
+			t[0]++
+			return s[0]
+		}
+		`,
+		big.NewInt(1),
+	},
+	{
+		"increment elem, struct",
+		`type T struct { V int }
+		func F%d() int {
+			s := []T{{0}}
+			s[0].V++
+			return s[0].V
+		}
+		`,
+		big.NewInt(1),
+	},
+	{
+		"increment elem, pointer",
+		`func F%d() int {
+			t := &T{0}
+			s := []*T{t}
+			s[0].V++
+			return t.V
+		}
+		`,
+		big.NewInt(1),
+	},
+	{
+		"increment elem, 2d slice",
+		`func F%d() int {
+			s := [][]int{{0}}
+			s[0][0]++
+			return s[0][0]
+		}
+		`,
+		big.NewInt(1),
+	},
+	{
+		"increment elem, func call",
+		`func f(s []int) []int { return s }
+		func F%d() int {
+			var s = []int{0}
+			f(s)[0]++
+			return s[0]
+		}
+		`,
+		big.NewInt(1),
+	},
 }
 
 func TestSliceOperations(t *testing.T) {

pkg/compiler/struct_test.go

@@ -431,6 +431,55 @@ var structTestCases = []testCase{
 		`,
 		big.NewInt(2),
 	},
+	{
+		"increment field, simple",
+		`type test struct { t int }
+			func F%d() int {
+				var myStruct test
+				myStruct.t++
+				return myStruct.t
+			}
+		`,
+		big.NewInt(1),
+	},
+	{
+		"increment field, nested",
+		`type nested struct { t int }
+			type external struct { n nested }
+			func F%d() int {
+				var myStruct external
+				myStruct.n.t++
+				return myStruct.n.t
+			}
+		`,
+		big.NewInt(1),
+	},
+	{
+		"increment field, selector chain",
+		`type C struct { v int }
+			type B struct { c C }
+			type A struct { b B }
+			func F%d() int {
+				var a A
+				a.b.c.v++
+				return a.b.c.v
+			}
+		`,
+		big.NewInt(1),
+	},
+	{
+		"incremental field, func call",
+		`func f(t *test) *test {
+				return t
+			}
+			func F%d() int {
+				var myStruct = &test{}
+				f(myStruct).t++
+				return myStruct.t
+			}
+		`,
+		big.NewInt(1),
+	},
 }
 
 func TestStructs(t *testing.T) {