Implement missing invoke resolution branch with context receivers

#KT-61937 fixed
Note: the current implementation only resolves for a single receiver
in each group. See KT-62712 and KT-62709 for the followup.

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java

@@ -12687,6 +12687,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/labelsFromClassNameForbidden.kt");
                 }
 
+                @Test
+                @TestMetadata("lambdaReceiverFromContext.kt")
+                public void testLambdaReceiverFromContext() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/lambdaReceiverFromContext.kt");
+                }
+
                 @Test
                 @TestMetadata("lazy.kt")
                 public void testLazy() throws Exception {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java

@@ -12687,6 +12687,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/labelsFromClassNameForbidden.kt");
                 }
 
+                @Test
+                @TestMetadata("lambdaReceiverFromContext.kt")
+                public void testLambdaReceiverFromContext() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/lambdaReceiverFromContext.kt");
+                }
+
                 @Test
                 @TestMetadata("lazy.kt")
                 public void testLazy() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java

@@ -12687,6 +12687,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/labelsFromClassNameForbidden.kt");
                 }
 
+                @Test
+                @TestMetadata("lambdaReceiverFromContext.kt")
+                public void testLambdaReceiverFromContext() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/lambdaReceiverFromContext.kt");
+                }
+
                 @Test
                 @TestMetadata("lazy.kt")
                 public void testLazy() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java

@@ -12693,6 +12693,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/labelsFromClassNameForbidden.kt");
                 }
 
+                @Test
+                @TestMetadata("lambdaReceiverFromContext.kt")
+                public void testLambdaReceiverFromContext() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/lambdaReceiverFromContext.kt");
+                }
+
                 @Test
                 @TestMetadata("lazy.kt")
                 public void testLazy() throws Exception {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirLightTreeBlackBoxCodegenTestGenerated.java

@@ -17883,6 +17883,12 @@ public class FirLightTreeBlackBoxCodegenTestGenerated extends AbstractFirLightTr
             runTest("compiler/testData/codegen/box/extensionClasses/generics.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaReceiverFromContext.kt")
+        public void testLambdaReceiverFromContext() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/lambdaReceiverFromContext.kt");
+        }
+
         @Test
         @TestMetadata("multiple.kt")
         public void testMultiple() throws Exception {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirLightTreeBlackBoxCodegenWithIrFakeOverrideGeneratorTestGenerated.java

@@ -17883,6 +17883,12 @@ public class FirLightTreeBlackBoxCodegenWithIrFakeOverrideGeneratorTestGenerated
             runTest("compiler/testData/codegen/box/extensionClasses/generics.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaReceiverFromContext.kt")
+        public void testLambdaReceiverFromContext() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/lambdaReceiverFromContext.kt");
+        }
+
         @Test
         @TestMetadata("multiple.kt")
         public void testMultiple() throws Exception {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirPsiBlackBoxCodegenTestGenerated.java

@@ -17883,6 +17883,12 @@ public class FirPsiBlackBoxCodegenTestGenerated extends AbstractFirPsiBlackBoxCo
             runTest("compiler/testData/codegen/box/extensionClasses/generics.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaReceiverFromContext.kt")
+        public void testLambdaReceiverFromContext() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/lambdaReceiverFromContext.kt");
+        }
+
         @Test
         @TestMetadata("multiple.kt")
         public void testMultiple() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/FirInvokeResolveTowerExtension.kt

@@ -443,6 +443,22 @@ private class InvokeFunctionResolveTask(
                 ExplicitReceiverKind.DISPATCH_RECEIVER
             )
         }
