support augmented assignment when ordinary binary operation method is overloaded

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;
+                final boolean keepReturnValue = !((FunctionDescriptor) op).getReturnType().equals(JetStandardClasses.getUnitType());
-                FunctionDescriptor fd = (FunctionDescriptor) op;
+                callAugAssignMethod(expression, (CallableMethod) callable, lhsType, keepReturnValue);
-                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);
-                }
             }
         }
     }
 
+    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);

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"
+}

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");