Debugger: more precise context expression in codeFragment

#KT-11380 Fixed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/CodeFragmentAnalyzer.kt

@@ -21,8 +21,6 @@ import org.jetbrains.kotlin.idea.project.ResolveElementCache
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.codeFragmentUtil.suppressDiagnosticsInDebugMode
-import org.jetbrains.kotlin.psi.psiUtil.parents
-import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
@@ -33,7 +31,6 @@ import org.jetbrains.kotlin.resolve.scopes.utils.addImportingScopes
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor
-import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import javax.inject.Inject
 
 class CodeFragmentAnalyzer(
@@ -81,15 +78,7 @@ class CodeFragmentAnalyzer(
                    is KtFunctionLiteral -> this.bodyExpression?.statements?.lastOrNull()
                    is KtDeclarationWithBody -> this.bodyExpression
                    is KtBlockExpression -> this.statements.lastOrNull()
-                   else -> {
+                   else -> null
-                       val previousSibling = this.siblings(forward = false, withItself = false).firstIsInstanceOrNull<KtExpression>()
-                       if (previousSibling != null) return previousSibling
-                       for (parent in this.parents) {
-                           if (parent is KtWhenEntry || parent is KtIfExpression || parent is KtBlockExpression) return this
-                           if (parent is KtExpression) return parent
-                       }
-                       null
-                   }
                } ?: this
     }
 

idea/src/org/jetbrains/kotlin/idea/debugger/evaluate/extractFunctionForDebuggerUtil.kt

@@ -24,6 +24,7 @@ import com.intellij.psi.PsiManager
 import com.intellij.psi.impl.PsiModificationTrackerImpl
 import org.jetbrains.kotlin.idea.actions.internal.KotlinInternalMode
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.caches.resolve.analyzeFullyAndGetResult
 import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.intentions.InsertExplicitTypeArgumentsIntention
@@ -35,6 +36,10 @@ import org.jetbrains.kotlin.idea.util.psi.patternMatching.toRange
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.codeFragmentUtil.suppressDiagnosticsInDebugMode
 import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
