From 2eea2a6d9fc5503b83bbd2ec37ef65d3d655ae97 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Fri, 8 Apr 2011 16:58:04 +0200
Subject: [PATCH] generate augmented assignments to local variables

---
 .../jet/codegen/ExpressionCodegen.java        | 29 +++++++++++++++++++
 .../org/jetbrains/jet/lexer/JetTokens.java    |  2 ++
 .../jet/codegen/NamespaceGenTest.java         |  6 ++++
 3 files changed, 37 insertions(+)

diff --git a/idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java b/idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
index 116c5b06d2d..32891d62c09 100644
--- a/idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
+++ b/idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
@@ -494,6 +494,10 @@ public class ExpressionCodegen extends JetVisitor {
             generateAssignmentExpression(expression);
             return;
         }
+        if (JetTokens.AUGMENTED_ASSIGNMENTS.contains(opToken)) {
+            generateAugmentedAssignment(expression);
+            return;
+        }
         DeclarationDescriptor op = bindingContext.resolveReferenceExpression(expression.getOperationReference());
         if (op instanceof FunctionDescriptor) {
             JetType returnType = bindingContext.getExpressionType(expression);
@@ -602,6 +606,31 @@ public class ExpressionCodegen extends JetVisitor {
         }
     }
 
+    private void generateAugmentedAssignment(JetBinaryExpression expression) {
+        final JetExpression lhs = expression.getLeft();
+        if (lhs instanceof JetReferenceExpression) {
+            DeclarationDescriptor op = bindingContext.resolveReferenceExpression(expression.getOperationReference());
+            final JetType leftType = bindingContext.getExpressionType(lhs);
+            final Type asmType = typeMapper.mapType(leftType);
+            if (isNumberPrimitive(asmType)) {
+                final DeclarationDescriptor declarationDescriptor = bindingContext.resolveReferenceExpression((JetReferenceExpression) lhs);
+                final int index = myMap.getIndex(declarationDescriptor);
+                assert index >= 0;
+                v.load(index, asmType);
+                gen(expression.getRight(), asmType);
+                int opcode = opcodeForMethod(op.getName());
+                v.visitInsn(asmType.getOpcode(opcode));
+                v.store(index, asmType);
+            }
+            else {
+                throw new UnsupportedOperationException("Don't know how to generate augmented assignment for non-numeric types");
+            }
+        }
+        else {
+            throw new UnsupportedOperationException("Don't know how to generate augmented assignment to " + lhs.getText());
+        }
+    }
+
     @Override
     public void visitPrefixExpression(JetPrefixExpression expression) {
         DeclarationDescriptor op = bindingContext.resolveReferenceExpression(expression.getOperationSign());
diff --git a/idea/src/org/jetbrains/jet/lexer/JetTokens.java b/idea/src/org/jetbrains/jet/lexer/JetTokens.java
index 7ec18826e4d..d93bda8d827 100644
--- a/idea/src/org/jetbrains/jet/lexer/JetTokens.java
+++ b/idea/src/org/jetbrains/jet/lexer/JetTokens.java
@@ -167,4 +167,6 @@ public interface JetTokens {
             QUEST, COLON,
             RANGE, EQ, MULTEQ, DIVEQ, PERCEQ, PLUSEQ, MINUSEQ,
             NOT_IN, NOT_IS, HASH, IDENTIFIER, LABEL_IDENTIFIER, ATAT, AT);
+
+    TokenSet AUGMENTED_ASSIGNMENTS = TokenSet.create(PLUSEQ, MINUSEQ, MULTEQ, PERCEQ, DIVEQ);
 }
diff --git a/idea/tests/org/jetbrains/jet/codegen/NamespaceGenTest.java b/idea/tests/org/jetbrains/jet/codegen/NamespaceGenTest.java
index 61f6f489448..9ad86f29b8b 100644
--- a/idea/tests/org/jetbrains/jet/codegen/NamespaceGenTest.java
+++ b/idea/tests/org/jetbrains/jet/codegen/NamespaceGenTest.java
@@ -413,6 +413,12 @@ public class NamespaceGenTest extends LightCodeInsightFixtureTestCase {
         binOpTest("fun foo(a: Double, b: Int): Double = a + b", 1.0, 2, 3.0);
     }
 
+    public void testAugAssign() throws Exception {
+        loadText("fun foo(a: Int): Int { var x = a; x += 5; return x; }");
+        final Method main = generateFunction();
+        assertEquals(10, main.invoke(null, 5));
+    }
+
     private void binOpTest(final String text, final Object arg1, final Object arg2, final Object expected) throws Exception {
         loadText(text);
         System.out.println(generateToText());