[Analysis API] Use SELF contexts for code fragment context expressions

Expressions can contribute smart casts, and those should normally
be ignored (a context element should be interpreted like the code
fragment prepends it).

^KT-70860 Fixed

Author: yanex

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/resultNothingBeforeSmartCast.ir.txt

@@ -5,8 +5,9 @@ MODULE_FRAGMENT
       CONSTRUCTOR visibility:public returnType:<root>.CodeFragment [primary]
         BLOCK_BODY
           DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in kotlin.Any'
-      FUN name:run visibility:public modality:FINAL returnType:kotlin.Any?
+      FUN name:run visibility:public modality:FINAL returnType:kotlin.String
         VALUE_PARAMETER kind:Regular name:p0 index:0 type:kotlin.String?
         EXPRESSION_BODY
-          BLOCK type=kotlin.Nothing? origin=null
-            GET_VAR 'p0: kotlin.String? declared in <root>.CodeFragment.run' type=kotlin.String? origin=null
+          BLOCK type=kotlin.String origin=null
+            TYPE_OP type=kotlin.String origin=IMPLICIT_CAST typeOperand=kotlin.String
+              GET_VAR 'p0: kotlin.String? declared in <root>.CodeFragment.run' type=kotlin.String? origin=null

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/resultNothingBeforeSmartCast.txt

@@ -18,8 +18,8 @@ public final class CodeFragment {
     MAXLOCALS = 1
 
   // access flags 0x19
-  public final static run(Ljava/lang/String;)Ljava/lang/Object;
-  @Lorg/jetbrains/annotations/Nullable;() // invisible
+  public final static run(Ljava/lang/String;)Ljava/lang/String;
+  @Lorg/jetbrains/annotations/NotNull;() // invisible
     // annotable parameter count: 1 (invisible)
     @Lorg/jetbrains/annotations/Nullable;() // invisible, parameter 0
    L0

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/resultNothingSmartCast.ir.txt

@@ -5,8 +5,9 @@ MODULE_FRAGMENT
       CONSTRUCTOR visibility:public returnType:<root>.CodeFragment [primary]
         BLOCK_BODY
           DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in kotlin.Any'
-      FUN name:run visibility:public modality:FINAL returnType:kotlin.Any?
+      FUN name:run visibility:public modality:FINAL returnType:kotlin.String
         VALUE_PARAMETER kind:Regular name:p0 index:0 type:kotlin.String?
         EXPRESSION_BODY
-          BLOCK type=kotlin.Nothing? origin=null
-            GET_VAR 'p0: kotlin.String? declared in <root>.CodeFragment.run' type=kotlin.String? origin=null
+          BLOCK type=kotlin.String origin=null
+            TYPE_OP type=kotlin.String origin=IMPLICIT_CAST typeOperand=kotlin.String
+              GET_VAR 'p0: kotlin.String? declared in <root>.CodeFragment.run' type=kotlin.String? origin=null

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/resultNothingSmartCast.kt

@@ -7,8 +7,8 @@ fun main() {
     var str: String? = null
     str = "not null"
     println(str)
-    str = null
-    <caret_context>println(str)
+    <caret_context>str = null
+    println(str)
 }
 
 

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/resultNothingSmartCast.txt

@@ -18,8 +18,8 @@ public final class CodeFragment {
     MAXLOCALS = 1
 
   // access flags 0x19
-  public final static run(Ljava/lang/String;)Ljava/lang/Object;
-  @Lorg/jetbrains/annotations/Nullable;() // invisible
+  public final static run(Ljava/lang/String;)Ljava/lang/String;
+  @Lorg/jetbrains/annotations/NotNull;() // invisible
     // annotable parameter count: 1 (invisible)
     @Lorg/jetbrains/annotations/Nullable;() // invisible, parameter 0
    L0

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt

@@ -54,9 +54,14 @@ import org.jetbrains.kotlin.fir.types.isResolved
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.psi
+import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtProperty
+import org.jetbrains.kotlin.psi.KtStatementExpression
+import org.jetbrains.kotlin.psi.KtVariableDeclaration
 import org.jetbrains.kotlin.utils.exceptions.checkWithAttachment
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment
@@ -778,7 +783,16 @@ private class LLFirBodyTargetResolver(target: LLFirResolveTarget) : LLFirAbstrac
 
         return if (contextKtFile != null) {
             val contextFirFile = resolutionFacade.getOrBuildFirFile(contextKtFile)
-            val elementContext = ContextCollector.process(resolutionFacade, contextFirFile, contextPsiElement)
+
+            // Avoid using body context of expressions/statements, as those can contribute additional smart casts.
+            // Still use body contexts for declarations (e.g., to be able to address parameters of a primary constructor).
+            val preferBodyContext = when (contextPsiElement) {
+                is KtDeclaration -> contextPsiElement !is KtProperty || !contextPsiElement.isLocal
+                is KtBlockExpression -> true
+                else -> false
+            }
+
+            val elementContext = ContextCollector.process(resolutionFacade, contextFirFile, contextPsiElement, preferBodyContext)
                 ?: errorWithAttachment("Cannot find enclosing context for ${contextPsiElement::class}") {
                     withPsiEntry("contextPsiElement", contextPsiElement)
                 }