From 0ab924a2984e6c6b6e4e59ba9fdda2ee1fa1e1b0 Mon Sep 17 00:00:00 2001
From: Alexey Sedunov <alexey.sedunov@jetbrains.com>
Date: Tue, 12 Dec 2017 15:35:24 +0300
Subject: [PATCH] Change Signature: Fix overrider search for suspend functions

 #KT-21288 Fixed
---
 .../declarationsSearch/overridersSearch.kt    | 10 ++---
 .../KotlinChangeSignatureData.kt              | 43 ++++++++++++-------
 .../suspendImpls/after/Common/Common.iml      | 19 ++++++++
 .../after/Common/src/test/test.kt             |  5 +++
 .../suspendImpls/after/JS/JS.iml              | 21 +++++++++
 .../suspendImpls/after/JS/src/test/test.kt    |  5 +++
 .../suspendImpls/after/JVM/JVM.iml            | 21 +++++++++
 .../suspendImpls/after/JVM/src/test/test.kt   |  5 +++
 .../suspendImpls/before/Common/Common.iml     | 19 ++++++++
 .../before/Common/src/test/test.kt            |  5 +++
 .../suspendImpls/before/JS/JS.iml             | 21 +++++++++
 .../suspendImpls/before/JS/src/test/test.kt   |  5 +++
 .../suspendImpls/before/JVM/JVM.iml           | 21 +++++++++
 .../suspendImpls/before/JVM/src/test/test.kt  |  5 +++
 .../KotlinMultiModuleChangeSignatureTest.kt   |  4 ++
 15 files changed, 189 insertions(+), 20 deletions(-)
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/Common/Common.iml
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/Common/src/test/test.kt
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JS/JS.iml
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JS/src/test/test.kt
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JVM/JVM.iml
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JVM/src/test/test.kt
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/Common/Common.iml
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/Common/src/test/test.kt
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JS/JS.iml
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JS/src/test/test.kt
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JVM/JVM.iml
 create mode 100644 idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JVM/src/test/test.kt

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/declarationsSearch/overridersSearch.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/declarationsSearch/overridersSearch.kt
index 31df71b7f5a..a53eeb04a72 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/declarationsSearch/overridersSearch.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/declarationsSearch/overridersSearch.kt
@@ -140,7 +140,7 @@ private fun forEachKotlinOverride(
     val baseDescriptors = runReadAction { members.mapNotNull { it.unsafeResolveToDescriptor() as? CallableMemberDescriptor }.filter { it.isOverridable } }
     if (baseDescriptors.isEmpty()) return true
 
-    HierarchySearchRequest(ktClass, scope.restrictToKotlinSources(), true).searchInheritors().forEach {
+    HierarchySearchRequest(ktClass, scope, true).searchInheritors().forEach {
         val inheritor = it.unwrapped as? KtClassOrObject ?: return@forEach
         val inheritorDescriptor = runReadAction { inheritor.unsafeResolveToDescriptor() as ClassDescriptor }
         val substitutor = getTypeSubstitutor(baseClassDescriptor.defaultType, inheritorDescriptor.defaultType) ?: return@forEach
@@ -159,15 +159,15 @@ private fun forEachKotlinOverride(
 
 fun KtNamedDeclaration.forEachOverridingElement(
         scope: SearchScope = runReadAction { useScope },
-        processor: (PsiElement) -> Boolean
+        processor: (PsiElement, PsiElement) -> Boolean
 ): Boolean {
     val ktClass = runReadAction { containingClassOrObject as? KtClass } ?: return true
 
-    toLightMethods().forEach {
-        if (!OverridingMethodsSearch.search(it, scope.excludeKotlinSources(), true).all(processor)) return false
+    toLightMethods().forEach { baseMethod ->
+        if (!OverridingMethodsSearch.search(baseMethod, scope.excludeKotlinSources(), true).all { processor(baseMethod, it) }) return false
     }
 
-    return forEachKotlinOverride(ktClass, listOf(this), scope) { _, overrider -> processor(overrider) }
+    return forEachKotlinOverride(ktClass, listOf(this), scope) { baseElement, overrider -> processor(baseElement, overrider) }
 }
 
 fun PsiMethod.forEachOverridingMethod(
diff --git a/idea/src/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinChangeSignatureData.kt b/idea/src/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinChangeSignatureData.kt
index 8a8ef55ff86..5d22325ea81 100644
--- a/idea/src/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinChangeSignatureData.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinChangeSignatureData.kt
@@ -17,13 +17,11 @@
 package org.jetbrains.kotlin.idea.refactoring.changeSignature
 
 import com.intellij.psi.PsiElement
-import com.intellij.psi.search.searches.OverridingMethodsSearch
+import com.intellij.psi.PsiMethod
 import com.intellij.refactoring.changeSignature.MethodDescriptor
 import com.intellij.refactoring.changeSignature.OverriderUsageInfo
 import com.intellij.usageView.UsageInfo
-import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.asJava.namedUnwrappedElement
-import org.jetbrains.kotlin.asJava.toLightMethods
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.unsafeResolveToDescriptor
@@ -31,9 +29,14 @@ import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.highlighter.markers.actualsForExpected
 import org.jetbrains.kotlin.idea.highlighter.markers.isExpectedOrExpectedClassMember
 import org.jetbrains.kotlin.idea.refactoring.changeSignature.usages.KotlinCallableDefinitionUsage
+import org.jetbrains.kotlin.idea.search.declarationsSearch.forEachOverridingElement
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.KtCallableDeclaration
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.utils.SmartList
 import java.util.*
 
 class KotlinChangeSignatureData(
@@ -112,18 +115,28 @@ class KotlinChangeSignatureData(
 
             if (primaryDeclaration !is KtCallableDeclaration) return@flatMapTo emptyList()
 
-            primaryDeclaration.toLightMethods().flatMap { baseMethod ->
-                OverridingMethodsSearch
-                        .search(baseMethod)
-                        .mapNotNullTo(HashSet<UsageInfo>()) { overridingMethod ->
-                            if (overridingMethod is KtLightMethod) {
-                                val overridingDeclaration = overridingMethod.namedUnwrappedElement as KtNamedDeclaration
-                                val overridingDescriptor = overridingDeclaration.unsafeResolveToDescriptor() as CallableDescriptor
-                                KotlinCallableDefinitionUsage<PsiElement>(overridingDeclaration, overridingDescriptor, primaryFunction, null)
-                            }
-                            else OverriderUsageInfo(overridingMethod, baseMethod, true, true, true)
-                        }
+            val results = SmartList<UsageInfo>()
+
+            primaryDeclaration.forEachOverridingElement { baseElement, overridingElement ->
+                val currentDeclaration = overridingElement.namedUnwrappedElement
+                results += when (currentDeclaration) {
+                    is KtDeclaration -> {
+                        val overridingDescriptor = currentDeclaration.unsafeResolveToDescriptor() as CallableDescriptor
+                        KotlinCallableDefinitionUsage(currentDeclaration, overridingDescriptor, primaryFunction, null)
+                    }
+
+                    is PsiMethod -> {
+                        val baseMethod = baseElement as? PsiMethod ?: return@forEachOverridingElement true
+                        OverriderUsageInfo(currentDeclaration, baseMethod, true, true, true)
+                    }
+
+                    else -> return@forEachOverridingElement true
+                }
+
+                true
             }
+
+            results
         }
     }
 
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/Common/Common.iml b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/Common/Common.iml
new file mode 100644
index 00000000000..0ecf9949f6c
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/Common/Common.iml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="JAVA_MODULE" version="4">
+  <component name="FacetManager">
+      <facet type="kotlin-language" name="Kotlin">
+        <configuration version="2" platform="Common (experimental) " useProjectSettings="false">
+            <compilerSettings/>
+            <compilerArguments/>
+        </configuration>
+      </facet>
+    </component>
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/Common/src/test/test.kt b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/Common/src/test/test.kt
new file mode 100644
index 00000000000..9e212761bab
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/Common/src/test/test.kt
@@ -0,0 +1,5 @@
+package test
+
+interface I {
+    suspend fun foo(s: String, n: Int): Int
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JS/JS.iml b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JS/JS.iml
new file mode 100644
index 00000000000..3de270f43fa
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JS/JS.iml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="JAVA_MODULE" version="4">
+  <component name="FacetManager">
+      <facet type="kotlin-language" name="Kotlin">
+        <configuration version="2" platform="JavaScript " useProjectSettings="false">
+            <implements>Common</implements>
+            <compilerSettings/>
+            <compilerArguments/>
+        </configuration>
+      </facet>
+    </component>
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+    <orderEntry type="module" module-name="Common" />
+  </component>
+</module>
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JS/src/test/test.kt b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JS/src/test/test.kt
new file mode 100644
index 00000000000..f230a1526ab
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JS/src/test/test.kt
@@ -0,0 +1,5 @@
+package test
+
+class C : I {
+    suspend override fun foo(s: String, n: Int) = s.length
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JVM/JVM.iml b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JVM/JVM.iml
new file mode 100644
index 00000000000..76d2b885834
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JVM/JVM.iml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="JAVA_MODULE" version="4">
+  <component name="FacetManager">
+      <facet type="kotlin-language" name="Kotlin">
+        <configuration version="2" platform="JVM 1.6" useProjectSettings="false">
+            <implements>Common</implements>
+            <compilerSettings/>
+            <compilerArguments/>
+        </configuration>
+      </facet>
+    </component>
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+    <orderEntry type="module" module-name="Common" />
+  </component>
+</module>
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JVM/src/test/test.kt b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JVM/src/test/test.kt
new file mode 100644
index 00000000000..860ed577749
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/after/JVM/src/test/test.kt
@@ -0,0 +1,5 @@
+package test
+
+class D : I {
+    suspend override fun foo(s: String, n: Int) = s.length
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/Common/Common.iml b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/Common/Common.iml
new file mode 100644
index 00000000000..0ecf9949f6c
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/Common/Common.iml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="JAVA_MODULE" version="4">
+  <component name="FacetManager">
+      <facet type="kotlin-language" name="Kotlin">
+        <configuration version="2" platform="Common (experimental) " useProjectSettings="false">
+            <compilerSettings/>
+            <compilerArguments/>
+        </configuration>
+      </facet>
+    </component>
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/Common/src/test/test.kt b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/Common/src/test/test.kt
new file mode 100644
index 00000000000..be3e51a1cde
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/Common/src/test/test.kt
@@ -0,0 +1,5 @@
+package test
+
+interface I {
+    suspend fun <caret>foo(s: String): Int
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JS/JS.iml b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JS/JS.iml
new file mode 100644
index 00000000000..3de270f43fa
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JS/JS.iml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="JAVA_MODULE" version="4">
+  <component name="FacetManager">
+      <facet type="kotlin-language" name="Kotlin">
+        <configuration version="2" platform="JavaScript " useProjectSettings="false">
+            <implements>Common</implements>
+            <compilerSettings/>
+            <compilerArguments/>
+        </configuration>
+      </facet>
+    </component>
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+    <orderEntry type="module" module-name="Common" />
+  </component>
+</module>
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JS/src/test/test.kt b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JS/src/test/test.kt
new file mode 100644
index 00000000000..ae4aaeb91ce
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JS/src/test/test.kt
@@ -0,0 +1,5 @@
+package test
+
+class C : I {
+    suspend override fun foo(s: String) = s.length
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JVM/JVM.iml b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JVM/JVM.iml
new file mode 100644
index 00000000000..76d2b885834
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JVM/JVM.iml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="JAVA_MODULE" version="4">
+  <component name="FacetManager">
+      <facet type="kotlin-language" name="Kotlin">
+        <configuration version="2" platform="JVM 1.6" useProjectSettings="false">
+            <implements>Common</implements>
+            <compilerSettings/>
+            <compilerArguments/>
+        </configuration>
+      </facet>
+    </component>
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+    <orderEntry type="module" module-name="Common" />
+  </component>
+</module>
\ No newline at end of file
diff --git a/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JVM/src/test/test.kt b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JVM/src/test/test.kt
new file mode 100644
index 00000000000..872bd272e05
--- /dev/null
+++ b/idea/testData/refactoring/changeSignatureMultiModule/suspendImpls/before/JVM/src/test/test.kt
@@ -0,0 +1,5 @@
+package test
+
+class D : I {
+    suspend override fun foo(s: String) = s.length
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinMultiModuleChangeSignatureTest.kt b/idea/tests/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinMultiModuleChangeSignatureTest.kt
index f36cd2e19d8..c03c9372c8a 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinMultiModuleChangeSignatureTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinMultiModuleChangeSignatureTest.kt
@@ -101,4 +101,8 @@ class KotlinMultiModuleChangeSignatureTest : KotlinMultiFileTestCase() {
     fun testImplSecondaryConstructor() = doTest("JS/src/test/test.kt") {
         addParameter("b", KotlinChangeSignatureTest.BUILT_INS.booleanType, "false")
     }
+
+    fun testSuspendImpls() = doTest("Common/src/test/test.kt") {
+        addParameter("n", KotlinChangeSignatureTest.BUILT_INS.intType, "0")
+    }
 }