Extract Function: Extract smart-cast value as parameter

#KT-7576 Fixed

compiler/frontend/src/org/jetbrains/kotlin/cfg/pseudocode/pseudocodeUtil.kt

@@ -191,13 +191,9 @@ public fun JetElement.getContainingPseudocode(context: BindingContext): Pseudoco
             ?: getNonStrictParentOfType<JetProperty>()
             ?: return null
 
-    val enclosingPseudocodeDeclaration = if (pseudocodeDeclaration is JetFunctionLiteral) {
+    val enclosingPseudocodeDeclaration = (pseudocodeDeclaration as? JetFunctionLiteral)?.let {
-        parents(withItself = false).firstOrNull { it is JetDeclaration && it !is JetFunctionLiteral } as? JetDeclaration
+        it.parents(withItself = false).firstOrNull { it is JetDeclaration && it !is JetFunctionLiteral } as? JetDeclaration
-        ?: pseudocodeDeclaration
+    } ?: pseudocodeDeclaration
-    }
-    else {
-        pseudocodeDeclaration
-    }
 
     val enclosingPseudocode = PseudocodeUtil.generatePseudocode(enclosingPseudocodeDeclaration, context)
     return enclosingPseudocode.getPseudocodeByElement(pseudocodeDeclaration)

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/jetPsiUtil.kt

@@ -667,3 +667,12 @@ public fun JetExpression.getAnnotationEntries(): List<JetAnnotationEntry> {
         else -> emptyList<JetAnnotationEntry>()
     }
 }
+
+public fun JetElement.getQualifiedExpressionForSelector(): JetQualifiedExpression? {
+    val parent = getParent()
+    return if (parent is JetQualifiedExpression && parent.getSelectorExpression() == this) parent else null
+}
+
+public fun JetElement.getQualifiedExpressionForSelectorOrThis(): JetElement {
+    return getQualifiedExpressionForSelector() ?: this
+}

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

