Extract Function: Add support for member extension properties

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/resolvedCallUtil.kt

@@ -84,3 +84,5 @@ public fun ResolvedCall<*>.getImplicitReceiverValue(): ReceiverValue {
         else -> ReceiverValue.NO_RECEIVER
     }
 }
+
+public fun ResolvedCall<*>.hasBothReceivers() = dispatchReceiver.exists() && extensionReceiver.exists()

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/ExtractableCodeDescriptor.kt

@@ -22,6 +22,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.util.containers.ContainerUtil
 import com.intellij.util.containers.MultiMap
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.idea.core.KotlinNameSuggester
@@ -32,6 +33,7 @@ import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.Analysis
 import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.*
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference.ShorteningMode
 import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.idea.util.approximateFlexibleTypes
 import org.jetbrains.kotlin.idea.util.isAnnotatedNotNull
 import org.jetbrains.kotlin.idea.util.isAnnotatedNullable
@@ -39,14 +41,16 @@ import org.jetbrains.kotlin.idea.util.psi.patternMatching.KotlinPsiRange
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.*
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElement
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.descriptorUtil.resolveTopLevelClass
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.typeUtil.builtIns
 import org.jetbrains.kotlin.types.typeUtil.isUnit
-import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import java.util.*
 
 interface Parameter {
@@ -95,15 +99,27 @@ class RenameReplacement(override val parameter: Parameter): ParameterReplacement
     }
 }
 
-class WrapInWithReplacement(override val parameter: Parameter): ParameterReplacement {
+abstract class WrapInWithReplacement : Replacement {
+    abstract val argumentText: String
+
     override fun invoke(descriptor: ExtractableCodeDescriptor, e: KtElement): KtElement {
-        val call = e.parents.firstIsInstance<KtCallExpression>().getQualifiedExpressionForSelectorOrThis()
+        val call = (e as? KtSimpleNameExpression)?.getQualifiedElement() ?: return e
-        val replacingExpression = KtPsiFactory(e).createExpressionByPattern("with($0) { $1 }", parameter.name, call)
+        val replacingExpression = KtPsiFactory(e).createExpressionByPattern("with($0) { $1 }", argumentText, call)
         val replace = call.replace(replacingExpression)
         return (replace as KtCallExpression).functionLiteralArguments.first().getFunctionLiteral().bodyExpression!!.statements.first()
     }
+}
 
-    override fun copy(parameter: Parameter) = WrapInWithReplacement(parameter)
+class WrapParameterInWithReplacement(override val parameter: Parameter): WrapInWithReplacement(), ParameterReplacement {
+    override val argumentText: String
+        get() = parameter.name
+
+    override fun copy(parameter: Parameter) = WrapParameterInWithReplacement(parameter)
+}
+
+class WrapCompanionInWithReplacement(val descriptor: ClassDescriptor): WrapInWithReplacement() {
+    override val argumentText: String
+        get() = IdeDescriptorRenderers.SOURCE_CODE.renderClassifierName(descriptor)
 }
 
 class AddPrefixReplacement(override val parameter: Parameter): ParameterReplacement {

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/ExtractionData.kt

@@ -44,6 +44,7 @@ import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
 import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getImplicitReceiverValue
+import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.hasBothReceivers
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.calls.tasks.isSynthesizedInvoke
 import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
@@ -75,7 +76,8 @@ data class ResolvedReferenceInfo(
         val refExpr: KtSimpleNameExpression,
         val resolveResult: ResolveResult,
         val smartCast: KotlinType?,
-        val possibleTypes: Set<KotlinType>
+        val possibleTypes: Set<KotlinType>,
+        val shouldSkipPrimaryReceiver: Boolean
 )
 
 internal var KtSimpleNameExpression.resolveResult: ResolveResult? by CopyableUserDataProperty(Key.create("RESOLVE_RESULT"))
@@ -202,6 +204,7 @@ data class ExtractionData(
 
             val smartCast: KotlinType?
             val possibleTypes: Set<KotlinType>
+            val shouldSkipPrimaryReceiver: Boolean
 
             // Qualified property reference: a.b
             val qualifiedExpression = newRef.getQualifiedExpressionForSelector()
@@ -211,13 +214,15 @@ data class ExtractionData(
                 possibleTypes = getPossibleTypes(smartCastTarget, originalResolveResult.resolvedCall, originalContext)
                 val receiverDescriptor =
                         (originalResolveResult.resolvedCall?.dispatchReceiver as? ImplicitReceiver)?.declarationDescriptor
-                if (smartCast == null
+                shouldSkipPrimaryReceiver = smartCast == null
-                    && !DescriptorUtils.isCompanionObject(receiverDescriptor)
+                                            && !DescriptorUtils.isCompanionObject(receiverDescriptor)
-                    && qualifiedExpression.receiverExpression !is KtSuperExpression) continue
+                                            && qualifiedExpression.receiverExpression !is KtSuperExpression
+                if (shouldSkipPrimaryReceiver && !(originalResolveResult.resolvedCall?.hasBothReceivers() ?: false)) continue
             }
             else {
                 smartCast = originalContext[BindingContext.SMARTCAST, originalResolveResult.originalRefExpr]
                 possibleTypes = getPossibleTypes(originalResolveResult.originalRefExpr, originalResolveResult.resolvedCall, originalContext)
+                shouldSkipPrimaryReceiver = false
             }
 
             val parent = newRef.parent
@@ -241,11 +246,11 @@ data class ExtractionData(
                     val functionResolveResult = originalResolveResult.copy(resolvedCall = originalFunctionCall!!,
                                                                            descriptor = originalFunctionCall.resultingDescriptor,
                                                                            declaration = invokeDeclaration)
-                    referencesInfo.add(ResolvedReferenceInfo(newRef, variableResolveResult, smartCast, possibleTypes))
+                    referencesInfo.add(ResolvedReferenceInfo(newRef, variableResolveResult, smartCast, possibleTypes, shouldSkipPrimaryReceiver))
-                    referencesInfo.add(ResolvedReferenceInfo(newRef, functionResolveResult, smartCast, possibleTypes))
+                    referencesInfo.add(ResolvedReferenceInfo(newRef, functionResolveResult, smartCast, possibleTypes, shouldSkipPrimaryReceiver))
                 }
                 else {
-                    referencesInfo.add(ResolvedReferenceInfo(newRef, originalResolveResult, smartCast, possibleTypes))
+                    referencesInfo.add(ResolvedReferenceInfo(newRef, originalResolveResult, smartCast, possibleTypes, shouldSkipPrimaryReceiver))
                 }
             }
         }

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/inferParameterInfo.kt

@@ -35,8 +35,10 @@ import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
 import org.jetbrains.kotlin.psi.psiUtil.isInsideOf
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.hasBothReceivers
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.calls.tasks.isSynthesizedInvoke
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
@@ -44,8 +46,8 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.getImportableDescriptor
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
-import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.resolve.scopes.utils.findFunction
 import org.jetbrains.kotlin.types.CommonSupertypes
 import org.jetbrains.kotlin.types.KotlinType
@@ -90,9 +92,22 @@ internal fun ExtractionData.inferParametersInfo(
             else -> extensionReceiver
         } as? ReceiverValue) ?: ReceiverValue.NO_RECEIVER
 
