[LL API] Fix value capturing for multiple receivers

There were problem when 'this' bound symbol was used twice for a call

KT-61144

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt

@@ -14,9 +14,13 @@ import org.jetbrains.kotlin.analysis.api.diagnostics.KtDiagnostic
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.impl.base.util.KtCompiledFileForOutputFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.*
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.collectDiagnosticsForFile
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.resolve
+import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CodeFragmentCapturedId
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CodeFragmentCapturedValueAnalyzer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CompilationPeerCollector
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
@@ -324,7 +328,9 @@ internal class KtFirCompilerFacility(
 
         val capturedSymbols = capturedData.symbols
         val capturedValues = capturedSymbols.map { it.value }
-        val injectedSymbols = capturedSymbols.map { InjectedValue(it.symbol, it.typeRef, it.value.isMutated) }
+        val injectedSymbols = capturedSymbols.map {
+            InjectedValue(it.symbol, it.contextReceiverNumber, it.typeRef, it.value.isMutated)
+        }
 
         codeFragment.conversionData = CodeFragmentConversionData(
             classId = ClassId(FqName.ROOT, Name.identifier(configuration[CODE_FRAGMENT_CLASS_NAME] ?: "CodeFragment")),
@@ -332,7 +338,9 @@ internal class KtFirCompilerFacility(
             injectedSymbols
         )
 
-        val injectedSymbolMapping = injectedSymbols.associateBy { it.symbol }
+        val injectedSymbolMapping = injectedSymbols.associateBy {
+            CodeFragmentCapturedId(it.symbol, it.contextReceiverNumber)
+        }
         val injectedValueProvider = InjectedSymbolProvider(mainKtFile, injectedSymbolMapping)
 
         return CodeFragmentMappings(capturedValues, capturedData.files, injectedValueProvider)
@@ -340,7 +348,7 @@ internal class KtFirCompilerFacility(
 
     private class InjectedSymbolProvider(
         private val mainKtFile: KtFile,
-        private val injectedSymbolMapping: Map<FirBasedSymbol<*>, InjectedValue>
+        private val injectedSymbolMapping: Map<CodeFragmentCapturedId, InjectedValue>
     ) : (FirReference, Fir2IrConversionScope) -> InjectedValue? {
         override fun invoke(calleeReference: FirReference, conversionScope: Fir2IrConversionScope): InjectedValue? {
             val irFile = conversionScope.containingFileIfAny()
@@ -350,11 +358,13 @@ internal class KtFirCompilerFacility(
                 return null
             }
 
-            val symbol = when (calleeReference) {
-                is FirThisReference -> calleeReference.boundSymbol
-                else -> calleeReference.toResolvedSymbol<FirBasedSymbol<*>>()
+            val id = when (calleeReference) {
+                is FirThisReference -> calleeReference.boundSymbol?.let {
+                    CodeFragmentCapturedId(it, calleeReference.contextReceiverNumber)
+                }
+                else -> calleeReference.toResolvedSymbol<FirBasedSymbol<*>>()?.let { CodeFragmentCapturedId(it) }
             }
-            return injectedSymbolMapping[symbol]
+            return injectedSymbolMapping[id]
         }
     }
 

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

@@ -220,6 +220,24 @@ public class FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated exten
                 runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/localMutated.kt");
             }
 
+            @Test
+            @TestMetadata("multipleClassAndFunctionContextReceivers.kt")
+            public void testMultipleClassAndFunctionContextReceivers() throws Exception {
+                runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/multipleClassAndFunctionContextReceivers.kt");
+            }
+
+            @Test
+            @TestMetadata("multipleClassContextReceivers.kt")
+            public void testMultipleClassContextReceivers() throws Exception {
+                runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/multipleClassContextReceivers.kt");
+            }
+
+            @Test
+            @TestMetadata("multipleFunctionContextReceivers.kt")
+            public void testMultipleFunctionContextReceivers() throws Exception {
+                runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/multipleFunctionContextReceivers.kt");
+            }
+
             @Test
             @TestMetadata("nestedOuterClass.kt")
             public void testNestedOuterClass() throws Exception {

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

@@ -0,0 +1,16 @@
+ContainingClass[name: <this>; isMutated: false; displayText: this@Test]
+    context(R|Ctx1|, R|Ctx2|)
+    class Test : R|kotlin/Any|
+    R|Test|
+ContainingClass[name: <this>; isMutated: false; displayText: this@Test]
+    context(R|Ctx1|, R|Ctx2|)
+    class Test : R|kotlin/Any|
+    R|Test|
+ContextReceiver[name: Ctx3; isMutated: false; displayText: this@Ctx3]
+    context(R|Ctx3|, R|Ctx4|)
+    fun foo(): R|kotlin/Unit|
+    R|Ctx3|
+ContextReceiver[name: Ctx4; isMutated: false; displayText: this@Ctx4]
+    context(R|Ctx3|, R|Ctx4|)
+    fun foo(): R|kotlin/Unit|
+    R|Ctx4|

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

@@ -0,0 +1 @@
+useWithCtx1Ctx2Ctx3Ctx4()

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

@@ -0,0 +1,21 @@
+MODULE_FRAGMENT name:<Sources of main>
+  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>.Test, p1:<root>.Test, p2:<root>.Ctx3, p3:<root>.Ctx4) returnType:kotlin.Int
+        VALUE_PARAMETER name:p0 index:0 type:<root>.Test
+        VALUE_PARAMETER name:p1 index:1 type:<root>.Test
+        VALUE_PARAMETER name:p2 index:2 type:<root>.Ctx3
+        VALUE_PARAMETER name:p3 index:3 type:<root>.Ctx4
+        EXPRESSION_BODY
+          BLOCK type=kotlin.Int origin=null
+            CALL 'public final fun useWithCtx1Ctx2Ctx3Ctx4 ($context_receiver_0: <root>.Ctx1, $context_receiver_1: <root>.Ctx2, $context_receiver_2: <root>.Ctx3, $context_receiver_3: <root>.Ctx4): kotlin.Int declared in <root>.MultipleClassAndFunctionContextReceiversKt' type=kotlin.Int origin=null
+              $context_receiver_0: GET_FIELD 'FIELD FIELD_FOR_CLASS_CONTEXT_RECEIVER name:contextReceiverField0 type:<root>.Ctx1 visibility:private [final]' type=<root>.Ctx1 origin=null
+                receiver: GET_VAR 'p0: <root>.Test declared in <root>.CodeFragment.run' type=<root>.Test origin=null
+              $context_receiver_1: GET_FIELD 'FIELD FIELD_FOR_CLASS_CONTEXT_RECEIVER name:contextReceiverField1 type:<root>.Ctx2 visibility:private [final]' type=<root>.Ctx2 origin=null
+                receiver: GET_VAR 'p1: <root>.Test declared in <root>.CodeFragment.run' type=<root>.Test origin=null
+              $context_receiver_2: GET_VAR 'p2: <root>.Ctx3 declared in <root>.CodeFragment.run' type=<root>.Ctx3 origin=null
+              $context_receiver_3: GET_VAR 'p3: <root>.Ctx4 declared in <root>.CodeFragment.run' type=<root>.Ctx4 origin=null

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

@@ -0,0 +1,16 @@
+// !LANGUAGE: +ContextReceivers
+class Ctx1
+class Ctx2
+class Ctx3
+class Ctx4
+
+context(Ctx1, Ctx2, Ctx3, Ctx4)
+fun useWithCtx1Ctx2Ctx3Ctx4() = 3
+
+context(Ctx1, Ctx2)
+class Test {
+    context(Ctx3, Ctx4)
+    fun foo() {
+        <caret>val x = 1
+    }
+}

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

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

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

@@ -0,0 +1,8 @@
+ContainingClass[name: <this>; isMutated: false; displayText: this@Test]
+    context(R|Ctx1|, R|Ctx2|)
+    class Test : R|kotlin/Any|
+    R|Test|
+ContainingClass[name: <this>; isMutated: false; displayText: this@Test]
+    context(R|Ctx1|, R|Ctx2|)
+    class Test : R|kotlin/Any|
+    R|Test|

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

@@ -0,0 +1 @@
+useWithCtx1Ctx2()

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

@@ -0,0 +1,17 @@
+MODULE_FRAGMENT name:<Sources of main>
+  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>.Test, p1:<root>.Test) returnType:kotlin.Int
+        VALUE_PARAMETER name:p0 index:0 type:<root>.Test
+        VALUE_PARAMETER name:p1 index:1 type:<root>.Test
+        EXPRESSION_BODY
+          BLOCK type=kotlin.Int origin=null
+            CALL 'public final fun useWithCtx1Ctx2 ($context_receiver_0: <root>.Ctx1, $context_receiver_1: <root>.Ctx2): kotlin.Int declared in <root>.MultipleClassContextReceiversKt' type=kotlin.Int origin=null
+              $context_receiver_0: GET_FIELD 'FIELD FIELD_FOR_CLASS_CONTEXT_RECEIVER name:contextReceiverField0 type:<root>.Ctx1 visibility:private [final]' type=<root>.Ctx1 origin=null
+                receiver: GET_VAR 'p0: <root>.Test declared in <root>.CodeFragment.run' type=<root>.Test origin=null
+              $context_receiver_1: GET_FIELD 'FIELD FIELD_FOR_CLASS_CONTEXT_RECEIVER name:contextReceiverField1 type:<root>.Ctx2 visibility:private [final]' type=<root>.Ctx2 origin=null
+                receiver: GET_VAR 'p1: <root>.Test declared in <root>.CodeFragment.run' type=<root>.Test origin=null

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

@@ -0,0 +1,13 @@
+// !LANGUAGE: +ContextReceivers
+class Ctx1
+class Ctx2
+
+context(Ctx1, Ctx2)
+fun useWithCtx1Ctx2() = 3
+
+context(Ctx1, Ctx2)
+class Test {
+    fun foo() {
+        <caret>val x = 1
+    }
+}

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

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

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

@@ -0,0 +1,8 @@
+ContextReceiver[name: Ctx1; isMutated: false; displayText: this@Ctx1]
+    context(R|Ctx1|, R|Ctx2|)
+    fun foo(): R|kotlin/Unit|
+    R|Ctx1|
+ContextReceiver[name: Ctx2; isMutated: false; displayText: this@Ctx2]
+    context(R|Ctx1|, R|Ctx2|)
+    fun foo(): R|kotlin/Unit|
+    R|Ctx2|

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

@@ -0,0 +1 @@
+useWithCtx1Ctx2()

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

@@ -0,0 +1,15 @@
+MODULE_FRAGMENT name:<Sources of main>
+  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>.Ctx1, p1:<root>.Ctx2) returnType:kotlin.Int
+        VALUE_PARAMETER name:p0 index:0 type:<root>.Ctx1
+        VALUE_PARAMETER name:p1 index:1 type:<root>.Ctx2
+        EXPRESSION_BODY
+          BLOCK type=kotlin.Int origin=null
+            CALL 'public final fun useWithCtx1Ctx2 ($context_receiver_0: <root>.Ctx1, $context_receiver_1: <root>.Ctx2): kotlin.Int declared in <root>.MultipleFunctionContextReceiversKt' type=kotlin.Int origin=null
+              $context_receiver_0: GET_VAR 'p0: <root>.Ctx1 declared in <root>.CodeFragment.run' type=<root>.Ctx1 origin=null
+              $context_receiver_1: GET_VAR 'p1: <root>.Ctx2 declared in <root>.CodeFragment.run' type=<root>.Ctx2 origin=null

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

@@ -0,0 +1,11 @@
+// !LANGUAGE: +ContextReceivers
+class Ctx1
+class Ctx2
+
+context(Ctx1, Ctx2)
+fun useWithCtx1Ctx2() = 3
+
+context(Ctx1, Ctx2)
+fun foo() {
+    <caret>val x = 1
+}

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

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

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/compile/CodeFragmentCapturedValueAnalyzer.kt

@@ -43,9 +43,12 @@ import kotlin.collections.LinkedHashMap
 class CodeFragmentCapturedSymbol(
     val value: CodeFragmentCapturedValue,
     val symbol: FirBasedSymbol<*>,
-    val typeRef: FirTypeRef
+    val typeRef: FirTypeRef,
+    val contextReceiverNumber: Int = -1
 )
 
+data class CodeFragmentCapturedId(val symbol: FirBasedSymbol<*>, val contextReceiverNumber: Int = -1)
+
 object CodeFragmentCapturedValueAnalyzer {
     fun analyze(resolveSession: LLFirResolveSession, codeFragment: FirCodeFragment): CodeFragmentCapturedValueData {
         val selfSymbols = CodeFragmentDeclarationCollector().apply { codeFragment.accept(this) }.symbols.toSet()
@@ -76,7 +79,7 @@ private class CodeFragmentCapturedValueVisitor(
     private val resolveSession: LLFirResolveSession,
     private val selfSymbols: Set<FirBasedSymbol<*>>,
 ) : FirDefaultVisitorVoid() {
-    private val collectedMappings = LinkedHashMap<FirBasedSymbol<*>, CodeFragmentCapturedSymbol>()
+    private val collectedMappings = LinkedHashMap<CodeFragmentCapturedId, CodeFragmentCapturedSymbol>()
     private val collectedFiles = LinkedHashSet<KtFile>()
 
     private val assignmentLhs = mutableListOf<FirBasedSymbol<*>>()
@@ -112,10 +115,11 @@ private class CodeFragmentCapturedValueVisitor(
                 if (symbol != null && symbol !in selfSymbols) {
                     val isCrossingInlineBounds = isCrossingInlineBounds(element, symbol)
                     val capturedValue = CodeFragmentCapturedValue.SuperClass(symbol.classId, isCrossingInlineBounds)
-                    register(symbol, CodeFragmentCapturedSymbol(capturedValue, symbol, element.superTypeRef))
+                    register(CodeFragmentCapturedSymbol(capturedValue, symbol, element.superTypeRef))
                 }
             }
             is FirThisReference -> {
+                val contextReceiverNumber = element.contextReceiverNumber
                 val symbol = element.boundSymbol
                 if (symbol != null && symbol !in selfSymbols) {
                     when (symbol) {
@@ -124,11 +128,10 @@ private class CodeFragmentCapturedValueVisitor(
                                 val isCrossingInlineBounds = isCrossingInlineBounds(element, symbol)
                                 val capturedValue = CodeFragmentCapturedValue.ContainingClass(symbol.classId, isCrossingInlineBounds)
                                 val typeRef = buildResolvedTypeRef { type = symbol.defaultType() }
-                                register(symbol, CodeFragmentCapturedSymbol(capturedValue, symbol, typeRef))
+                                register(CodeFragmentCapturedSymbol(capturedValue, symbol, typeRef, contextReceiverNumber))
                             }
                         }
                         is FirFunctionSymbol<*> -> {
-                            val contextReceiverNumber = element.contextReceiverNumber
                             if (contextReceiverNumber >= 0) {
                                 val contextReceiver = symbol.resolvedContextReceivers[contextReceiverNumber]
                                 val labelName = contextReceiver.labelName
@@ -136,7 +139,9 @@ private class CodeFragmentCapturedValueVisitor(
                                     val isCrossingInlineBounds = isCrossingInlineBounds(element, symbol)
                                     val capturedValue = CodeFragmentCapturedValue
                                         .ContextReceiver(contextReceiverNumber, labelName, isCrossingInlineBounds)
-                                    register(symbol, CodeFragmentCapturedSymbol(capturedValue, symbol, contextReceiver.typeRef))
+                                    register(
+                                        CodeFragmentCapturedSymbol(capturedValue, symbol, contextReceiver.typeRef, contextReceiverNumber)
+                                    )
                                 }
                             } else {
                                 val labelName = element.labelName
@@ -146,7 +151,7 @@ private class CodeFragmentCapturedValueVisitor(
                                 val typeRef = symbol.receiverParameter?.typeRef ?: error("Receiver parameter not found")
                                 val isCrossingInlineBounds = isCrossingInlineBounds(element, symbol)
                                 val capturedValue = CodeFragmentCapturedValue.ExtensionReceiver(labelName, isCrossingInlineBounds)
-                                register(symbol, CodeFragmentCapturedSymbol(capturedValue, symbol, typeRef))
+                                register(CodeFragmentCapturedSymbol(capturedValue, symbol, typeRef, contextReceiverNumber))
                             }
                         }
                     }
@@ -172,7 +177,7 @@ private class CodeFragmentCapturedValueVisitor(
             is FirValueParameterSymbol -> {
                 val isCrossingInlineBounds = isCrossingInlineBounds(element, symbol)
                 val capturedValue = CodeFragmentCapturedValue.Local(symbol.name, symbol.isMutated, isCrossingInlineBounds)
-                register(symbol, CodeFragmentCapturedSymbol(capturedValue, symbol, symbol.resolvedReturnTypeRef))
+                register(CodeFragmentCapturedSymbol(capturedValue, symbol, symbol.resolvedReturnTypeRef))
             }
             is FirPropertySymbol -> {
                 if (symbol.isLocal) {
@@ -182,14 +187,14 @@ private class CodeFragmentCapturedValueVisitor(
                         symbol.hasDelegate -> CodeFragmentCapturedValue.LocalDelegate(symbol.name, symbol.isMutated, isCrossingInlineBounds)
                         else -> CodeFragmentCapturedValue.Local(symbol.name, symbol.isMutated, isCrossingInlineBounds)
                     }
-                    register(symbol, CodeFragmentCapturedSymbol(capturedValue, symbol, symbol.resolvedReturnTypeRef))
+                    register(CodeFragmentCapturedSymbol(capturedValue, symbol, symbol.resolvedReturnTypeRef))
                 }
             }
             is FirBackingFieldSymbol -> {
                 val propertyName = symbol.propertySymbol.name
                 val isCrossingInlineBounds = isCrossingInlineBounds(element, symbol)
                 val capturedValue = CodeFragmentCapturedValue.BackingField(propertyName, symbol.isMutated, isCrossingInlineBounds)
-                register(symbol, CodeFragmentCapturedSymbol(capturedValue, symbol, symbol.resolvedReturnTypeRef))
+                register(CodeFragmentCapturedSymbol(capturedValue, symbol, symbol.resolvedReturnTypeRef))
             }
             is FirNamedFunctionSymbol -> {
                 if (symbol.isLocal) {
@@ -201,12 +206,13 @@ private class CodeFragmentCapturedValueVisitor(
         if (symbol.callableId == StandardClassIds.Callables.coroutineContext) {
             val isCrossingInlineBounds = isCrossingInlineBounds(element, symbol)
             val capturedValue = CodeFragmentCapturedValue.CoroutineContext(isCrossingInlineBounds)
-            register(symbol, CodeFragmentCapturedSymbol(capturedValue, symbol, symbol.resolvedReturnTypeRef))
+            register(CodeFragmentCapturedSymbol(capturedValue, symbol, symbol.resolvedReturnTypeRef))
         }
     }
 
-    private fun register(symbol: FirBasedSymbol<*>, mapping: CodeFragmentCapturedSymbol) {
-        val previousMapping = collectedMappings[symbol]
+    private fun register(mapping: CodeFragmentCapturedSymbol) {
+        val id = CodeFragmentCapturedId(mapping.symbol, mapping.contextReceiverNumber)
+        val previousMapping = collectedMappings[id]
 
         if (previousMapping != null) {
             val previousValue = previousMapping.value
@@ -220,8 +226,8 @@ private class CodeFragmentCapturedValueVisitor(
             }
         }
 
-        collectedMappings[symbol] = mapping
-        registerFile(symbol)
+        collectedMappings[id] = mapping
+        registerFile(mapping.symbol)
     }
 
     private fun registerFile(symbol: FirBasedSymbol<*>) {

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

@@ -120,6 +120,24 @@ public class CodeFragmentCapturingTestGenerated extends AbstractCodeFragmentCapt
         runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/localMutated.kt");
     }
 
+    @Test
+    @TestMetadata("multipleClassAndFunctionContextReceivers.kt")
+    public void testMultipleClassAndFunctionContextReceivers() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/multipleClassAndFunctionContextReceivers.kt");
+    }
+
+    @Test
+    @TestMetadata("multipleClassContextReceivers.kt")
+    public void testMultipleClassContextReceivers() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/multipleClassContextReceivers.kt");
+    }
+
+    @Test
+    @TestMetadata("multipleFunctionContextReceivers.kt")
+    public void testMultipleFunctionContextReceivers() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/multipleFunctionContextReceivers.kt");
+    }
+
     @Test
     @TestMetadata("nestedOuterClass.kt")
     public void testNestedOuterClass() throws Exception {

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/CodeFragmentUtils.kt

@@ -21,7 +21,7 @@ class CodeFragmentConversionData(
     val injectedValues: List<InjectedValue>
 )
 
-class InjectedValue(val symbol: FirBasedSymbol<*>, val typeRef: FirTypeRef, val isMutated: Boolean) {
+class InjectedValue(val symbol: FirBasedSymbol<*>, val contextReceiverNumber: Int, val typeRef: FirTypeRef, val isMutated: Boolean) {
     val irParameterSymbol: IrValueParameterSymbol = IrValueParameterSymbolImpl()
 }