compiler: allow using OP= with struct fields and slice elements

Do it in a generic way, there is no need in restricting to only
variables.
Port of #954.
This commit is contained in:
Evgenii Stratonikov 2020-05-18 11:45:20 +03:00
parent fdb217ec81
commit 1b105a9f1d
3 changed files with 40 additions and 15 deletions

View file

@ -384,34 +384,36 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
case *ast.AssignStmt: case *ast.AssignStmt:
multiRet := len(n.Rhs) != len(n.Lhs) multiRet := len(n.Rhs) != len(n.Lhs)
c.saveSequencePoint(n) 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++ { for i := 0; i < len(n.Lhs); i++ {
switch t := n.Lhs[i].(type) { switch t := n.Lhs[i].(type) {
case *ast.Ident: case *ast.Ident:
switch n.Tok { if n.Tok == token.DEFINE {
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 !multiRet { if !multiRet {
c.registerDebugVariable(t.Name, n.Rhs[i]) c.registerDebugVariable(t.Name, n.Rhs[i])
} }
if t.Name != "_" { if t.Name != "_" {
c.scope.newLocal(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: case *ast.SelectorExpr:
switch expr := t.X.(type) { switch expr := t.X.(type) {
case *ast.Ident: 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 { if strct, ok := c.typeOf(expr).Underlying().(*types.Struct); ok {
c.emitLoadVar(expr.Name) // load the struct c.emitLoadVar(expr.Name) // load the struct
i := indexOfStruct(strct, t.Sel.Name) // get the index of the field 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. // Assignments to index expressions.
// slice[0] = 10 // slice[0] = 10
case *ast.IndexExpr: case *ast.IndexExpr:
ast.Walk(c, n.Rhs[i]) if !isAssignOp {
ast.Walk(c, n.Rhs[i])
}
name := t.X.(*ast.Ident).Name name := t.X.(*ast.Ident).Name
c.emitLoadVar(name) c.emitLoadVar(name)
ast.Walk(c, t.Index) ast.Walk(c, t.Index)

View file

@ -33,6 +33,16 @@ var sliceTestCases = []testCase{
`, `,
big.NewInt(42), 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", "complex test",
` `

View file

@ -134,6 +134,17 @@ var structTestCases = []testCase{
}`, }`,
big.NewInt(14), 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", "assign a struct field to a struct field",
` `