Add "replace map.put with assignment" to J2K

Related to KT-21502

idea/src/org/jetbrains/kotlin/idea/inspections/ReplacePutWithAssignmentInspection.kt

@@ -21,68 +21,68 @@ import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.builtins.DefaultBuiltIns
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.callExpression
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.endOffset
-import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getExplicitReceiverValue
 import org.jetbrains.kotlin.resolve.descriptorUtil.isSubclassOf
 
 class ReplacePutWithAssignmentInspection : AbstractKotlinInspection() {
-    private val compatibleNames = setOf("put")
-
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor {
         return object : KtVisitorVoid() {
             override fun visitDotQualifiedExpression(expression: KtDotQualifiedExpression) {
                 super.visitDotQualifiedExpression(expression)
-                val callExpression = expression.callExpression ?: return
+                if (isActiveFor(expression)) {
-                if (callExpression.valueArguments.size != 2) return
+                    holder.registerProblem(
-
+                            expression.callExpression!!.calleeExpression!!,
-                val calleeExpression = callExpression.calleeExpression as? KtNameReferenceExpression ?: return
-                if (calleeExpression.getReferencedName() !in compatibleNames) return
-
-                val context = expression.analyze()
-                if (expression.isUsedAsExpression(context)) return
-                val resolvedCall = expression.getResolvedCall(context) ?: return
-                val receiverType = resolvedCall.getExplicitReceiverValue()?.type ?: return
-                val receiverClass = receiverType.constructor.declarationDescriptor as? ClassDescriptor ?: return
-                if (receiverClass.isSubclassOf(DefaultBuiltIns.Instance.mutableMap)) {
-                    val argumentOffset = expression.startOffset
-                    val problemDescriptor = holder.manager.createProblemDescriptor(
-                            calleeExpression,
-                            TextRange(expression.startOffset - argumentOffset,
-                                      callExpression.endOffset - argumentOffset),
                             "map.put() can be converted to assignment",
                             ProblemHighlightType.WEAK_WARNING,
-                            isOnTheFly,
                             ReplacePutWithAssignmentQuickfix()
                     )
-                    holder.registerProblem(problemDescriptor)
                 }
             }
         }
     }
-}
 
-class ReplacePutWithAssignmentQuickfix : LocalQuickFix {
+    companion object {
-    override fun getName() = "Convert put to assignment"
+        private val compatibleNames = setOf("put")
 
-    override fun getFamilyName() = name
+        fun isActiveFor(expression: KtDotQualifiedExpression): Boolean {
+            val callExpression = expression.callExpression
+            if (callExpression?.valueArguments?.size != 2) return false
 
-    override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+            val calleeExpression = callExpression.calleeExpression as? KtNameReferenceExpression ?: return false
-        val element = descriptor.psiElement as KtNameReferenceExpression
+            if (calleeExpression.getReferencedName() !in compatibleNames) return false
-        val valueArguments = (element.parent as? KtCallExpression)?.valueArguments ?: return
+
-        val qualifiedExpression = element.parent.parent as? KtDotQualifiedExpression ?: return
+            val context = expression.analyze()
-        qualifiedExpression.replace(KtPsiFactory(element).createExpressionByPattern("$0[$1] = $2",
+            if (expression.isUsedAsExpression(context)) return false
-                                                                                    qualifiedExpression.receiverExpression,
+            val resolvedCall = expression.getResolvedCall(context)
-                                                                                    valueArguments[0]?.getArgumentExpression() ?: return,
+            val receiverType = resolvedCall?.getExplicitReceiverValue()?.type ?: return false
-                                                                                    valueArguments[1]?.getArgumentExpression() ?: return))
+            val receiverClass = receiverType.constructor.declarationDescriptor as? ClassDescriptor ?: return false
+            return receiverClass.isSubclassOf(DefaultBuiltIns.Instance.mutableMap)
+        }
+
+        fun simplify(expression: KtDotQualifiedExpression) {
+            val valueArguments = expression.callExpression?.valueArguments ?: return
+            expression.replace(KtPsiFactory(expression).createExpressionByPattern("$0[$1] = $2",
+                                                                                  expression.receiverExpression,
+                                                                                  valueArguments[0]?.getArgumentExpression() ?: return,
+                                                                                  valueArguments[1]?.getArgumentExpression() ?: return))
+        }
+    }
+
+    class ReplacePutWithAssignmentQuickfix : LocalQuickFix {
+        override fun getName() = "Convert put to assignment"
+
+        override fun getFamilyName() = name
+
+        override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+            simplify(descriptor.psiElement.parent.parent as? KtDotQualifiedExpression ?: return)
+        }
     }
 }

idea/src/org/jetbrains/kotlin/idea/j2k/J2kPostProcessings.kt

@@ -25,10 +25,7 @@ import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.core.setVisibility
-import org.jetbrains.kotlin.idea.inspections.RedundantSamConstructorInspection
+import org.jetbrains.kotlin.idea.inspections.*
-import org.jetbrains.kotlin.idea.inspections.UnnecessaryVariableInspection
-import org.jetbrains.kotlin.idea.inspections.UseExpressionBodyInspection
-import org.jetbrains.kotlin.idea.inspections.findExistingEditor
 import org.jetbrains.kotlin.idea.intentions.*
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions.FoldIfToReturnAsymmetricallyIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions.FoldIfToReturnIntention
@@ -81,6 +78,7 @@ object J2KPostProcessingRegistrar {
         _processings.add(UnresolvedVariableReferenceFromInitializerToThisReferenceProcessing())
         _processings.add(RemoveRedundantSamAdaptersProcessing())
         _processings.add(RemoveRedundantCastToNullableProcessing())
+        _processings.add(ReplacePutWithAssignmentProcessing())
         _processings.add(UseExpressionBodyProcessing())
         _processings.add(UnnecessaryVariableProcessing())
 
@@ -274,6 +272,19 @@ object J2KPostProcessingRegistrar {
         }
     }
 
+    private class ReplacePutWithAssignmentProcessing : J2kPostProcessing {
+        override val writeActionNeeded = true
+
+        override fun createAction(element: KtElement, diagnostics: Diagnostics): (() -> Unit)? {
+            if (element !is KtDotQualifiedExpression) return null
+            if (!ReplacePutWithAssignmentInspection.isActiveFor(element)) return null
+
+            return {
+                ReplacePutWithAssignmentInspection.simplify(element)
+            }
+        }
+    }
+
     private class UseExpressionBodyProcessing : J2kPostProcessing {
         override val writeActionNeeded = true
 

idea/testData/inspectionsLocal/replacePutWithAssignment/nonMap.kt

@@ -6,5 +6,5 @@ class A {
 }
 
 fun foo() {
-    A().put<caret>(1, "foo")
+    A().<caret>put(1, "foo")
 }

idea/testData/inspectionsLocal/replacePutWithAssignment/putAsExpression.kt

@@ -3,4 +3,4 @@
 
 val map = mutableMapOf(42 to "foo")
 
-fun foo() = map.put<caret>(60, "bar")
+fun foo() = map.<caret>put(60, "bar")

idea/testData/inspectionsLocal/replacePutWithAssignment/putOnParameter.kt

@@ -1,5 +1,5 @@
 // WITH_RUNTIME
 
 fun foo(map: MutableMap<Int, String>) {
-    map.put<caret>(42, "foo")
+    map.<caret>put(42, "foo")
 }

idea/testData/inspectionsLocal/replacePutWithAssignment/putOnThis.kt

@@ -2,6 +2,6 @@
 
 class MyMap() : HashMap<String, String>() {
     init {
-        this.put<caret>("foo", "bar")
+        this.<caret>put("foo", "bar")
     }
 }

idea/testData/inspectionsLocal/replacePutWithAssignment/putOnVal.kt

@@ -2,5 +2,5 @@
 
 fun foo() {
     val map = mutableMapOf(42 to "foo")
-    map.put<caret>(60, "bar")
+    map.<caret>put(60, "bar")
 }

idea/testData/inspectionsLocal/replacePutWithAssignment/putOnVar.kt

@@ -2,5 +2,5 @@
 
 fun foo() {
     var map = mutableMapOf(42 to "foo")
-    map.put<caret>(60, "bar")
+    map.<caret>put(60, "bar")
 }

j2k/testData/fileOrElement/typeParameters/needTypeArgs.kt

@@ -12,7 +12,7 @@ internal class A {
 
     fun <K, V> getMap2(k: K, v: V): Map<K, V> {
         val map = HashMap<K, V>()
-        map.put(k, v)
+        map[k] = v
         return map
     }
 }