@@ -77,7 +77,7 @@ class RenameReplacement(override val parameter: Parameter): ParameterReplacement
 
     [suppress("PARAMETER_NAME_CHANGED_ON_OVERRIDE")]
     override fun invoke(e: JetElement): JetElement {
-        val expressionToReplace = e.getParent() as? JetThisExpression ?: e
+        var expressionToReplace = (e.getParent() as? JetThisExpression ?: e).getQualifiedExpressionForSelectorOrThis()
         val psiFactory = JetPsiFactory(e)
         val replacement =
                 if (expressionToReplace is JetOperationReferenceExpression) {

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

@@ -30,18 +30,22 @@ import java.util.Collections
 import org.jetbrains.kotlin.psi.psiUtil.isInsideOf
 import java.util.ArrayList
 import com.intellij.psi.PsiNamedElement
+import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.resolve.scopes.receivers.ThisReceiver
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.idea.core.compareDescriptors
+import org.jetbrains.kotlin.idea.core.refactoring.getContextForContainingDeclarationBody
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.idea.util.psi.patternMatching.JetPsiRange
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
+import org.jetbrains.kotlin.types.JetType
 
 data class ExtractionOptions(
         val inferUnitTypeForUnusedValues: Boolean = true,
@@ -65,7 +69,8 @@ data class ResolveResult(
 data class ResolvedReferenceInfo(
         val refExpr: JetSimpleNameExpression,
         val offsetInBody: Int,
-        val resolveResult: ResolveResult
+        val resolveResult: ResolveResult,
+        val smartCast: JetType?
 )
 
 data class ExtractionData(
@@ -105,6 +110,10 @@ data class ExtractionData(
     val originalStartOffset: Int?
         get() = originalElements.firstOrNull()?.let { e -> e.getTextRange()!!.getStartOffset() }
 
+    val commonParent = PsiTreeUtil.findCommonParent(originalElements) as JetElement
+
+    val bindingContext: BindingContext? by Delegates.lazy { commonParent.getContextForContainingDeclarationBody() }
+
     private val itFakeDeclaration by Delegates.lazy { JetPsiFactory(originalFile).createParameter("it: Any?") }
 
     val refOffsetToDeclaration by Delegates.lazy {
@@ -115,33 +124,44 @@ data class ExtractionData(
         }
 
         val originalStartOffset = originalStartOffset
+        val context = bindingContext
 
-        if (originalStartOffset != null) {
+        if (originalStartOffset != null && context != null) {
             val resultMap = HashMap<Int, ResolveResult>()
 
-            for ((ref, context) in JetFileReferencesResolver.resolve(originalFile, getExpressions())) {
+            val visitor = object: JetTreeVisitorVoid() {
-                if (ref !is JetSimpleNameExpression) continue
+                override fun visitQualifiedExpression(expression: JetQualifiedExpression) {
-                if (ref.getParent() is JetValueArgumentName) continue
+                    if (context[BindingContext.SMARTCAST, expression] != null) {
+                        expression.getSelectorExpression()?.accept(this)
+                        return
+                    }
 
-                val resolvedCall = ref.getResolvedCall(context)?.let {
+                    super.visitQualifiedExpression(expression)
-                    (it as? VariableAsFunctionResolvedCall)?.functionCall ?: it
                 }
 
-                val descriptor = context[BindingContext.REFERENCE_TARGET, ref]
+                override fun visitSimpleNameExpression(ref: JetSimpleNameExpression) {
-                if (descriptor == null) continue
+                    if (ref !is JetSimpleNameExpression) return
+                    if (ref.getParent() is JetValueArgumentName) return
 
-                val declaration = DescriptorToSourceUtilsIde.getAnyDeclaration(project, descriptor) as? PsiNamedElement
+                    val resolvedCall = ref.getResolvedCall(context)
-                        ?: if (isExtractableIt(descriptor, context)) itFakeDeclaration else continue
+                    val descriptor = context[BindingContext.REFERENCE_TARGET, ref] ?: return
+                    val declaration = DescriptorToSourceUtilsIde.getAnyDeclaration(project, descriptor) as? PsiNamedElement
+                                      ?: if (isExtractableIt(descriptor, context)) itFakeDeclaration else return
 
-                val offset = ref.getTextRange()!!.getStartOffset() - originalStartOffset
+                    val offset = ref.getTextRange()!!.getStartOffset() - originalStartOffset
-                resultMap[offset] = ResolveResult(ref, declaration, descriptor, resolvedCall)
+                    resultMap[offset] = ResolveResult(ref, declaration, descriptor, resolvedCall)
+                }
             }
+            getExpressions().forEach { it.accept(visitor) }
+
             resultMap
         }
         else Collections.emptyMap<Int, ResolveResult>()
     }
 
     fun getBrokenReferencesInfo(body: JetBlockExpression): List<ResolvedReferenceInfo> {
+        val originalContext = bindingContext ?: return listOf()
+
         val startOffset = body.getBlockContentOffset()
 
         val referencesInfo = ArrayList<ResolvedReferenceInfo>()
@@ -150,22 +170,33 @@ data class ExtractionData(
             if (ref !is JetSimpleNameExpression) continue
 
             val offset = ref.getTextRange()!!.getStartOffset() - startOffset
-            val originalResolveResult = refOffsetToDeclaration[offset]
+            val originalResolveResult = refOffsetToDeclaration[offset] ?: continue
-            if (originalResolveResult == null) continue
 
-            val parent = ref.getParent()
+            val smartCast: JetType?
-            if (parent is JetQualifiedExpression && parent.getSelectorExpression() == ref) {
+
+            // Qualified property reference: a.b
+            val qualifiedExpression = ref.getQualifiedExpressionForSelector()
+            if (qualifiedExpression != null) {
+                smartCast = originalContext[BindingContext.SMARTCAST, originalResolveResult.originalRefExpr.getParent() as JetExpression]
                 val receiverDescriptor =
                         (originalResolveResult.resolvedCall?.getDispatchReceiver() as? ThisReceiver)?.getDeclarationDescriptor()
-                if (!DescriptorUtils.isCompanionObject(receiverDescriptor) && parent.getReceiverExpression() !is JetSuperExpression) continue
+                if (smartCast == null
+                    && !DescriptorUtils.isCompanionObject(receiverDescriptor)
+                    && qualifiedExpression.getReceiverExpression() !is JetSuperExpression) continue
             }
+            else {
+                smartCast = originalContext[BindingContext.SMARTCAST, originalResolveResult.originalRefExpr]
+            }
+
+            val parent = ref.getParent()
+
             // Skip P in type references like 'P.Q'
             if (parent is JetUserType && (parent.getParent() as? JetUserType)?.getQualifier() == parent) continue
 
             val descriptor = context[BindingContext.REFERENCE_TARGET, ref]
-            if (!compareDescriptors(project, originalResolveResult.descriptor, descriptor)
+            val isBadRef = !compareDescriptors(project, originalResolveResult.descriptor, descriptor) || smartCast != null
-                    && !originalResolveResult.declaration.isInsideOf(originalElements)) {
+            if (isBadRef && !originalResolveResult.declaration.isInsideOf(originalElements)) {
-                referencesInfo.add(ResolvedReferenceInfo(ref, offset, originalResolveResult))
+                referencesInfo.add(ResolvedReferenceInfo(ref, offset, originalResolveResult, smartCast))
             }
         }
 

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

@@ -633,10 +633,10 @@ private fun ExtractionData.inferParametersInfo(
             info.replacementMap[refInfo.offsetInBody] = FqNameReplacement(originalDescriptor.importableFqNameSafe)
         }
         else {
-            val extractThis = hasThisReceiver || thisExpr != null
+            val extractThis = (hasThisReceiver && refInfo.smartCast == null) || thisExpr != null
-            val extractLocalVar =
+            val extractOrdinaryParameter =
                     originalDeclaration is JetMultiDeclarationEntry ||
-                            (originalDeclaration is JetProperty && originalDeclaration.isLocal()) ||
+                            originalDeclaration is JetProperty ||
                             originalDeclaration is JetParameter
 
             val extractFunctionRef =
@@ -647,11 +647,15 @@ private fun ExtractionData.inferParametersInfo(
 
             val descriptorToExtract = (if (extractThis) thisDescriptor else null) ?: originalDescriptor
 
-            val extractParameter = extractThis || extractLocalVar || extractFunctionRef
+            val extractParameter = extractThis || extractOrdinaryParameter || extractFunctionRef
             if (extractParameter) {
                 val parameterExpression = when {
-                    receiverToExtract is ExpressionReceiver -> receiverToExtract.getExpression()
+                    receiverToExtract is ExpressionReceiver -> {
-                    receiverToExtract.exists() -> null
+                        val receiverExpression = receiverToExtract.getExpression()
+                        // If p.q has a smart-cast, then extract entire qualified expression
+                        if (refInfo.smartCast != null) receiverExpression.getParent() as JetExpression else receiverExpression
+                    }
+                    receiverToExtract.exists() && refInfo.smartCast == null -> null
                     else -> (originalRef.getParent() as? JetThisExpression) ?: originalRef
                 }
 
@@ -683,8 +687,10 @@ private fun ExtractionData.inferParametersInfo(
                                 val label = if (descriptorToExtract is ClassDescriptor) "@${descriptorToExtract.getName().asString()}" else ""
                                 "this$label"
                             }
-                            else
+                            else {
-                                (thisExpr ?: ref).getText() ?: throw AssertionError("'this' reference shouldn't be empty: code fragment = $codeFragmentText")
+                                val argumentExpr = (thisExpr ?: ref).getQualifiedExpressionForSelectorOrThis()
+                                argumentExpr.getText() ?: throw AssertionError("'this' reference shouldn't be empty: code fragment = $codeFragmentText")
+                            }
                     if (extractFunctionRef) {
                         val receiverTypeText = (originalDeclaration as JetCallableDeclaration).getReceiverTypeReference()?.getText() ?: ""
                         argumentText = "$receiverTypeText::$argumentText"
@@ -807,24 +813,18 @@ fun ExtractionData.getDefaultVisibility(): String {
 }
 
 fun ExtractionData.performAnalysis(): AnalysisResult {
-    if (originalElements.isEmpty()) {
+    if (originalElements.isEmpty()) return AnalysisResult(null, Status.CRITICAL_ERROR, listOf(ErrorMessage.NO_EXPRESSION))
-        return AnalysisResult(null, Status.CRITICAL_ERROR, listOf(ErrorMessage.NO_EXPRESSION))
-    }
 
     val noContainerError = AnalysisResult(null, Status.CRITICAL_ERROR, listOf(ErrorMessage.NO_CONTAINER))
 
-    val commonParent = PsiTreeUtil.findCommonParent(originalElements) as JetElement
+    val bindingContext = bindingContext ?: return noContainerError
-
-    val bindingContext = commonParent.getContextForContainingDeclarationBody()
-    if (bindingContext == null) return noContainerError
-
-    val targetScope = JetScopeUtils.getResolutionScope(targetSibling, bindingContext)
 
     val pseudocode = commonParent.getContainingPseudocode(bindingContext) ?: return noContainerError
     val localInstructions = getLocalInstructions(pseudocode)
 
     val modifiedVarDescriptorsWithExpressions = localInstructions.getModifiedVarDescriptors(bindingContext)
 
+    val targetScope = JetScopeUtils.getResolutionScope(targetSibling, bindingContext)
     val paramsInfo = inferParametersInfo(commonParent, pseudocode, bindingContext, targetScope, modifiedVarDescriptorsWithExpressions.keySet())
     if (paramsInfo.errorMessage != null) {
         return AnalysisResult(null, Status.CRITICAL_ERROR, listOf(paramsInfo.errorMessage!!))

idea/testData/refactoring/extractFunction/basic/extensionValUnderSmartCast.kt

@@ -0,0 +1,14 @@
+// SUGGESTED_NAMES: i, getKm
+// PARAM_TYPES: A
+// PARAM_DESCRIPTOR: val a: A defined in test
+class A {
+
+}
+
+val A.meters: Int? get() = 1
+
+fun test() {
+    val a = A()
+    if (a.meters == null) return
+    val km = <selection>a.meters / 10</selection>
+}

idea/testData/refactoring/extractFunction/basic/extensionValUnderSmartCast.kt.after

@@ -0,0 +1,16 @@
+// SUGGESTED_NAMES: i, getKm
+// PARAM_TYPES: A
+// PARAM_DESCRIPTOR: val a: A defined in test
+class A {
+
+}
+
+val A.meters: Int? get() = 1
+
+fun test() {
+    val a = A()
+    if (a.meters == null) return
+    val km = i(a)
+}
+
+private fun i(a: A) = a.meters / 10

idea/testData/refactoring/extractFunction/basic/memberValUnderSmartCast.kt

@@ -0,0 +1,12 @@
+// SUGGESTED_NAMES: i, getKm
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: internal final val meters: kotlin.Int? defined in A
+class A {
+    val meters: Int? = 1
+}
+
+fun test() {
+    val a = A()
+    if (a.meters == null) return
+    val km = <selection>a.meters / 10</selection>
+}

idea/testData/refactoring/extractFunction/basic/memberValUnderSmartCast.kt.after

@@ -0,0 +1,14 @@
+// SUGGESTED_NAMES: i, getKm
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: internal final val meters: kotlin.Int? defined in A
+class A {
+    val meters: Int? = 1
+}
+
+fun test() {
+    val a = A()
+    if (a.meters == null) return
+    val km = i(a.meters)
+}
+
+private fun i(meters: Int) = meters / 10

idea/testData/refactoring/extractFunction/basic/memberValUnderSmartCastAndImplicitReceiver.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+// SUGGESTED_NAMES: i, getKm
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: internal final val meters: kotlin.Int? defined in A
+class A {
+    val meters: Int? = 1
+}
+
+fun test() {
+    val a = A()
+    with (a) {
+        if (meters == null) return
+        val km = <selection>meters / 10</selection>
+    }
+}

idea/testData/refactoring/extractFunction/basic/memberValUnderSmartCastAndImplicitReceiver.kt.after

@@ -0,0 +1,17 @@
+// WITH_RUNTIME
+// SUGGESTED_NAMES: i, getKm
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: internal final val meters: kotlin.Int? defined in A
+class A {
+    val meters: Int? = 1
+}
+
+fun test() {
+    val a = A()
+    with (a) {
+        if (meters == null) return
+        val km = i(meters)
+    }
+}
+
+private fun i(meters: Int) = meters / 10

idea/testData/refactoring/extractFunction/basic/topLevelValUnderSmartCast.kt

@@ -0,0 +1,9 @@
+// SUGGESTED_NAMES: i, getKm
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: internal val meters: kotlin.Int? defined in root package
+val meters: Int? = 1
+
+fun test() {
+    if (meters == null) return
+    val km = <selection>meters / 10</selection>
+}

idea/testData/refactoring/extractFunction/basic/topLevelValUnderSmartCast.kt.after

@@ -0,0 +1,11 @@
+// SUGGESTED_NAMES: i, getKm
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: internal val meters: kotlin.Int? defined in root package
+val meters: Int? = 1
+
+fun test() {
+    if (meters == null) return
+    val km = i(meters)
+}
+
+private fun i(meters: Int) = meters / 10

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

@@ -348,6 +348,12 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 doExtractFunctionTest(fileName);
             }
 
+            @TestMetadata("extensionValUnderSmartCast.kt")
+            public void testExtensionValUnderSmartCast() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/extensionValUnderSmartCast.kt");
+                doExtractFunctionTest(fileName);
+            }
+
             @TestMetadata("fakeOverride.kt")
             public void testFakeOverride() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/fakeOverride.kt");
@@ -456,6 +462,18 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 doExtractFunctionTest(fileName);
             }
 
+            @TestMetadata("memberValUnderSmartCast.kt")
+            public void testMemberValUnderSmartCast() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/memberValUnderSmartCast.kt");
+                doExtractFunctionTest(fileName);
+            }
+
+            @TestMetadata("memberValUnderSmartCastAndImplicitReceiver.kt")
+            public void testMemberValUnderSmartCastAndImplicitReceiver() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/memberValUnderSmartCastAndImplicitReceiver.kt");
+                doExtractFunctionTest(fileName);
+            }
+
             @TestMetadata("misdirectedRef.kt")
             public void testMisdirectedRef() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/misdirectedRef.kt");
@@ -498,6 +516,12 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 doExtractFunctionTest(fileName);
             }
 
+            @TestMetadata("topLevelValUnderSmartCast.kt")
+            public void testTopLevelValUnderSmartCast() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/topLevelValUnderSmartCast.kt");
+                doExtractFunctionTest(fileName);
+            }
+
             @TestMetadata("traitFunction.kt")
             public void testTraitFunction() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/traitFunction.kt");