Refactorings: Support header/impl functions in member hierarchy transformations

#KT-18904 Fixed
 #KT-18905 Fixed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/usagesSearch/utils.kt

@@ -22,6 +22,8 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiReference
 import com.intellij.psi.search.SearchScope
+import com.intellij.psi.search.searches.MethodReferencesSearch
+import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.psi.util.MethodSignatureUtil
 import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.asJava.elements.KtLightMethod
@@ -250,3 +252,13 @@ fun PsiReference.isCallableOverrideUsage(declaration: KtNamedDeclaration): Boole
         }
     }
 }
+
+fun PsiElement.searchReferencesOrMethodReferences(): Collection<PsiReference> {
+    val lightMethods = toLightMethods()
+    return if (lightMethods.isNotEmpty()) {
+        lightMethods.flatMapTo(LinkedHashSet()) { MethodReferencesSearch.search(it) }
+    }
+    else {
+        ReferencesSearch.search(this).findAll()
+    }
+}

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertFunctionTypeParameterToReceiverIntention.kt

@@ -23,11 +23,9 @@ import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiNamedElement
 import com.intellij.psi.search.LocalSearchScope
-import com.intellij.psi.search.searches.MethodReferencesSearch
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.refactoring.util.RefactoringUIUtil
 import com.intellij.util.containers.MultiMap
-import org.jetbrains.kotlin.asJava.toLightMethods
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
@@ -41,6 +39,7 @@ import org.jetbrains.kotlin.idea.refactoring.introduce.introduceVariable.KotlinI
 import org.jetbrains.kotlin.idea.references.KtReference
 import org.jetbrains.kotlin.idea.references.KtSimpleReference
 import org.jetbrains.kotlin.idea.runSynchronouslyWithProgress
+import org.jetbrains.kotlin.idea.search.usagesSearch.searchReferencesOrMethodReferences
 import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.*
@@ -196,8 +195,7 @@ class ConvertFunctionTypeParameterToReceiverIntention : SelfTargetingRangeIntent
                             conflicts.putValue(callable, "Can't modify $renderedCallable")
                         }
 
