[PSI2IR, IR] Support functional types and literals for IR

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

@@ -16238,6 +16238,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                     runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/fromKEEP/dp.kt");
                 }
 
+                @Test
+                @TestMetadata("functionalType.kt")
+                public void testFunctionalType() throws Exception {
+                    runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/fromKEEP/functionalType.kt");
+                }
+
                 @Test
                 @TestMetadata("monoidSum.kt")
                 public void testMonoidSum() throws Exception {

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/reflection/KFunctionProxy.kt

@@ -52,11 +52,15 @@ internal class KFunctionProxy(
         // TODO check arity
         var index = 0
         val dispatchReceiver = state.irFunction.dispatchReceiverParameter?.let { environment.convertToState(args[index++], it.type) }
+        val contextReceiverParametersCount = state.irFunction.contextReceiverParametersCount
+        val contextReceivers = state.irFunction.valueParameters.take(contextReceiverParametersCount).map {
+            environment.convertToState(args[index++], it.type)
+        }
         val extensionReceiver = state.irFunction.extensionReceiverParameter?.let { environment.convertToState(args[index++], it.type) }
-        val argsVariables = state.irFunction.valueParameters.map { parameter ->
+        val argsVariables = state.irFunction.valueParameters.drop(contextReceiverParametersCount).map { parameter ->
             environment.convertToState(args[index++], parameter.type)
         }
-        val valueArguments = listOfNotNull(dispatchReceiver, extensionReceiver) + argsVariables
+        val valueArguments = listOfNotNull(dispatchReceiver) + contextReceivers + listOfNotNull(extensionReceiver) + argsVariables
         return callInterceptor.interceptProxy(state.irFunction, valueArguments)
     }
 

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/FunctionGenerator.kt

@@ -327,17 +327,16 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
 
         val bodyGenerator = createBodyGenerator(irFunction.symbol)
 
-        val contextReceiversNumber = ktContextReceiverParameterElements.size
-        assert(functionDescriptor.contextReceiverParameters.size == contextReceiversNumber)
-        irFunction.contextReceiverParametersCount = contextReceiversNumber
+        val contextReceiverParametersCount = functionDescriptor.contextReceiverParameters.size
+        irFunction.contextReceiverParametersCount = contextReceiverParametersCount
         irFunction.valueParameters += functionDescriptor.contextReceiverParameters.mapIndexed { i, additionalReceiver ->
-            declareParameter(additionalReceiver, ktContextReceiverParameterElements[i], irFunction, null, i)
+            declareParameter(additionalReceiver, ktContextReceiverParameterElements.getOrNull(i) ?: ktParameterOwner, irFunction, null, i)
         }
 
         // Declare all the value parameters up first.
         irFunction.valueParameters += functionDescriptor.valueParameters.mapIndexed { i, valueParameterDescriptor ->
             val ktParameter = DescriptorToSourceUtils.getSourceFromDescriptor(valueParameterDescriptor) as? KtParameter
-            declareParameter(valueParameterDescriptor, ktParameter, irFunction, null, i + contextReceiversNumber)
+            declareParameter(valueParameterDescriptor, ktParameter, irFunction, null, i + contextReceiverParametersCount)
         }
         // Only after value parameters have been declared, generate default values. This ensures
         // that forward references to other parameters works in default value lambdas. For example:
@@ -345,7 +344,7 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
         // fun f(f1: () -> String = { f2() },
         //       f2: () -> String) = f1()
         if (withDefaultValues) {
-            irFunction.valueParameters.drop(contextReceiversNumber).forEachIndexed { index, irValueParameter ->
+            irFunction.valueParameters.drop(contextReceiverParametersCount).forEachIndexed { index, irValueParameter ->
                 val valueParameterDescriptor = functionDescriptor.valueParameters[index]
                 val ktParameter = DescriptorToSourceUtils.getSourceFromDescriptor(valueParameterDescriptor) as? KtParameter
                 irValueParameter.defaultValue = ktParameter?.defaultValue?.let { defaultValue ->

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/descriptors/IrDescriptorBasedFunctionFactory.kt

@@ -367,6 +367,7 @@ class IrDescriptorBasedFunctionFactory(
             newFunction.overriddenSymbols = descriptor.overriddenDescriptors.map { symbolTable.referenceSimpleFunction(it.original) }
             newFunction.dispatchReceiverParameter = descriptor.dispatchReceiverParameter?.let { newFunction.createValueParameter(it) }
             newFunction.extensionReceiverParameter = descriptor.extensionReceiverParameter?.let { newFunction.createValueParameter(it) }
+            newFunction.contextReceiverParametersCount = descriptor.contextReceiverParameters.size
             newFunction.valueParameters = descriptor.valueParameters.map { newFunction.createValueParameter(it) }
             newFunction.correspondingPropertySymbol = property
 

compiler/testData/codegen/box/extensionFunctions/contextReceivers/fromKEEP/functionalType.kt

@@ -0,0 +1,18 @@
+// !LANGUAGE: +ContextReceivers
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_FIR: JVM_IR
+
+class Param
+class O {
+    val o = "o"
+}
+class K {
+    val k = "k"
+}
+
+context(O)
+fun <T> K.f(g: context(O) K.(Param) -> T) = g(this@O, this@K, Param())
+
+fun box() = with(O()) {
+    K().f { o + k }
+}

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

@@ -16238,6 +16238,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                     runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/fromKEEP/dp.kt");
                 }
 
+                @Test
+                @TestMetadata("functionalType.kt")
+                public void testFunctionalType() throws Exception {
+                    runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/fromKEEP/functionalType.kt");
+                }
+
                 @Test
                 @TestMetadata("monoidSum.kt")
                 public void testMonoidSum() throws Exception {