Create From Usage ("set" operation): Support expressions of the form a[i]++

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetOperationExpression.java

@@ -16,9 +16,10 @@
 
 package org.jetbrains.jet.lang.psi;
 
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 
-public interface JetOperationExpression {
+public interface JetOperationExpression extends PsiElement {
     @NotNull
     JetSimpleNameExpression getOperationReference();
 }

idea/idea-analysis/src/org/jetbrains/jet/plugin/intentions/SimplifyNegatedBinaryExpressionIntention.kt

@@ -29,6 +29,7 @@ import org.jetbrains.jet.lexer.JetTokens
 import org.jetbrains.jet.lang.psi.JetPsiUtil
 import org.jetbrains.jet.plugin.JetBundle
 import org.jetbrains.jet.lang.psi.JetPsiFactory
+import com.intellij.psi.PsiElement
 
 public class SimplifyNegatedBinaryExpressionIntention : JetSelfTargetingIntention<JetPrefixExpression>("simplify.negated.binary.expression", javaClass()) {
 
@@ -87,7 +88,7 @@ public class SimplifyNegatedBinaryExpressionIntention : JetSelfTargetingIntentio
                             expression.getRight()
                     )
                     else -> throw IllegalStateException(
-                         "Expression is neither a JetIsExpression or JetBinaryExpression (checked by isApplicableTo): ${expression.getText()}"
+                         "Expression is neither a JetIsExpression or JetBinaryExpression (checked by isApplicableTo): ${(expression : PsiElement).getText()}"
                     )
                 }
         )

idea/src/org/jetbrains/jet/plugin/intentions/ConvertToForEachFunctionCallIntention.kt

@@ -23,6 +23,7 @@ import org.jetbrains.jet.lang.psi.JetBlockExpression
 import org.jetbrains.jet.lang.psi.JetElement
 import org.jetbrains.jet.lang.psi.JetParameter
 import org.jetbrains.jet.lang.psi.JetOperationExpression
+import com.intellij.psi.PsiElement
 
 public class ConvertToForEachFunctionCallIntention : JetSelfTargetingIntention<JetForExpression>("convert.to.for.each.function.call.intention", javaClass()) {
     override fun isApplicableTo(element: JetForExpression): Boolean {
@@ -49,7 +50,7 @@ public class ConvertToForEachFunctionCallIntention : JetSelfTargetingIntention<J
             val loopRange = element.getLoopRange()!!
 
             return when (loopRange) {
-                is JetOperationExpression -> "(${loopRange.getText()})"
+                is JetOperationExpression -> "(${(loopRange : PsiElement).getText()})"
                 else -> loopRange.getText() ?: throw AssertionError("LoopRange in ForExpression shouldn't be empty: expressionText = ${element.getText()}")
             }
         }

idea/src/org/jetbrains/jet/plugin/quickfix/createFromUsage/createFunction/CreateSetFunctionActionFactory.kt

@@ -11,6 +11,12 @@ import org.jetbrains.jet.lang.psi.JetBinaryExpression
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns
 import org.jetbrains.jet.plugin.quickfix.createFromUsage.callableBuilder.*
 import java.util.Collections
+import org.jetbrains.jet.lang.psi.JetOperationExpression
+import org.jetbrains.jet.lang.psi.JetUnaryExpression
+import org.jetbrains.jet.lang.types.expressions.OperatorConventions
+import org.jetbrains.jet.lang.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache
+import org.jetbrains.jet.lang.types.ErrorUtils
 
 object CreateSetFunctionActionFactory : JetSingleIntentionActionFactory() {
     override fun createAction(diagnostic: Diagnostic): IntentionAction? {
@@ -18,16 +24,29 @@ object CreateSetFunctionActionFactory : JetSingleIntentionActionFactory() {
         val arrayExpr = accessExpr.getArrayExpression() ?: return null
         val arrayType = TypeInfo(arrayExpr, Variance.IN_VARIANCE)
 
+        val builtIns = KotlinBuiltIns.getInstance()
+
         val parameters = accessExpr.getIndexExpressions().mapTo(ArrayList<ParameterInfo>()) {
             ParameterInfo(TypeInfo(it, Variance.IN_VARIANCE))
         }
 
-        val assignmentExpr = QuickFixUtil.getParentElementOfType(diagnostic, javaClass<JetBinaryExpression>()) ?: return null
-        val rhs = assignmentExpr.getRight() ?: return null
-        val valType = TypeInfo(rhs, Variance.IN_VARIANCE)
+        val assignmentExpr = QuickFixUtil.getParentElementOfType(diagnostic, javaClass<JetOperationExpression>()) ?: return null
+        val valType = when (assignmentExpr) {
+            is JetBinaryExpression -> {
+                TypeInfo(assignmentExpr.getRight() ?: return null, Variance.IN_VARIANCE)
+            }
+            is JetUnaryExpression -> {
+                if (assignmentExpr.getOperationToken() !in OperatorConventions.INCREMENT_OPERATIONS) return null
+
+                val context = AnalyzerFacadeWithCache.getContextForElement(assignmentExpr)
+                val rhsType = assignmentExpr.getResolvedCall(context)?.getResultingDescriptor()?.getReturnType()
+                TypeInfo(if (rhsType == null || ErrorUtils.containsErrorType(rhsType)) builtIns.getAnyType() else rhsType, Variance.IN_VARIANCE)
+            }
+            else -> return null
+        }
         parameters.add(ParameterInfo(valType, "value"))
 
-        val returnType = TypeInfo(KotlinBuiltIns.getInstance().getUnitType(), Variance.OUT_VARIANCE)
+        val returnType = TypeInfo(builtIns.getUnitType(), Variance.OUT_VARIANCE)
         return CreateCallableFromUsageFix(accessExpr, FunctionInfo("set", arrayType, returnType, Collections.emptyList(), parameters))
     }
 }

idea/testData/quickfix/createFromUsage/createFunction/set/afterSetterForIncrement.kt

@@ -0,0 +1,13 @@
+// "Create function 'set' from usage" "true"
+class A {
+    fun get(s: String): Int = 1
+
+    fun set(s: String, value: Int) {
+        throw UnsupportedOperationException("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+}
+
+fun foo() {
+    var a = A()
+    a["1"]++
+}

idea/testData/quickfix/createFromUsage/createFunction/set/beforeSetterForIncrement.kt

@@ -0,0 +1,9 @@
+// "Create function 'set' from usage" "true"
+class A {
+    fun get(s: String): Int = 1
+}
+
+fun foo() {
+    var a = A()
+    a<caret>["1"]++
+}

idea/tests/org/jetbrains/jet/plugin/quickfix/QuickFixTestGenerated.java

@@ -1197,6 +1197,13 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                     String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createFunction/set/beforeCreateSetFromUsage2.kt");
                     doTest(fileName);
                 }
+                
+                @TestMetadata("beforeSetterForIncrement.kt")
+                public void testSetterForIncrement() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createFunction/set/beforeSetterForIncrement.kt");
+                    doTest(fileName);
+                }
+                
             }
 
             @TestMetadata("idea/testData/quickfix/createFromUsage/createFunction/unaryOperations")