-                        val references = callable.toLightMethods().flatMapTo(LinkedHashSet()) { MethodReferencesSearch.search(it) }
+                        usageLoop@ for (ref in callable.searchReferencesOrMethodReferences()) {
-                        usageLoop@ for (ref in references) {
                             val refElement = ref.element ?: continue
                             when (ref) {
                                 is KtSimpleReference<*> -> processExternalUsage(conflicts, refElement, usages)

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertFunctionTypeReceiverToParameterIntention.kt

@@ -21,11 +21,9 @@ import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.LocalSearchScope
-import com.intellij.psi.search.searches.MethodReferencesSearch
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.refactoring.util.RefactoringUIUtil
 import com.intellij.util.containers.MultiMap
-import org.jetbrains.kotlin.asJava.toLightMethods
 import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
@@ -38,6 +36,7 @@ import org.jetbrains.kotlin.idea.refactoring.checkConflictsInteractively
 import org.jetbrains.kotlin.idea.refactoring.getAffectedCallables
 import org.jetbrains.kotlin.idea.references.KtSimpleReference
 import org.jetbrains.kotlin.idea.runSynchronouslyWithProgress
+import org.jetbrains.kotlin.idea.search.usagesSearch.searchReferencesOrMethodReferences
 import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.*
@@ -135,8 +134,7 @@ class ConvertFunctionTypeReceiverToParameterIntention : SelfTargetingRangeIntent
                             conflicts.putValue(callable, "Can't modify $renderedCallable")
                         }
 
-                        val references = callable.toLightMethods().flatMapTo(LinkedHashSet()) { MethodReferencesSearch.search(it) }
+                        for (ref in callable.searchReferencesOrMethodReferences()) {
-                        for (ref in references) {
                             if (ref !is KtSimpleReference<*>) continue
                             processExternalUsage(ref, usages)
                         }

idea/src/org/jetbrains/kotlin/idea/refactoring/CallableRefactoring.kt

@@ -35,8 +35,12 @@ import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.core.getDeepestSuperDeclarations
 import org.jetbrains.kotlin.idea.core.getDirectlyOverriddenDeclarations
+import org.jetbrains.kotlin.idea.highlighter.markers.headerImplementations
+import org.jetbrains.kotlin.idea.highlighter.markers.isHeaderOrHeaderClassMember
+import org.jetbrains.kotlin.idea.highlighter.markers.liftToHeader
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.renderer.DescriptorRenderer
@@ -48,10 +52,13 @@ import java.util.*
 
 abstract class CallableRefactoring<out T: CallableDescriptor>(
         val project: Project,
-        val callableDescriptor: T,
+        callableDescriptor: T,
         val commandName: String) {
     private val LOG = Logger.getInstance(CallableRefactoring::class.java)
 
+    @Suppress("UNCHECKED_CAST")
+    val callableDescriptor = callableDescriptor.liftToHeader() as? T ?: callableDescriptor
+
     private val kind = (callableDescriptor as? CallableMemberDescriptor)?.kind ?: CallableMemberDescriptor.Kind.DECLARATION
 
     protected open fun forcePerformForSelectedFunctionOnly(): Boolean {
@@ -69,7 +76,7 @@ abstract class CallableRefactoring<out T: CallableDescriptor>(
             else -> {
                 throw IllegalStateException("Unexpected callable kind: $kind")
             }
-        }
+        }.map { it.liftToHeader() as? CallableDescriptor ?: it }
     }
 
     private fun showSuperFunctionWarningDialog(superCallables: Collection<CallableDescriptor>,
@@ -176,9 +183,16 @@ abstract class CallableRefactoring<out T: CallableDescriptor>(
 
 fun getAffectedCallables(project: Project, descriptorsForChange: Collection<CallableDescriptor>): List<PsiElement> {
     val baseCallables = descriptorsForChange.mapNotNull { DescriptorToSourceUtilsIde.getAnyDeclaration(project, it) }
-    return baseCallables + baseCallables.flatMap { it.toLightMethods() }.flatMapTo(HashSet<PsiElement>()) { psiMethod ->
+    return baseCallables + baseCallables.flatMapTo(HashSet<PsiElement>()) { callable ->
-        val overrides = OverridingMethodsSearch.search(psiMethod).findAll()
+        if (callable is KtDeclaration && callable.isHeaderOrHeaderClassMember()) {
-        overrides.map { method -> method.namedUnwrappedElement ?: method}
+            callable.headerImplementations()
+        }
+        else {
+            callable.toLightMethods().flatMap { psiMethod ->
+                val overrides = OverridingMethodsSearch.search(psiMethod).findAll()
+                overrides.map { method -> method.namedUnwrappedElement ?: method}
+            }
+        }
     }
 }
 

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByHeader/header/header.kt

@@ -0,0 +1,7 @@
+// "Convert '(Int) -> Int' to 'Int.() -> Int'" "true"
+
+header fun foo(n: Int, action: (<caret>Int) -> Int): Int
+
+fun test() {
+    foo(1) { n -> n + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByHeader/header/header.kt.after

@@ -0,0 +1,7 @@
+// "Convert '(Int) -> Int' to 'Int.() -> Int'" "true"
+
+header fun foo(n: Int, action: Int.() -> Int): Int
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByHeader/js/js.kt

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: (Int) -> Int) = action(n)
+
+fun test() {
+    foo(1) { n -> n + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByHeader/js/js.kt.after

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: Int.() -> Int) = n.action()
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByHeader/jvm/jvm.kt

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: (Int) -> Int) = action(n)
+
+fun test() {
+    foo(1) { n -> n + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByHeader/jvm/jvm.kt.after

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: Int.() -> Int) = n.action()
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByImpl/header/header.kt

@@ -0,0 +1,5 @@
+header fun foo(n: Int, action: (Int) -> Int): Int
+
+fun test() {
+    foo(1) { n -> n + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByImpl/header/header.kt.after

@@ -0,0 +1,5 @@
+header fun foo(n: Int, action: Int.() -> Int): Int
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByImpl/js/js.kt

@@ -0,0 +1,7 @@
+// "Convert '(Int) -> Int' to 'Int.() -> Int'" "true"
+
+impl fun foo(n: Int, action: (<caret>Int) -> Int) = action(n)
+
+fun test() {
+    foo(1) { n -> n + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByImpl/js/js.kt.after

@@ -0,0 +1,7 @@
+// "Convert '(Int) -> Int' to 'Int.() -> Int'" "true"
+
+impl fun foo(n: Int, action: Int.() -> Int) = n.action()
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByImpl/jvm/jvm.kt

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: (Int) -> Int) = action(n)
+
+fun test() {
+    foo(1) { n -> n + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeParameterToReceiverByImpl/jvm/jvm.kt.after

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: Int.() -> Int) = n.action()
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByHeader/header/header.kt

@@ -0,0 +1,7 @@
+// "Convert 'Int.() -> Int' to '(Int) -> Int'" "true"
+
+header fun foo(n: Int, action: <caret>Int.() -> Int): Int
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByHeader/header/header.kt.after

@@ -0,0 +1,7 @@
+// "Convert 'Int.() -> Int' to '(Int) -> Int'" "true"
+
+header fun foo(n: Int, action: (Int) -> Int): Int
+
+fun test() {
+    foo(1) { i -> i + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByHeader/js/js.kt

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: Int.() -> Int) = n.action()
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByHeader/js/js.kt.after

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: (Int) -> Int) = action(n)
+
+fun test() {
+    foo(1) { i -> i + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByHeader/jvm/jvm.kt

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: Int.() -> Int) = n.action()
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByHeader/jvm/jvm.kt.after

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: (Int) -> Int) = action(n)
+
+fun test() {
+    foo(1) { i -> i + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByImpl/header/header.kt

@@ -0,0 +1,5 @@
+header fun foo(n: Int, action: Int.() -> Int): Int
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByImpl/header/header.kt.after

@@ -0,0 +1,5 @@
+header fun foo(n: Int, action: (Int) -> Int): Int
+
+fun test() {
+    foo(1) { i -> i + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByImpl/js/js.kt

@@ -0,0 +1,7 @@
+// "Convert 'Int.() -> Int' to '(Int) -> Int'" "true"
+
+impl fun foo(n: Int, action: <caret>Int.() -> Int) = n.action()
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByImpl/js/js.kt.after

@@ -0,0 +1,7 @@
+// "Convert 'Int.() -> Int' to '(Int) -> Int'" "true"
+
+impl fun foo(n: Int, action: (Int) -> Int) = action(n)
+
+fun test() {
+    foo(1) { i -> i + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByImpl/jvm/jvm.kt

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: Int.() -> Int) = n.action()
+
+fun test() {
+    foo(1) { this + 1 }
+}

idea/testData/multiModuleQuickFix/functionTypeReceiverToParameterByImpl/jvm/jvm.kt.after

@@ -0,0 +1,5 @@
+impl fun foo(n: Int, action: (Int) -> Int) = action(n)
+
+fun test() {
+    foo(1) { i -> i + 1 }
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiModuleTest.kt

@@ -141,4 +141,16 @@ class QuickFixMultiModuleTest : AbstractQuickFixMultiModuleTest() {
 
     @Test
     fun testMemberFunParameterToReceiverByImpl() = doTestHeaderWithJvmAndJs()
+
+    @Test
+    fun testFunctionTypeParameterToReceiverByHeader() = doTestHeaderWithJvmAndJs()
+
+    @Test
+    fun testFunctionTypeParameterToReceiverByImpl() = doTestHeaderWithJvmAndJs()
+
+    @Test
+    fun testFunctionTypeReceiverToParameterByHeader() = doTestHeaderWithJvmAndJs()
+
+    @Test
+    fun testFunctionTypeReceiverToParameterByImpl() = doTestHeaderWithJvmAndJs()
 }