Extraction Engine: Replace receiver occurrences with 'this' when necessary

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

@@ -71,7 +71,7 @@ data class TypeParameter(
         val originalConstraints: List<JetTypeConstraint>
 )
 
-interface Replacement: Function1<JetElement, JetElement>
+interface Replacement: Function2<ExtractableCodeDescriptor, JetElement, JetElement>
 
 interface ParameterReplacement : Replacement {
     val parameter: Parameter
@@ -81,19 +81,17 @@ interface ParameterReplacement : Replacement {
 class RenameReplacement(override val parameter: Parameter): ParameterReplacement {
     override fun copy(parameter: Parameter) = RenameReplacement(parameter)
 
-    override fun invoke(e: JetElement): JetElement {
+    override fun invoke(descriptor: ExtractableCodeDescriptor, e: JetElement): JetElement {
         var expressionToReplace = (e.getParent() as? JetThisExpression ?: e).let { it.getQualifiedExpressionForSelector() ?: it }
         val parameterName = JetPsiUtil.unquoteIdentifier(parameter.nameForRef)
         val replacingName =
                 if (e.getText().startsWith('`') || !KotlinNameSuggester.isIdentifier(parameterName)) "`$parameterName`" else parameterName
         val psiFactory = JetPsiFactory(e)
-        val replacement =
+        val replacement = when {
-                if (expressionToReplace is JetOperationReferenceExpression) {
+            parameter == descriptor.receiverParameter -> psiFactory.createExpression("this")
-                    psiFactory.createOperationName(replacingName)
+            expressionToReplace is JetOperationReferenceExpression -> psiFactory.createOperationName(replacingName)
-                }
+            else -> psiFactory.createSimpleName(replacingName)
-                else {
+        }
-                    psiFactory.createSimpleName(replacingName)
-                }
         return expressionToReplace.replaced(replacement)
     }
 }
@@ -101,17 +99,18 @@ class RenameReplacement(override val parameter: Parameter): ParameterReplacement
 class AddPrefixReplacement(override val parameter: Parameter): ParameterReplacement {
     override fun copy(parameter: Parameter) = AddPrefixReplacement(parameter)
 
-    override fun invoke(e: JetElement): JetElement {
+    override fun invoke(descriptor: ExtractableCodeDescriptor, e: JetElement): JetElement {
-        val selector = (e.getParent() as? JetCallExpression) ?: e
+        if (descriptor.receiverParameter == parameter) return e
-        val newExpr = selector.replace(JetPsiFactory(e).createExpression("${parameter.nameForRef}.${selector.getText()}")
-        ) as JetQualifiedExpression
 
-        return with(newExpr.getSelectorExpression()!!) { if (this is JetCallExpression) getCalleeExpression()!! else this }
+        val selector = (e.parent as? JetCallExpression) ?: e
+        val replacingExpression = JetPsiFactory(e).createExpression("${parameter.nameForRef}.${selector.text}")
+        val newExpr = (selector.replace(replacingExpression) as JetQualifiedExpression).selectorExpression!!
+        return (newExpr as? JetCallExpression)?.calleeExpression ?: newExpr
     }
 }
 
 class FqNameReplacement(val fqName: FqName): Replacement {
-    override fun invoke(e: JetElement): JetElement {
+    override fun invoke(descriptor: ExtractableCodeDescriptor, e: JetElement): JetElement {
         val thisExpr = e.getParent() as? JetThisExpression
         if (thisExpr != null) {
             return thisExpr.replaced(JetPsiFactory(e).createExpression(fqName.asString())).getQualifiedElementSelector()!!

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

@@ -1005,8 +1005,9 @@ fun ExtractableCodeDescriptor.validate(): ExtractableCodeDescriptorWithConflicts
         for ((originalOffset, resolveResult) in extractionData.refOffsetToDeclaration) {
             if (resolveResult.declaration.isInsideOf(extractionData.originalElements)) continue
 
-            val currentRefExpr = result.nameByOffset[originalOffset] as JetSimpleNameExpression?
+            val currentRefExpr = result.nameByOffset[originalOffset]?.let {
-            if (currentRefExpr == null) continue
+                (it as? JetThisExpression)?.instanceReference ?: it as? JetSimpleNameExpression
+            } ?: continue
 
             if (currentRefExpr.getParent() is JetThisExpression) continue
 

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

@@ -455,7 +455,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
     fun adjustDeclarationBody(declaration: JetNamedDeclaration) {
         val body = declaration.getGeneratedBody()
 
-        val exprReplacementMap = HashMap<JetElement, (JetElement) -> JetElement>()
+        val exprReplacementMap = HashMap<JetElement, (ExtractableCodeDescriptor, JetElement) -> JetElement>()
         val originalOffsetByExpr = LinkedHashMap<JetElement, Int>()
 
         val bodyOffset = body.getBlockContentOffset()
@@ -472,9 +472,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
             originalOffsetByExpr[expr!!] = offsetInBody
 
             descriptor.replacementMap[offsetInBody]?.let { replacement ->
-                if (replacement !is ParameterReplacement || replacement.parameter != descriptor.receiverParameter) {
+                exprReplacementMap[expr] = replacement
-                    exprReplacementMap[expr] = replacement
-                }
             }
         }
 
@@ -505,7 +503,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
 
         for ((expr, originalOffset) in originalOffsetByExpr) {
             if (expr.isValid()) {
-                nameByOffset.put(originalOffset, exprReplacementMap[expr]?.invoke(expr) ?: expr)
+                nameByOffset.put(originalOffset, exprReplacementMap[expr]?.invoke(descriptor, expr) ?: expr)
             }
         }
 

idea/testData/refactoring/extractFunction/parameters/extractThis/qualifiedThis.kt

@@ -0,0 +1,11 @@
+// PARAM_TYPES: kotlin.String, Comparable<String>, CharSequence, kotlin.Any
+// PARAM_DESCRIPTOR: public fun kotlin.String.test(): kotlin.Unit defined in root package
+fun String.foo(f: () -> Unit) {
+    f()
+}
+
+fun String.test() {
+    "sss".foo {
+        println(<selection>this@test</selection>)
+    }
+}

idea/testData/refactoring/extractFunction/parameters/extractThis/qualifiedThis.kt.after

@@ -0,0 +1,13 @@
+// PARAM_TYPES: kotlin.String, Comparable<String>, CharSequence, kotlin.Any
+// PARAM_DESCRIPTOR: public fun kotlin.String.test(): kotlin.Unit defined in root package
+fun String.foo(f: () -> Unit) {
+    f()
+}
+
+fun String.test() {
+    "sss".foo {
+        println(s())
+    }
+}
+
+private fun String.s() = this

idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/JetExtractionTestGenerated.java

@@ -2003,6 +2003,12 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                     doExtractFunctionTest(fileName);
                 }
 
+                @TestMetadata("qualifiedThis.kt")
+                public void testQualifiedThis() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/parameters/extractThis/qualifiedThis.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
                 @TestMetadata("qualifiedThisAsArgument.kt")
                 public void testQualifiedThisAsArgument() throws Exception {
                     String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/parameters/extractThis/qualifiedThisAsArgument.kt");