+import org.jetbrains.kotlin.renderer.DescriptorRenderer
+import org.jetbrains.kotlin.resolve.BindingContext.IMPLICIT_RECEIVER_SMARTCAST
+import org.jetbrains.kotlin.resolve.BindingContext.SMARTCAST
 
 fun getFunctionForExtractedFragment(
         codeFragment: KtCodeFragment,
@@ -109,25 +114,68 @@ fun getFunctionForExtractedFragment(
 }
 
 fun addDebugExpressionIntoTmpFileForExtractFunction(originalFile: KtFile, codeFragment: KtCodeFragment, line: Int): List<KtExpression> {
-    val context = codeFragment.getOriginalContext()
+    codeFragment.markContextElement()
-    context?.IS_CONTEXT_ELEMENT = true
+    codeFragment.markSmartCasts()
+
     val tmpFile = originalFile.copy() as KtFile
     tmpFile.suppressDiagnosticsInDebugMode = true
-    context?.IS_CONTEXT_ELEMENT = false
 
-    val contextElement = getExpressionToAddDebugExpressionBefore(tmpFile, context, line) ?: return emptyList()
+    val contextElement = getExpressionToAddDebugExpressionBefore(tmpFile, codeFragment.getOriginalContext(), line) ?: return emptyList()
 
     addImportsToFile(codeFragment.importsAsImportList(), tmpFile)
 
-    return addDebugExpressionBeforeContextElement(codeFragment, contextElement)
+    val contentElementsInTmpFile = addDebugExpressionBeforeContextElement(codeFragment, contextElement)
+    contentElementsInTmpFile.forEach { it.insertSmartCasts() }
+
+    codeFragment.clearContextElement()
+    codeFragment.clearSmartCasts()
+
+    return contentElementsInTmpFile
 }
 
 private var PsiElement.IS_CONTEXT_ELEMENT: Boolean by NotNullableCopyableUserDataProperty(Key.create("IS_CONTEXT_ELEMENT"), false)
 
+private fun KtCodeFragment.markContextElement() {
+    getOriginalContext()?.IS_CONTEXT_ELEMENT = true
+}
+
+private fun KtCodeFragment.clearContextElement() {
+    getOriginalContext()?.IS_CONTEXT_ELEMENT = false
+}
+
 private fun KtFile.findContextElement(): KtElement? {
     return this.findDescendantOfType { it.IS_CONTEXT_ELEMENT == true }
 }
 
+private var PsiElement.DEBUG_SMART_CAST: PsiElement? by CopyableUserDataProperty(Key.create("DEBUG_SMART_CAST"))
+
+private fun KtCodeFragment.markSmartCasts() {
+    val bindingContext = analyzeFullyAndGetResult().bindingContext
+    val factory = KtPsiFactory(project)
+
+    getContentElement()?.forEachDescendantOfType<KtExpression> { expression ->
+        val smartCast = bindingContext.get(SMARTCAST, expression) ?: bindingContext.get(IMPLICIT_RECEIVER_SMARTCAST, expression)
+        if (smartCast != null) {
+            val smartCastedExpression = factory.createExpressionByPattern(
+                    "($0 as ${DescriptorRenderer.FQ_NAMES_IN_TYPES.renderType(smartCast)})",
+                    expression) as KtParenthesizedExpression
+
+            expression.DEBUG_SMART_CAST = smartCastedExpression
+        }
+    }
+}
+
+private fun KtExpression.insertSmartCasts() {
+    forEachDescendantOfType<KtExpression> {
+        val replacement = it.DEBUG_SMART_CAST
+        if (replacement != null) runReadAction { it.replace(replacement) }
+    }
+}
+
+private fun KtCodeFragment.clearSmartCasts() {
+    getContentElement()?.forEachDescendantOfType<KtExpression> { it.DEBUG_SMART_CAST = null }
+}
+
 private fun addImportsToFile(newImportList: KtImportList?, tmpFile: KtFile) {
     if (newImportList != null && newImportList.imports.isNotEmpty()) {
         val tmpFileImportList = tmpFile.importList

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

@@ -0,0 +1,15 @@
+LineBreakpoint created at smartcasts.kt:19
+LineBreakpoint created at smartcasts.kt:29
+!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! smartcasts.SmartcastsKt
+Connected to the target VM, address: '!HOST_NAME!:PORT_NAME!', transport: 'socket'
+smartcasts.kt:19
+Compile bytecode for derived.prop
+smartcasts.kt:19
+Compile bytecode for derived.prop
+smartcasts.kt:29
+Compile bytecode for nullable.prop
+smartcasts.kt:29
+Compile bytecode for nullable.prop
+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/smartcasts.kt

@@ -0,0 +1,35 @@
+package smartcasts
+
+fun main(args: Array<String>) {
+    test1(Derived())
+    test1(Base())
+
+    test2(Derived())
+    test2(null)
+}
+
+// EXPRESSION: derived.prop
+// RESULT: 1: I
+
+// EXPRESSION: derived.prop
+// RESULT: java.lang.ClassCastException: smartcasts.Base cannot be cast to smartcasts.Derived: Ljava/lang/ClassCastException;
+fun test1(derived: Base) =
+        derived is Derived &&
+        //Breakpoint!
+        derived.prop == 1
+
+// EXPRESSION: nullable.prop
+// RESULT: 1: I
+
+// EXPRESSION: nullable.prop
+// RESULT: Method threw 'kotlin.TypeCastException' exception.
+fun test2(nullable: Derived?) =
+        nullable != null &&
+        //Breakpoint!
+        nullable.prop == 1
+
+class Derived : Base() {
+    val prop = 1
+}
+
+open class Base

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

@@ -814,6 +814,12 @@ public class KotlinEvaluateExpressionTestGenerated extends AbstractKotlinEvaluat
             doMultipleBreakpointsTest(fileName);
         }
 
+        @TestMetadata("smartcasts.kt")
+        public void testSmartcasts() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/debugger/tinyApp/src/evaluate/multipleBreakpoints/smartcasts.kt");
+            doMultipleBreakpointsTest(fileName);
+        }
+
         @TestMetadata("whenEntry.kt")
         public void testWhenEntry() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/debugger/tinyApp/src/evaluate/multipleBreakpoints/whenEntry.kt");