diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go index c2c312c6b..032d3d2f1 100644 --- a/pkg/compiler/codegen.go +++ b/pkg/compiler/codegen.go @@ -317,38 +317,36 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor { case *ast.AssignStmt: multiRet := len(n.Rhs) != len(n.Lhs) c.saveSequencePoint(n) + // Assign operations are grouped https://github.com/golang/go/blob/master/src/go/types/stmt.go#L160 + isAssignOp := token.ADD_ASSIGN <= n.Tok && n.Tok <= token.AND_NOT_ASSIGN + if isAssignOp { + // RHS can contain exactly one expression, thus there is no need to iterate. + ast.Walk(c, n.Lhs[0]) + ast.Walk(c, n.Rhs[0]) + c.convertToken(n.Tok) + } for i := 0; i < len(n.Lhs); i++ { switch t := n.Lhs[i].(type) { case *ast.Ident: - switch n.Tok { - case token.ADD_ASSIGN, token.SUB_ASSIGN, token.MUL_ASSIGN, token.QUO_ASSIGN, token.REM_ASSIGN: - c.emitLoadLocal(t.Name) - ast.Walk(c, n.Rhs[0]) // can only add assign to 1 expr on the RHS - c.convertToken(n.Tok) + if n.Tok == token.DEFINE && !multiRet { + c.registerDebugVariable(t.Name, n.Rhs[i]) + } + if !isAssignOp && (i == 0 || !multiRet) { + ast.Walk(c, n.Rhs[i]) + } + if t.Name == "_" { + emit.Opcode(c.prog.BinWriter, opcode.DROP) + } else { l := c.scope.loadLocal(t.Name) c.emitStoreLocal(l) - case token.DEFINE: - if !multiRet { - c.registerDebugVariable(t.Name, n.Rhs[i]) - } - fallthrough - default: - if i == 0 || !multiRet { - ast.Walk(c, n.Rhs[i]) - } - - if t.Name == "_" { - emit.Opcode(c.prog.BinWriter, opcode.DROP) - } else { - l := c.scope.loadLocal(t.Name) - c.emitStoreLocal(l) - } } case *ast.SelectorExpr: switch expr := t.X.(type) { case *ast.Ident: - ast.Walk(c, n.Rhs[i]) + if !isAssignOp { + ast.Walk(c, n.Rhs[i]) + } typ := c.typeInfo.ObjectOf(expr).Type().Underlying() if strct, ok := typ.(*types.Struct); ok { c.emitLoadLocal(expr.Name) // load the struct @@ -363,7 +361,9 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor { // Assignments to index expressions. // slice[0] = 10 case *ast.IndexExpr: - ast.Walk(c, n.Rhs[i]) + if !isAssignOp { + ast.Walk(c, n.Rhs[i]) + } name := t.X.(*ast.Ident).Name c.emitLoadLocal(name) switch ind := t.Index.(type) { diff --git a/pkg/compiler/slice_test.go b/pkg/compiler/slice_test.go index 40600ea28..ca4c25798 100644 --- a/pkg/compiler/slice_test.go +++ b/pkg/compiler/slice_test.go @@ -33,6 +33,16 @@ var sliceTestCases = []testCase{ `, big.NewInt(42), }, + { + "increase slice element with +=", + `package foo + func Main() int { + a := []int{1, 2, 3} + a[1] += 40 + return a[1] + }`, + big.NewInt(42), + }, { "complex test", ` diff --git a/pkg/compiler/struct_test.go b/pkg/compiler/struct_test.go index 1dd63e07f..d945ab775 100644 --- a/pkg/compiler/struct_test.go +++ b/pkg/compiler/struct_test.go @@ -134,6 +134,17 @@ var structTestCases = []testCase{ }`, big.NewInt(14), }, + { + "increase struct field with +=", + `package foo + type token struct { x int } + func Main() int { + t := token{x: 2} + t.x += 3 + return t.x + }`, + big.NewInt(5), + }, { "assign a struct field to a struct field", `