Fix augmented assignment evaluation

Author: msimacek

graalpython/com.oracle.graal.python.test/src/com/oracle/graal/python/test/compiler/CompilerTests.java

@@ -95,6 +95,16 @@ public void testAugAssignment() {
         doTest("a += 12.0");
     }
 
+    @Test
+    public void testAugAssignmentAttr() {
+        doTest("a.b += 12.0");
+    }
+
+    @Test
+    public void testAugAssignmentSubscr() {
+        doTest("a[b] += 12.0");
+    }
+
     @Test
     public void testAnnAssignment() {
         doTest("a: int = 12");

graalpython/com.oracle.graal.python.test/src/tests/test_assign.py

@@ -105,6 +105,71 @@ def test_destructuring():
     assert a == -1 and b == -1 and s == [0, 1, 2, 3, 4, 5, 6, 7] and c == 8 and d == 9
 
 
+def test_augassign_evaluation_subsc():
+    calls = []
+
+    class C(list):
+        def __getitem__(self, item):
+            calls.append("get")
+            return super().__getitem__(item)
+
+        def __setitem__(self, key, value):
+            calls.append("set")
+            super().__setitem__(key, value)
+
+    class I:
+        def __iadd__(self, other):
+            calls.append("iadd")
+            return 3
+
+    def index():
+        calls.append("index")
+        return 0
+
+    def value():
+        calls.append("value")
+        return 1
+
+    x = C([I()])
+
+    def container():
+        calls.append("container")
+        return x
+
+    container()[index()] += value()
+    assert calls == ["container", "index", "get", "value", "iadd", "set"]
+
+
+def test_augassign_evaluation_attr():
+    calls = []
+
+    class C(list):
+        def __getattr__(self, item):
+            calls.append("get")
+            return I()
+
+        def __setattr__(self, key, value):
+            calls.append("set")
+
+    class I:
+        def __iadd__(self, other):
+            calls.append("iadd")
+            return 3
+
+    def value():
+        calls.append("value")
+        return 1
+
+    x = C()
+
+    def container():
+        calls.append("container")
+        return x
+
+    container().attr += value()
+    assert calls == ["container", "get", "value", "iadd", "set"]
+
+
 def test_assigning_hidden_keys():
     class A():
         def __init__(self):

graalpython/com.oracle.graal.python.test/testData/goldenFiles/CompilerTests/testAugAssignmentAttr.co

@@ -0,0 +1,10 @@
+Disassembly of <module>:
+000000         0 LOAD_NAME                         0    (a)
+000000         2 DUP_TOP
+000000         3 LOAD_ATTR                         1    (b)
+000007         5 LOAD_DOUBLE_O                     0    (12.0)
+000000         7 BINARY_OP                         1    (INPLACE_ADD)
+000000         9 ROT_TWO
+000000        10 STORE_ATTR                        1    (b)
+000000        12 LOAD_NONE
+000000        13 RETURN_VALUE

graalpython/com.oracle.graal.python.test/testData/goldenFiles/CompilerTests/testAugAssignmentSubscr.co

@@ -0,0 +1,13 @@
+Disassembly of <module>:
+000000         0 LOAD_NAME                         0    (a)
+000000         2 DUP_TOP
+000002         3 LOAD_NAME                         1    (b)
+000000         5 DUP_TOP
+000000         6 ROT_THREE
+000000         7 BINARY_SUBSCR
+000008         8 LOAD_DOUBLE_O                     0    (12.0)
+000000        10 BINARY_OP                         1    (INPLACE_ADD)
+000000        12 ROT_THREE
+000000        13 STORE_SUBSCR
+000000        14 LOAD_NONE
+000000        15 RETURN_VALUE

graalpython/com.oracle.graal.python.test/testData/goldenFiles/CompilerTests/testBenchmark.co

@@ -49,8 +49,8 @@ Disassembly of docompute:
 000183        71 LOAD_BYTE_I                       1                                      can quicken
 000178        73 BINARY_OP                         1    (INPLACE_ADD)                     can quicken, generalizes: 71, 69
 000178        75 STORE_FAST                        3    (j)                               generalizes: 73
-000178        77 JUMP_BACKWARD                     58   (to 19)
-000178   >>   79 JUMP_BACKWARD                     72   (to 7)
+000086        77 JUMP_BACKWARD                     58   (to 19)
+000086   >>   79 JUMP_BACKWARD                     72   (to 7)
 000197   >>   81 LOAD_FAST                         2    (sum_)
 000190        83 RETURN_VALUE
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/Compiler.java

@@ -140,7 +140,6 @@
 import java.util.List;
 
 import com.oracle.graal.python.builtins.objects.PNone;
-import com.oracle.graal.python.pegparser.CopyWithContextVisitor;
 import com.oracle.graal.python.pegparser.ErrorCallback;
 import com.oracle.graal.python.pegparser.ErrorCallback.ErrorType;
 import com.oracle.graal.python.pegparser.FExprParser;
@@ -2020,11 +2019,30 @@ public Void visit(StmtTy.AsyncWith node) {
 
     @Override
     public Void visit(StmtTy.AugAssign node) {
-        SourceRange savedLocation = setLocation(node);
-        node.target.accept(new CopyWithContextVisitor(ExprContextTy.Load)).accept(this);
-        setLocation(savedLocation);
-        node.value.accept(this);
+        ExprTy target = node.target;
+        SourceRange savedLocation = setLocation(target);
+        if (target instanceof ExprTy.Attribute) {
+            ExprTy.Attribute attr = (ExprTy.Attribute) target;
+            attr.value.accept(this);
+            addOp(DUP_TOP);
+            addOpName(LOAD_ATTR, unit.names, attr.attr);
+        } else if (target instanceof ExprTy.Subscript) {
+            ExprTy.Subscript subscript = (ExprTy.Subscript) target;
+            subscript.value.accept(this);
+            addOp(DUP_TOP);
+            subscript.slice.accept(this);
+            addOp(DUP_TOP);
+            addOp(ROT_THREE);
+            addOp(BINARY_SUBSCR);
+        } else if (target instanceof ExprTy.Name) {
+            ExprTy.Name name = (ExprTy.Name) target;
+            addNameOp(name.id, ExprContextTy.Load);
+        } else {
+            // TODO this raises SystemError under CPython
+            throw new IllegalArgumentException("invalid node type for augmented assignment");
+        }
         setLocation(node);
+        node.value.accept(this);
         switch (node.op) {
             case Add:
                 addOp(BINARY_OP, BinaryOps.INPLACE_ADD.ordinal());
@@ -2068,7 +2086,20 @@ public Void visit(StmtTy.AugAssign node) {
             default:
                 throw new IllegalStateException("Unknown binary inplace operation " + node.op);
         }
-        return node.target.accept(this);
+        setLocation(target);
+        if (target instanceof ExprTy.Attribute) {
+            ExprTy.Attribute attr = (ExprTy.Attribute) target;
+            addOp(ROT_TWO);
+            addOpName(STORE_ATTR, unit.names, attr.attr);
+        } else if (target instanceof ExprTy.Subscript) {
+            addOp(ROT_THREE);
+            addOp(STORE_SUBSCR);
+        } else {
+            ExprTy.Name name = (ExprTy.Name) target;
+            addNameOp(name.id, ExprContextTy.Store);
+        }
+        setLocation(savedLocation);
+        return null;
     }
 
     @Override