-        extractReceiver(receiverToExtract, info, targetScope, refInfo, extractedDescriptorToParameter, pseudocode, bindingContext, false)
+        val twoReceivers = resolvedCall != null && resolvedCall.hasBothReceivers()
-        if (options.canWrapInWith && resolvedCall != null && isMemberExtensionFunction(resolvedCall, ref)) {
+        val dispatchReceiverDescriptor = (resolvedCall?.dispatchReceiver as? ImplicitReceiver)?.declarationDescriptor
-            extractReceiver(resolvedCall.dispatchReceiver, info, targetScope, refInfo, extractedDescriptorToParameter, pseudocode, bindingContext, true)
+        if (twoReceivers
+            && resolvedCall!!.extensionReceiver is ExpressionReceiver
+            && DescriptorUtils.isCompanionObject(dispatchReceiverDescriptor)) {
+            info.replacementMap.putValue(refInfo.resolveResult.originalRefExpr,
+                                         WrapCompanionInWithReplacement(dispatchReceiverDescriptor as ClassDescriptor))
+            continue
+        }
+
+        if (!refInfo.shouldSkipPrimaryReceiver) {
+            extractReceiver(receiverToExtract, info, targetScope, refInfo, extractedDescriptorToParameter, pseudocode, bindingContext, false)
+        }
+
+        if (options.canWrapInWith && twoReceivers) {
+            extractReceiver(resolvedCall!!.dispatchReceiver, info, targetScope, refInfo, extractedDescriptorToParameter, pseudocode, bindingContext, true)
         }
     }
 
@@ -132,7 +147,7 @@ private fun ExtractionData.extractReceiver(
         extractedDescriptorToParameter: HashMap<DeclarationDescriptor, MutableParameter>,
         pseudocode: Pseudocode,
         bindingContext: BindingContext,
-        isMemberExtensionFunction: Boolean
+        isMemberExtension: Boolean
 ) {
     val (originalRef, originalDeclaration, originalDescriptor, resolvedCall) = refInfo.resolveResult
 
@@ -162,12 +177,17 @@ private fun ExtractionData.extractReceiver(
         } as? ClassifierDescriptor
     }
 
