From 8b1118f9071ebb3c1b7ff4592220ff71cd2ae0f8 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Thu, 14 Jul 2011 20:36:15 +0200
Subject: [PATCH] support augmented assignment when ordinary binary operation
 method is overloaded

---
 .../jet/codegen/ExpressionCodegen.java        | 31 +++++++++++--------
 .../classes/overloadPlusToPlusAssign.jet      | 30 ++++++++++++++++++
 .../jetbrains/jet/codegen/ClassGenTest.java   |  4 +++
 3 files changed, 52 insertions(+), 13 deletions(-)
 create mode 100644 idea/testData/codegen/classes/overloadPlusToPlusAssign.jet

diff --git a/idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java b/idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
index 958e027e8f2..021f4406352 100644
--- a/idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
+++ b/idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
@@ -1278,6 +1278,9 @@ public class ExpressionCodegen extends JetVisitor {
                 v.visitInsn(lhsType.getOpcode(opcodeForMethod(op.getName())));  // receiver result
                 value.store(v);
             }
+            else if (callable instanceof CallableMethod) {
+                callAugAssignMethod(expression, (CallableMethod) callable, lhsType, true);
+            }
             else {
                 throw new UnsupportedOperationException("Augmented assignment for non-primitive types not yet implemented");
             }
@@ -1288,23 +1291,25 @@ public class ExpressionCodegen extends JetVisitor {
                 intrinsic.generate(this, v, Type.VOID_TYPE, expression, Arrays.asList(lhs, expression.getRight()), false);
             }
             else {
-                CallableMethod method = (CallableMethod) callable;
-                FunctionDescriptor fd = (FunctionDescriptor) op;
-                StackValue value = generateIntermediateValue(lhs);
-                final boolean keepReturnValue = !fd.getReturnType().equals(JetStandardClasses.getUnitType());
-                if (keepReturnValue) {
-                    value.dupReceiver(v, 0);
-                }
-                value.put(lhsType, v);
-                genToJVMStack(expression.getRight());
-                method.invoke(v);
-                if (keepReturnValue) {
-                    value.store(v);
-                }
+                final boolean keepReturnValue = !((FunctionDescriptor) op).getReturnType().equals(JetStandardClasses.getUnitType());
+                callAugAssignMethod(expression, (CallableMethod) callable, lhsType, keepReturnValue);
             }
         }
     }
 
+    private void callAugAssignMethod(JetBinaryExpression expression, CallableMethod callable, Type lhsType, final boolean keepReturnValue) {
+        StackValue value = generateIntermediateValue(expression.getLeft());
+        if (keepReturnValue) {
+            value.dupReceiver(v, 0);
+        }
+        value.put(lhsType, v);
+        genToJVMStack(expression.getRight());
+        callable.invoke(v);
+        if (keepReturnValue) {
+            value.store(v);
+        }
+    }
+
     public void generateStringBuilderConstructor() {
         Type type = Type.getObjectType(CLASS_STRING_BUILDER);
         v.anew(type);
diff --git a/idea/testData/codegen/classes/overloadPlusToPlusAssign.jet b/idea/testData/codegen/classes/overloadPlusToPlusAssign.jet
new file mode 100644
index 00000000000..e81515821b0
--- /dev/null
+++ b/idea/testData/codegen/classes/overloadPlusToPlusAssign.jet
@@ -0,0 +1,30 @@
+import java.util.*
+
+class ArrayWrapper<T>() {
+    val contents = ArrayList<T>()
+
+    fun add(item: T) {
+        contents.add(item)
+    }
+
+    fun plus(rhs: ArrayWrapper<T>): ArrayWrapper<T> {
+        val result = ArrayWrapper<T>()
+        result.contents.addAll(contents)
+        result.contents.addAll(rhs.contents)
+        return result
+    }
+
+    fun get(index: Int): T {
+        return contents.get(index)
+    }
+}
+
+fun box(): String {
+    val v1 = ArrayWrapper<String>()
+    val v2 = ArrayWrapper<String>()
+    v1.add("foo")
+    val v3 = v1
+    v2.add("bar")
+    v1 += v2
+    return if (v1.contents.size() == 2 && v3.contents.size() == 1) "OK" else "fail"
+}
diff --git a/idea/tests/org/jetbrains/jet/codegen/ClassGenTest.java b/idea/tests/org/jetbrains/jet/codegen/ClassGenTest.java
index aac3d2ceb7a..2a9fd6c98e7 100644
--- a/idea/tests/org/jetbrains/jet/codegen/ClassGenTest.java
+++ b/idea/tests/org/jetbrains/jet/codegen/ClassGenTest.java
@@ -148,6 +148,10 @@ public class ClassGenTest extends CodegenTestCase {
         blackBoxFile("classes/overloadPlusAssignReturn.jet");
     }
 
+    public void testOverloadPlusToPlusAssign() throws Exception {
+        blackBoxFile("classes/overloadPlusToPlusAssign.jet");
+    }
+
     public void testEnumClass() throws Exception {
         loadText("enum class Direction { NORTH; SOUTH; EAST; WEST }");
         final Class direction = loadAllClasses(generateClassesInFile()).get("Direction");