diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go index 39659d6b3..781d5a5a0 100644 --- a/pkg/compiler/codegen.go +++ b/pkg/compiler/codegen.go @@ -384,34 +384,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.emitLoadVar(t.Name) - ast.Walk(c, n.Rhs[0]) // can only add assign to 1 expr on the RHS - c.convertToken(n.Tok) - c.emitStoreVar(t.Name) - case token.DEFINE: + if n.Tok == token.DEFINE { if !multiRet { c.registerDebugVariable(t.Name, n.Rhs[i]) } if t.Name != "_" { c.scope.newLocal(t.Name) } - fallthrough - default: - if i == 0 || !multiRet { - ast.Walk(c, n.Rhs[i]) - } - c.emitStoreVar(t.Name) } + if !isAssignOp && (i == 0 || !multiRet) { + ast.Walk(c, n.Rhs[i]) + } + c.emitStoreVar(t.Name) 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]) + } if strct, ok := c.typeOf(expr).Underlying().(*types.Struct); ok { c.emitLoadVar(expr.Name) // load the struct i := indexOfStruct(strct, t.Sel.Name) // get the index of the field @@ -425,7 +427,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.emitLoadVar(name) ast.Walk(c, t.Index) diff --git a/pkg/compiler/slice_test.go b/pkg/compiler/slice_test.go index a853f6e00..629b1b52d 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 aa0da998e..f50390899 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", `