-    if (referencedClassifierDescriptor != null && !(isMemberExtensionFunction && options.canWrapInWith)) {
+    if (referencedClassifierDescriptor != null) {
         if (!referencedClassifierDescriptor.defaultType.processTypeIfExtractable(
                 info.typeParameters, info.nonDenotableTypes, options, targetScope, referencedClassifierDescriptor is TypeParameterDescriptor
         )) return
 
-        if (referencedClassifierDescriptor is ClassDescriptor) {
+        if (options.canWrapInWith
+            && resolvedCall != null
+            && resolvedCall.hasBothReceivers()
+            && DescriptorUtils.isCompanionObject(referencedClassifierDescriptor)) {
+            info.replacementMap.putValue(originalRef, WrapCompanionInWithReplacement(referencedClassifierDescriptor as ClassDescriptor))
+        } else if (referencedClassifierDescriptor is ClassDescriptor) {
             info.replacementMap.putValue(originalRef, FqNameReplacement(originalDescriptor.getImportableDescriptor().fqNameSafe))
         }
     }
@@ -255,7 +275,7 @@ private fun ExtractionData.extractReceiver(
             }
 
             val replacement = when {
-                isMemberExtensionFunction -> WrapInWithReplacement(parameter)
+                isMemberExtension -> WrapParameterInWithReplacement(parameter)
                 hasThisReceiver && extractThis -> AddPrefixReplacement(parameter)
                 else -> RenameReplacement(parameter)
             }
@@ -304,11 +324,3 @@ private fun suggestParameterType(
                else -> null
            } ?: builtIns.defaultParameterType
 }
-
-private fun isMemberExtensionFunction(resolvedCall: ResolvedCall<*>, ref: KtSimpleNameExpression): Boolean {
-    // TODO temporary hack because we couldn't correctly extract member extension function with two explicit receivers
-    if (ref.parent !is KtCallExpression) return false
-
-    val resultingDescriptor = resolvedCall.resultingDescriptor
-    return resultingDescriptor is FunctionDescriptor && resolvedCall.extensionReceiver != ReceiverValue.NO_RECEIVER && resolvedCall.dispatchReceiver != ReceiverValue.NO_RECEIVER
-}

idea/testData/debugger/tinyApp/outs/extensionMemberProperty.out

@@ -0,0 +1,53 @@
+LineBreakpoint created at extensionMemberProperty.kt:13
+LineBreakpoint created at extensionMemberProperty.kt:19
+LineBreakpoint created at extensionMemberProperty.kt:25
+LineBreakpoint created at extensionMemberProperty.kt:31
+LineBreakpoint created at extensionMemberProperty.kt:41
+LineBreakpoint created at extensionMemberProperty.kt:46
+LineBreakpoint created at extensionMemberProperty.kt:51
+LineBreakpoint created at extensionMemberProperty.kt:56
+LineBreakpoint created at extensionMemberProperty.kt:74
+LineBreakpoint created at extensionMemberProperty.kt:80
+LineBreakpoint created at extensionMemberProperty.kt:86
+LineBreakpoint created at extensionMemberProperty.kt:92
+LineBreakpoint created at extensionMemberProperty.kt:102
+LineBreakpoint created at extensionMemberProperty.kt:107
+LineBreakpoint created at extensionMemberProperty.kt:112
+LineBreakpoint created at extensionMemberProperty.kt:117
+!JDK_HOME!\bin\java -agentlib:jdwp=transport=dt_socket,address=!HOST_NAME!:!HOST_PORT!,suspend=y,server=n -Dfile.encoding=!FILE_ENCODING! -classpath !OUTPUT_PATH!;!KOTLIN_RUNTIME!;!CUSTOM_LIBRARY!;!RT_JAR! extensionMemberProperty.ExtensionMemberPropertyKt
+Connected to the target VM, address: '!HOST_NAME!:PORT_NAME!', transport: 'socket'
+extensionMemberProperty.kt:13
+Compile bytecode for extClass.testPublic
+extensionMemberProperty.kt:19
+Compile bytecode for testPublic
+extensionMemberProperty.kt:25
+Compile bytecode for extClass.testPrivate
+extensionMemberProperty.kt:31
+Compile bytecode for testPrivate
+extensionMemberProperty.kt:41
+Compile bytecode for testPublic
+extensionMemberProperty.kt:46
+Compile bytecode for this.testPublic
+extensionMemberProperty.kt:51
+Compile bytecode for testPrivate
+extensionMemberProperty.kt:56
+Compile bytecode for this.testPrivate
+extensionMemberProperty.kt:74
+Compile bytecode for extClass.testCompPublic
+extensionMemberProperty.kt:80
+Compile bytecode for testCompPublic
+extensionMemberProperty.kt:86
+Compile bytecode for extClass.testCompPrivate
+extensionMemberProperty.kt:92
+Compile bytecode for testCompPrivate
+extensionMemberProperty.kt:102
+Compile bytecode for testCompPublic
+extensionMemberProperty.kt:107
+Compile bytecode for this.testCompPublic
+extensionMemberProperty.kt:112
+Compile bytecode for testCompPrivate
+extensionMemberProperty.kt:117
+Compile bytecode for this.testCompPrivate
+Disconnected from the target VM, address: '!HOST_NAME!:PORT_NAME!', transport: 'socket'
+
+Process finished with exit code 0

idea/testData/debugger/tinyApp/src/evaluate/multipleBreakpoints/extensionMemberProperty.kt

@@ -0,0 +1,124 @@
+package extensionMemberProperty
+
+fun main(args: Array<String>) {
+    MemberClass().testMember(ExtClass())
+    MemberClass.testCompanion(ExtClass())
+}
+
+class MemberClass {
+    fun testMember(extClass: ExtClass) {
+        // EXPRESSION: extClass.testPublic
+        // RESULT: 1: I
+        //Breakpoint!
+        extClass.testPublic
+
+        with(extClass) {
+            // EXPRESSION: testPublic
+            // RESULT: 1: I
+            //Breakpoint!
+            testPublic
+        }
+
+        // EXPRESSION: extClass.testPrivate
+        // RESULT: 1: I
+        //Breakpoint!
+        extClass.testPrivate
+
+        with(extClass) {
+            // EXPRESSION: testPrivate
+            // RESULT: 1: I
+            //Breakpoint!
+            testPrivate
+        }
+
+        extClass.testExtMember()
+    }
+
+    fun ExtClass.testExtMember() {
+        // EXPRESSION: testPublic
+        // RESULT: 1: I
+        //Breakpoint!
+        testPublic
+
+        // EXPRESSION: this.testPublic
+        // RESULT: 1: I
+        //Breakpoint!
+        this.testPublic
+
+        // EXPRESSION: testPrivate
+        // RESULT: 1: I
+        //Breakpoint!
+        testPrivate
+
+        // EXPRESSION: this.testPrivate
+        // RESULT: 1: I
+        //Breakpoint!
+        this.testPrivate
+    }
+
+    public val ExtClass.testPublic: Int
+        get() = a
+    private val ExtClass.testPrivate: Int
+        get() = a
+
+    companion object {
+        public val ExtClass.testCompPublic: Int
+            get() = a
+        private val ExtClass.testCompPrivate: Int
+            get() = a
+
+        fun testCompanion(extClass: ExtClass) {
+            // EXPRESSION: extClass.testCompPublic
+            // RESULT: 1: I
+            //Breakpoint!
+            extClass.testCompPublic
+
+            with(extClass) {
+                // EXPRESSION: testCompPublic
+                // RESULT: 1: I
+                //Breakpoint!
+                testCompPublic
+            }
+
+            // EXPRESSION: extClass.testCompPrivate
+            // RESULT: 1: I
+            //Breakpoint!
+            extClass.testCompPrivate
+
+            with(extClass) {
+                // EXPRESSION: testCompPrivate
+                // RESULT: 1: I
+                //Breakpoint!
+                testCompPrivate
+            }
+
+            extClass.testExtCompanion()
+        }
+
+        fun ExtClass.testExtCompanion() {
+            // EXPRESSION: testCompPublic
+            // RESULT: 1: I
+            //Breakpoint!
+            testCompPublic
+
+            // EXPRESSION: this.testCompPublic
+            // RESULT: 1: I
+            //Breakpoint!
+            this.testCompPublic
+
+            // EXPRESSION: testCompPrivate
+            // RESULT: 1: I
+            //Breakpoint!
+            testCompPrivate
+
+            // EXPRESSION: this.testCompPrivate
+            // RESULT: 1: I
+            //Breakpoint!
+            this.testCompPrivate
+        }
+    }
+}
+
+class ExtClass {
+    val a = 1
+}

idea/tests/org/jetbrains/kotlin/idea/debugger/evaluate/KotlinEvaluateExpressionTestGenerated.java

@@ -712,6 +712,12 @@ public class KotlinEvaluateExpressionTestGenerated extends AbstractKotlinEvaluat
             doMultipleBreakpointsTest(fileName);
         }
 
+        @TestMetadata("extensionMemberProperty.kt")
+        public void testExtensionMemberProperty() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/debugger/tinyApp/src/evaluate/multipleBreakpoints/extensionMemberProperty.kt");
+            doMultipleBreakpointsTest(fileName);
+        }
+
         @TestMetadata("funFromOuterClassInLamdba.kt")
         public void testFunFromOuterClassInLamdba() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/debugger/tinyApp/src/evaluate/multipleBreakpoints/funFromOuterClassInLamdba.kt");