+        for ((depth, contextReceiverGroup) in towerDataElementsForName.contextReceiverGroups) {
+            val towerGroup =
+                TowerGroup
+                    .ContextReceiverGroup(depth)
+                    .InvokeExtensionWithImplicitReceiver
+                    .withGivenInvokeReceiverGroup(InvokeResolvePriority.INVOKE_EXTENSION)
+            val towerLevel = invokeReceiverValue.toMemberScopeTowerLevel()
+            // TODO: resolve for all receivers in the group, but implement the ambiguity diagnostics first. See KT-62712 and KT-69709
+            contextReceiverGroup.singleOrNull()?.let { contextReceiverValue ->
+                processLevel(
+                    towerLevel,
+                    info.withReceiverAsArgument(contextReceiverValue.receiverExpression), towerGroup,
+                    ExplicitReceiverKind.EXTENSION_RECEIVER
+                )
+            }
+        }
     }
 
     private suspend fun processLevelForRegularInvoke(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -542,7 +542,12 @@ class BodyResolveContext(
         val forMembersResolution =
             statics
                 .addLocalScope(parameterScope)
-                .addContextReceiverGroup(towerElementsForScript.implicitReceivers)
+                // TODO: temporary solution for avoiding problem described in KT-62712, flatten back after fix
+                .let { baseCtx ->
+                    towerElementsForScript.implicitReceivers.fold(baseCtx) { ctx, it ->
+                        ctx.addContextReceiverGroup(listOf(it))
+                    }
+                }
 
         val newContexts = FirRegularTowerDataContexts(
             regular = forMembersResolution,

compiler/testData/codegen/box/extensionClasses/lambdaReceiverFromContext.kt

@@ -0,0 +1,22 @@
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_K1: ANY
+// !LANGUAGE: +ContextReceivers
+
+class Ctx(val value: Int)
+
+fun Ctx.foo() = value + 4
+
+context(Ctx)
+class A {
+    fun bar(body: Ctx.() -> Int): Int {
+        return foo() *
+                body()
+    }
+}
+
+fun box(): String {
+    val res = with(Ctx(3)) {
+        A().bar { this.value * 2 }
+    }
+    return if (res == 42) "OK" else "NOK: $res"
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/lambdaReceiverFromContext.fir.kt

@@ -0,0 +1,13 @@
+// !LANGUAGE: +ContextReceivers
+
+class Ctx
+
+fun Ctx.foo() {}
+
+context(Ctx)
+class A {
+    fun bar(body: Ctx.() -> Unit) {
+        foo()
+        body()
+    }
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/lambdaReceiverFromContext.kt

@@ -0,0 +1,13 @@
+// !LANGUAGE: +ContextReceivers
+
+class Ctx
+
+fun Ctx.foo() {}
+
+context(Ctx)
+class A {
+    fun bar(body: Ctx.() -> Unit) {
+        foo()
+        body<!NO_VALUE_FOR_PARAMETER!>()<!>
+    }
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -12693,6 +12693,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/labelsFromClassNameForbidden.kt");
                 }
 
+                @Test
+                @TestMetadata("lambdaReceiverFromContext.kt")
+                public void testLambdaReceiverFromContext() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/lambdaReceiverFromContext.kt");
+                }
+
                 @Test
                 @TestMetadata("lazy.kt")
                 public void testLazy() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -17883,6 +17883,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/extensionClasses/generics.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaReceiverFromContext.kt")
+        public void testLambdaReceiverFromContext() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/lambdaReceiverFromContext.kt");
+        }
+
         @Test
         @TestMetadata("multiple.kt")
         public void testMultiple() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenWithIrInlinerTestGenerated.java

@@ -17883,6 +17883,12 @@ public class IrBlackBoxCodegenWithIrInlinerTestGenerated extends AbstractIrBlack
             runTest("compiler/testData/codegen/box/extensionClasses/generics.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaReceiverFromContext.kt")
+        public void testLambdaReceiverFromContext() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/lambdaReceiverFromContext.kt");
+        }
+
         @Test
         @TestMetadata("multiple.kt")
         public void testMultiple() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -14829,6 +14829,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/extensionClasses/generics.kt");
         }
 
+        @TestMetadata("lambdaReceiverFromContext.kt")
+        public void testLambdaReceiverFromContext() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/lambdaReceiverFromContext.kt");
+        }
+
         @TestMetadata("multiple.kt")
         public void testMultiple() throws Exception {
             runTest("compiler/testData/codegen/box/extensionClasses/multiple.kt");

libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/ScriptingHostTest.kt

@@ -292,6 +292,43 @@ class ScriptingHostTest : TestCase() {
         Assert.assertEquals(result, output)
     }
 
+    fun testScriptImplicitReceiversTransitiveVisibility() {
+        val result = listOf("42")
+        val script = """
+            import kotlin.script.experimental.jvmhost.test.forScript.p1.TestClass
+
+            fun TestClass.foo() = 6
+            
+            fun test(body: TestClass.() -> Int): Int {
+                return foo() *
+                    body()
+            }
+
+            println(test { 7 })
+        """.trimIndent()
+        val definition = createJvmScriptDefinitionFromTemplate<SimpleScriptTemplate>(
+            compilation = {
+                updateClasspath(classpathFromClass<kotlin.script.experimental.jvmhost.test.forScript.p1.TestClass>())
+                implicitReceivers(
+                    kotlin.script.experimental.jvmhost.test.forScript.p1.TestClass::class,
+                )
+            },
+            evaluation = {
+                implicitReceivers(
+                    kotlin.script.experimental.jvmhost.test.forScript.p1.TestClass(""),
+                )
+            }
+        )
+        val output = captureOut {
+            val evalRes = BasicJvmScriptingHost().eval(
+                script.toScriptSource(), definition.compilationConfiguration, definition.evaluationConfiguration
+            )
+            val retVal = evalRes.valueOrThrow().returnValue
+            if (retVal is ResultValue.Error) throw retVal.error
+        }.lines()
+        Assert.assertEquals(result, output)
+    }
+
     @Test
     fun testProvidedPropertiesNullability() {
         val stringType = KotlinType(String::class)