generate augmented assignments to local variables

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

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

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