Improve "ReplaceWith" for class with generic parameter

#KT-27089 Fixed

idea/src/org/jetbrains/kotlin/idea/codeInliner/ClassUsageReplacementStrategy.kt

@@ -17,10 +17,15 @@
 package org.jetbrains.kotlin.idea.codeInliner
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
 import org.jetbrains.kotlin.idea.core.ShortenReferences
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.idea.stubindex.KotlinClassShortNameIndex
+import org.jetbrains.kotlin.idea.util.replaceOrCreateTypeArgumentList
+import org.jetbrains.kotlin.ir.expressions.typeParametersCount
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
@@ -87,15 +92,29 @@ class ClassUsageReplacementStrategy(
 
     private fun replaceConstructorCallWithOtherTypeConstruction(callExpression: KtCallExpression): KtElement {
         val referenceExpression = typeReplacement?.referenceExpression ?: error("Couldn't find referenceExpression")
-        callExpression.calleeExpression?.replace(referenceExpression)
+        val classFromReplacement = KotlinClassShortNameIndex
+            .getInstance()[referenceExpression.text, callExpression.project, callExpression.resolveScope]
+            .firstOrNull()
+
+        val replacementTypeArgumentList = typeReplacement.typeArgumentList
+        val replacementTypeArgumentCount = classFromReplacement?.typeParameters?.size
+            ?: replacementTypeArgumentList?.arguments?.size
 
         val typeArgumentList = callExpression.typeArgumentList
-        if (typeArgumentList != null && typeReplacement.typeArguments.size != typeArgumentList.arguments.size) {
+        val callDescriptor = callExpression.resolveToCall()
-            val newTypeArgumentList = typeReplacement.typeArgumentList
+            ?.resultingDescriptor
-            if (newTypeArgumentList != null) typeArgumentList.replace(newTypeArgumentList.copy())
+            ?.let { if (it is TypeAliasConstructorDescriptor) it.underlyingConstructorDescriptor else it }
-            else typeArgumentList.delete()
+        val typeArgumentCount = callDescriptor?.typeParametersCount ?: typeArgumentList?.arguments?.size
+
+        if (typeArgumentCount != replacementTypeArgumentCount) {
+            if (replacementTypeArgumentList == null) typeArgumentList?.delete()
+            else callExpression.replaceOrCreateTypeArgumentList(
+                replacementTypeArgumentList.copy() as KtTypeArgumentList
+            )
         }
 
+        callExpression.calleeExpression?.replace(referenceExpression)
+
         val expressionToReplace = callExpression.getQualifiedExpressionForSelectorOrThis()
         val newExpression = if (typeReplacementQualifierAsExpression != null)
             factory.createExpressionByPattern("$0.$1", typeReplacementQualifierAsExpression, callExpression)

idea/src/org/jetbrains/kotlin/idea/quickfix/replaceWith/DeprecatedSymbolUsageFix.kt

@@ -61,7 +61,7 @@ class DeprecatedSymbolUsageFix(
             if (targetDescriptors.isEmpty()) return false
 
             return targetDescriptors.all {
-                fetchReplaceWithPattern(it, import.project) != null
+                fetchReplaceWithPattern(it, import.project, null) != null
             }
         }
     }

idea/src/org/jetbrains/kotlin/idea/quickfix/replaceWith/DeprecatedSymbolUsageFixBase.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.DiagnosticFactory
 import org.jetbrains.kotlin.diagnostics.Errors
+import org.jetbrains.kotlin.idea.caches.KotlinShortNamesCache
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.codeInliner.CallableUsageReplacementStrategy
 import org.jetbrains.kotlin.idea.codeInliner.ClassUsageReplacementStrategy
@@ -37,9 +38,12 @@ import org.jetbrains.kotlin.idea.quickfix.KotlinQuickFixAction
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.references.resolveMainReferenceToDescriptors
 import org.jetbrains.kotlin.idea.search.restrictToKotlinSources
+import org.jetbrains.kotlin.idea.util.replaceOrCreateTypeArgumentList
+import org.jetbrains.kotlin.ir.expressions.typeParametersCount
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.isCallee
+import org.jetbrains.kotlin.psi.psiUtil.referenceExpression
 import org.jetbrains.kotlin.renderer.DescriptorRenderer
 import org.jetbrains.kotlin.resolve.annotations.argumentValue
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
@@ -76,7 +80,11 @@ abstract class DeprecatedSymbolUsageFixBase(
     protected abstract fun invoke(replacementStrategy: UsageReplacementStrategy, project: Project, editor: Editor?)
 
     companion object {
-        fun fetchReplaceWithPattern(descriptor: DeclarationDescriptor, project: Project): ReplaceWith? {
+        fun fetchReplaceWithPattern(
+            descriptor: DeclarationDescriptor,
+            project: Project,
+            contextElement: KtSimpleNameExpression?
+        ): ReplaceWith? {
             val annotation = descriptor.annotations.findAnnotation(KotlinBuiltIns.FQ_NAMES.deprecated) ?: return null
             val replaceWithValue =
                 annotation.argumentValue(kotlin.Deprecated::replaceWith.name)?.safeAs<AnnotationValue>()?.value ?: return null
@@ -92,7 +100,35 @@ abstract class DeprecatedSymbolUsageFixBase(
                 }
             ) return null
 
-            return ReplaceWith(pattern, imports)
+            return ReplaceWith(pattern.applyContextElement(contextElement, descriptor), imports)
+        }
+
+        private fun String.applyContextElement(
+            element: KtSimpleNameExpression?,
+            descriptor: DeclarationDescriptor
+        ): String {
+            if (element == null) return this
+            val expressionFromPattern = KtPsiFactory(element).createExpressionIfPossible(this) as? KtCallExpression ?: return this
+            val methodFromPattern = expressionFromPattern.referenceExpression()?.let { name ->
+                KotlinShortNamesCache(element.project).getMethodsByName(
+                    name.text,
+                    element.resolveScope
+                ).firstOrNull()
+            }
+
+            val patternTypeArgumentList = expressionFromPattern.typeArgumentList
+            val patternTypeArgumentCount = methodFromPattern?.typeParameterList?.typeParameters?.size
+                ?: patternTypeArgumentList?.arguments?.size
+                ?: return this
+
+            val typeArgumentList = (element.parent as? KtCallExpression)?.typeArgumentList
+            return if (descriptor is CallableDescriptor && descriptor.typeParametersCount == patternTypeArgumentCount ||
+                patternTypeArgumentCount == typeArgumentList?.arguments?.size
+            ) {
+                if (typeArgumentList != null) expressionFromPattern.replaceOrCreateTypeArgumentList(typeArgumentList.copy() as KtTypeArgumentList)
+                else patternTypeArgumentList?.delete()
+                expressionFromPattern.text
+            } else this
         }
 
         data class Data(
@@ -120,7 +156,7 @@ abstract class DeprecatedSymbolUsageFixBase(
                 else -> throw IllegalStateException("Bad QuickFixRegistrar configuration")
             }
 
-            val replacement = fetchReplaceWithPattern(descriptor, nameExpression.project) ?: return null
+            val replacement = fetchReplaceWithPattern(descriptor, nameExpression.project, nameExpression) ?: return null
             return Data(nameExpression, replacement, descriptor)
         }
 
@@ -135,10 +171,10 @@ abstract class DeprecatedSymbolUsageFixBase(
             val bindingContext = resolutionFacade.analyze(element, BodyResolveMode.PARTIAL)
             var target = element.mainReference.resolveToDescriptors(bindingContext).singleOrNull() ?: return null
 
-            var replacePatternFromSymbol = fetchReplaceWithPattern(target, resolutionFacade.project)
+            var replacePatternFromSymbol = fetchReplaceWithPattern(target, resolutionFacade.project, element)
             if (replacePatternFromSymbol == null && target is ConstructorDescriptor) {
                 target = target.containingDeclaration
-                replacePatternFromSymbol = fetchReplaceWithPattern(target, resolutionFacade.project)
+                replacePatternFromSymbol = fetchReplaceWithPattern(target, resolutionFacade.project, element)
             }
 
             // check that ReplaceWith hasn't changed
@@ -165,7 +201,7 @@ abstract class DeprecatedSymbolUsageFixBase(
                                     ?.getStrictParentOfType<KtTypeAlias>()
                                 if (typeAlias != null) {
                                     val usedConstructorWithOwnReplaceWith = usedConstructorsWithOwnReplaceWith(
-                                        element.project, target, typeAlias
+                                        element.project, target, typeAlias, element
                                     )
 
                                     if (usedConstructorWithOwnReplaceWith != null) {
@@ -204,10 +240,13 @@ abstract class DeprecatedSymbolUsageFixBase(
         }
 
         private fun usedConstructorsWithOwnReplaceWith(
-            project: Project, classifier: ClassifierDescriptorWithTypeParameters, typeAlias: PsiElement
+            project: Project,
+            classifier: ClassifierDescriptorWithTypeParameters,
+            typeAlias: PsiElement,
+            contextElement: KtSimpleNameExpression
         ): ConstructorDescriptor? {
             val specialReplaceWithForConstructor = classifier.constructors.filter {
-                fetchReplaceWithPattern(it, project) != null
+                fetchReplaceWithPattern(it, project, contextElement) != null
             }.toSet()
 
             if (specialReplaceWithForConstructor.isEmpty()) {

idea/src/org/jetbrains/kotlin/idea/util/expressionExt.kt

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.util
+
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtTypeArgumentList
+
+
+fun KtCallExpression.replaceOrCreateTypeArgumentList(newTypeArgumentList: KtTypeArgumentList) {
+    if (typeArgumentList != null) typeArgumentList?.replace(newTypeArgumentList)
+    else addAfter(
+        newTypeArgumentList,
+        calleeExpression
+    )
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument.kt

@@ -1,4 +1,4 @@
-// "Replace with 'NewClass<T>'" "true"
+// "Replace with 'NewClass<Int>'" "true"
 package ppp
 
 @Deprecated("renamed", ReplaceWith("NewClass<T>"))

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument.kt.after

@@ -1,4 +1,4 @@
-// "Replace with 'NewClass<T>'" "true"
+// "Replace with 'NewClass<Int>'" "true"
 package ppp
 
 @Deprecated("renamed", ReplaceWith("NewClass<T>"))

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument4.kt

@@ -0,0 +1,9 @@
+// "Replace with 'Factory<Int>()'" "true"
+// WITH_RUNTIME
+
+class Foo<T> @Deprecated("", ReplaceWith("Factory<T>()")) constructor()
+fun <T> Factory(): Foo<T> = TODO()
+
+fun baz() {
+    val foo = <caret>Foo<Int>()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument4.kt.after

@@ -0,0 +1,9 @@
+// "Replace with 'Factory<Int>()'" "true"
+// WITH_RUNTIME
+
+class Foo<T> @Deprecated("", ReplaceWith("Factory<T>()")) constructor()
+fun <T> Factory(): Foo<T> = TODO()
+
+fun baz() {
+    val foo = <caret>Factory<Int>()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument5.kt

@@ -0,0 +1,9 @@
+// "Replace with 'Factory<Int>()'" "true"
+// WITH_RUNTIME
+
+class Foo<T> @Deprecated("", ReplaceWith("Factory()")) constructor()
+fun <T> Factory(): Foo<T> = TODO()
+
+fun baz() {
+    val foo = <caret>Foo<Int>()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument5.kt.after

@@ -0,0 +1,9 @@
+// "Replace with 'Factory<Int>()'" "true"
+// WITH_RUNTIME
+
+class Foo<T> @Deprecated("", ReplaceWith("Factory()")) constructor()
+fun <T> Factory(): Foo<T> = TODO()
+
+fun baz() {
+    val foo = <caret>Factory<Int>()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument6.kt

@@ -0,0 +1,15 @@
+// "Replace with 'test.New'" "true"
+// WITH_RUNTIME
+
+package test
+
+abstract class Main<T>
+
+@Deprecated("", ReplaceWith("test.New"))
+class Old<T> : Main<T>()
+
+class New<T> : Main<T>()
+
+fun test() {
+    val main = <caret>Old<Int>()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument6.kt.after

@@ -0,0 +1,15 @@
+// "Replace with 'test.New'" "true"
+// WITH_RUNTIME
+
+package test
+
+abstract class Main<T>
+
+@Deprecated("", ReplaceWith("test.New"))
+class Old<T> : Main<T>()
+
+class New<T> : Main<T>()
+
+fun test() {
+    val main = <caret>New<Int>()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument7.kt

@@ -0,0 +1,14 @@
+// "Replace with 'New'" "true"
+// WITH_RUNTIME
+// ERROR: Type inference failed: Not enough information to infer parameter T in constructor New<T>()<br>Please specify it explicitly.<br>
+
+abstract class Main<T>
+
+@Deprecated("", ReplaceWith("New"))
+class Old<T, F> : Main<T>()
+
+class New<T> : Main<T>()
+
+fun test() {
+    val main = <caret>Old<Int, String>()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument7.kt.after

@@ -0,0 +1,14 @@
+// "Replace with 'New'" "true"
+// WITH_RUNTIME
+// ERROR: Type inference failed: Not enough information to infer parameter T in constructor New<T>()<br>Please specify it explicitly.<br>
+
+abstract class Main<T>
+
+@Deprecated("", ReplaceWith("New"))
+class Old<T, F> : Main<T>()
+
+class New<T> : Main<T>()
+
+fun test() {
+    val main = <caret>New()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgumentWithoutSpecifyType.kt

@@ -0,0 +1,9 @@
+// "Replace with 'Factory()'" "true"
+// WITH_RUNTIME
+
+class Foo<T> @Deprecated("", ReplaceWith("Factory<T>()")) constructor()
+fun <T> Factory(): Foo<T> = TODO()
+
+fun baz() {
+    val foo: Foo<Int> = <caret>Foo()
+}

idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgumentWithoutSpecifyType.kt.after

@@ -0,0 +1,9 @@
+// "Replace with 'Factory()'" "true"
+// WITH_RUNTIME
+
+class Foo<T> @Deprecated("", ReplaceWith("Factory<T>()")) constructor()
+fun <T> Factory(): Foo<T> = TODO()
+
+fun baz() {
+    val foo: Foo<Int> = <caret>Factory()
+}

idea/testData/quickfix/deprecatedSymbolUsage/typeArguments/emptyVarargRuntime.kt

@@ -1,4 +1,4 @@
-// "Replace with 'newFun(*elements)'" "true"
+// "Replace with 'newFun<String>(*elements)'" "true"
 // WITH_RUNTIME
 
 @Deprecated("", ReplaceWith("newFun(*elements)"))

idea/testData/quickfix/deprecatedSymbolUsage/typeArguments/emptyVarargRuntime.kt.after

@@ -1,4 +1,4 @@
-// "Replace with 'newFun(*elements)'" "true"
+// "Replace with 'newFun<String>(*elements)'" "true"
 // WITH_RUNTIME
 
 @Deprecated("", ReplaceWith("newFun(*elements)"))

idea/testData/quickfix/deprecatedSymbolUsage/typeArguments/explicitInPatternImplicitInUsage.kt

@@ -1,4 +1,4 @@
-// "Replace with 'newFun<T>()'" "true"
+// "Replace with 'newFun()'" "true"
 
 @Deprecated("", ReplaceWith("newFun<T>()"))
 fun <T : Any> oldFun(): T? {

idea/testData/quickfix/deprecatedSymbolUsage/typeArguments/explicitInPatternImplicitInUsage.kt.after

@@ -1,4 +1,4 @@
-// "Replace with 'newFun<T>()'" "true"
+// "Replace with 'newFun()'" "true"
 
 @Deprecated("", ReplaceWith("newFun<T>()"))
 fun <T : Any> oldFun(): T? {

idea/testData/quickfix/deprecatedSymbolUsage/typeArguments/explicitTypeArg.kt

@@ -1,4 +1,4 @@
-// "Replace with 'newFun<T>()'" "true"
+// "Replace with 'newFun<String>()'" "true"
 
 @Deprecated("", ReplaceWith("newFun<T>()"))
 fun <T> oldFun() {

idea/testData/quickfix/deprecatedSymbolUsage/typeArguments/explicitTypeArg.kt.after

@@ -1,4 +1,4 @@
-// "Replace with 'newFun<T>()'" "true"
+// "Replace with 'newFun<String>()'" "true"
 
 @Deprecated("", ReplaceWith("newFun<T>()"))
 fun <T> oldFun() {

idea/testData/quickfix/deprecatedSymbolUsage/typeArguments/keepOriginalIfQualified.kt

@@ -1,4 +1,4 @@
-// "Replace with 'newFun<T>()'" "true"
+// "Replace with 'newFun<kotlin.Int>()'" "true"
 
 @Deprecated("", ReplaceWith("newFun<T>()"))
 fun <T> oldFun() {

idea/testData/quickfix/deprecatedSymbolUsage/typeArguments/keepOriginalIfQualified.kt.after

@@ -1,4 +1,4 @@
-// "Replace with 'newFun<T>()'" "true"
+// "Replace with 'newFun<kotlin.Int>()'" "true"
 
 @Deprecated("", ReplaceWith("newFun<T>()"))
 fun <T> oldFun() {
@@ -8,5 +8,5 @@ fun <T> oldFun() {
 fun <T> newFun(){}
 
 fun foo() {
-    <caret>newFun<kotlin.Int>()
+    <caret>newFun<Int>()
 }

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -6056,6 +6056,31 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                 runTest("idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument3.kt");
             }
 
+            @TestMetadata("constructorUsageWithTypeArgument4.kt")
+            public void testConstructorUsageWithTypeArgument4() throws Exception {
+                runTest("idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument4.kt");
+            }
+
+            @TestMetadata("constructorUsageWithTypeArgument5.kt")
+            public void testConstructorUsageWithTypeArgument5() throws Exception {
+                runTest("idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument5.kt");
+            }
+
+            @TestMetadata("constructorUsageWithTypeArgument6.kt")
+            public void testConstructorUsageWithTypeArgument6() throws Exception {
+                runTest("idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument6.kt");
+            }
+
+            @TestMetadata("constructorUsageWithTypeArgument7.kt")
+            public void testConstructorUsageWithTypeArgument7() throws Exception {
+                runTest("idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgument7.kt");
+            }
+
+            @TestMetadata("constructorUsageWithTypeArgumentWithoutSpecifyType.kt")
+            public void testConstructorUsageWithTypeArgumentWithoutSpecifyType() throws Exception {
+                runTest("idea/testData/quickfix/deprecatedSymbolUsage/classUsages/constructorUsageWithTypeArgumentWithoutSpecifyType.kt");
+            }
+
             @TestMetadata("imports.kt")
             public void testImports() throws Exception {
                 runTest("idea/testData/quickfix/deprecatedSymbolUsage/classUsages/imports.kt");