[Analysis API] Keep the collected 'FirTowerDataContext' mutable

'ContextCollector' is used for computing context of 'FirCodeFragment's.
Code fragments themselves might contain additional smart cast operations
that modify the context receiver stack.

^KT-63056 Fixed

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerated.java

@@ -196,6 +196,12 @@ public class FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerate
                 runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverPropertyLabeled.kt");
             }
 
+            @Test
+            @TestMetadata("extensionReceiverSmartCasted.kt")
+            public void testExtensionReceiverSmartCasted() throws Exception {
+                runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.kt");
+            }
+
             @Test
             @TestMetadata("initializer.kt")
             public void testInitializer() throws Exception {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated.java

@@ -196,6 +196,12 @@ public class FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated exten
                 runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverPropertyLabeled.kt");
             }
 
+            @Test
+            @TestMetadata("extensionReceiverSmartCasted.kt")
+            public void testExtensionReceiverSmartCasted() throws Exception {
+                runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.kt");
+            }
+
             @Test
             @TestMetadata("initializer.kt")
             public void testInitializer() throws Exception {

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.capturing.txt

@@ -0,0 +1,3 @@
+ExtensionReceiver[name: apply; isMutated: false; displayText: this@apply]
+    apply@fun R|Foo|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>
+    R|Foo|

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.fragment.kt

@@ -0,0 +1,2 @@
+this as FooImpl
+n

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

@@ -0,0 +1,16 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:fragment.kt
+    CLASS CLASS name:CodeFragment modality:FINAL visibility:public superTypes:[kotlin.Any]
+      $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.CodeFragment
+      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 <> (p0:<root>.Foo) returnType:kotlin.Int
+        VALUE_PARAMETER name:p0 index:0 type:<root>.Foo
+        EXPRESSION_BODY
+          BLOCK type=kotlin.Int origin=null
+            TYPE_OP type=<root>.FooImpl origin=CAST typeOperand=<root>.FooImpl
+              GET_VAR 'p0: <root>.Foo declared in <root>.CodeFragment.run' type=<root>.Foo origin=null
+            CALL 'public final fun <get-n> (): kotlin.Int declared in <root>.FooImpl' type=kotlin.Int origin=GET_PROPERTY
+              $this: TYPE_OP type=<root>.FooImpl origin=IMPLICIT_CAST typeOperand=<root>.FooImpl
+                GET_VAR 'p0: <root>.Foo declared in <root>.CodeFragment.run' type=<root>.Foo origin=null

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.kt

@@ -0,0 +1,13 @@
+interface Foo
+
+class FooImpl : Foo {
+    val n: Int = 5
+}
+
+fun makeFoo(): Foo = FooImpl()
+
+fun main() {
+    makeFoo().apply {
+        <caret>Unit
+    }
+}

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.txt

@@ -0,0 +1,5 @@
+public final class CodeFragment {
+    // source: 'fragment.kt'
+    public method <init>(): void
+    public final static method run(p0: Foo): int
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirTowerDataContextCollector.kt

@@ -51,7 +51,7 @@ internal class FirTowerDataContextAllElementsCollector : FirResolveContextCollec
     override fun addStatementContext(statement: FirStatement, context: BodyResolveContext) {
         val closestParentExpression = statement.psi?.closestParentExpressionWithSameContextOrSelf() ?: return
         // FIR body transform may alter the context if there are implicit receivers with smartcast
-        elementsToContext[closestParentExpression] = context.towerDataContext.createSnapshot()
+        elementsToContext[closestParentExpression] = context.towerDataContext.createSnapshot(keepMutable = false)
     }
 
     override fun addDeclarationContext(declaration: FirDeclaration, context: BodyResolveContext) {
@@ -62,7 +62,7 @@ internal class FirTowerDataContextAllElementsCollector : FirResolveContextCollec
         // we do not collect contexts for such declarations
         if (psi is KtDelegatedSuperTypeEntry) return
 
-        elementsToContext[psi] = context.towerDataContext.createSnapshot()
+        elementsToContext[psi] = context.towerDataContext.createSnapshot(keepMutable = false)
     }
 
     override fun addClassHeaderContext(declaration: FirRegularClass, context: FirTowerDataContext) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt

@@ -264,7 +264,7 @@ private class ContextCollectorVisitor(
             }
         }
 
-        val towerDataContextSnapshot = context.towerDataContext.createSnapshot()
+        val towerDataContextSnapshot = context.towerDataContext.createSnapshot(keepMutable = true)
 
         for ((index, oldType) in oldReceiverTypes) {
             implicitReceiverStack.replaceReceiverType(index, oldType)

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/CodeFragmentCapturingTestGenerated.java

@@ -72,6 +72,12 @@ public class CodeFragmentCapturingTestGenerated extends AbstractCodeFragmentCapt
         runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverPropertyLabeled.kt");
     }
 
+    @Test
+    @TestMetadata("extensionReceiverSmartCasted.kt")
+    public void testExtensionReceiverSmartCasted() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.kt");
+    }
+
     @Test
     @TestMetadata("initializer.kt")
     public void testInitializer() throws Exception {