Implement simple cache for dynamically created wrapper for lambda

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/IrInterpreter.kt

@@ -573,7 +573,7 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
     }
 
     private fun interpretFunctionExpression(expression: IrFunctionExpression) {
-        val function = KFunctionState(expression.function, expression.type.classOrNull!!.owner)
+        val function = KFunctionState(expression.function, expression.type.classOrNull!!.owner, environment)
         if (expression.function.isLocal) callStack.storeUpValues(function)
         callStack.pushState(function)
     }
@@ -586,6 +586,7 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
 
         val function = KFunctionState(
             reference,
+            environment,
             dispatchReceiver?.let { Variable(irFunction.getDispatchReceiver()!!, it) },
             extensionReceiver?.let { Variable(irFunction.getExtensionReceiver()!!, it) }
         )

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/IrInterpreterEnvironment.kt

@@ -7,10 +7,12 @@ package org.jetbrains.kotlin.ir.interpreter
 
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.interpreter.stack.CallStack
 import org.jetbrains.kotlin.ir.interpreter.state.Common
 import org.jetbrains.kotlin.ir.interpreter.state.Complex
+import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.isSubclassOf
 
@@ -22,6 +24,7 @@ class IrInterpreterEnvironment(
     internal val irExceptions = mutableListOf<IrClass>()
     internal var mapOfEnums = mutableMapOf<IrSymbol, Complex>()
     internal var mapOfObjects = mutableMapOf<IrSymbol, Complex>()
+    internal var cachedLambdasAndReferences = mutableMapOf<IrFunctionSymbol, IrFunctionSymbol>()
 
     init {
         mapOfObjects[irBuiltIns.unitClass] = Common(irBuiltIns.unitClass.owner)

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/intrinsics/IntrinsicImplementations.kt

@@ -222,7 +222,7 @@ internal object ArrayConstructor : IntrinsicBase() {
 
         val initSymbol = irFunction.valueParameters[1].symbol
         val state = callStack.getState(initSymbol).let {
-            (it as? KFunctionState) ?: (it as KPropertyState).convertGetterToKFunctionState()
+            (it as? KFunctionState) ?: (it as KPropertyState).convertGetterToKFunctionState(environment)
         }
         // if property was converted, then we must replace symbol in memory to get correct receiver later
         callStack.setState(initSymbol, state)

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/state/reflection/KClassState.kt

@@ -51,7 +51,10 @@ internal class KClassState(val classReference: IrClass, override val irClass: Ir
                             else -> TODO()
                         }
                     }
-                    is IrFunction -> KFunctionProxy(KFunctionState(it, callInterceptor.irBuiltIns), callInterceptor)
+                    is IrFunction -> {
+                        val irClass = callInterceptor.irBuiltIns.kFunctionN(it.valueParameters.size)
+                        KFunctionProxy(KFunctionState(it, irClass, callInterceptor.environment), callInterceptor)
+                    }
                     else -> TODO()
                 }
             }
@@ -62,7 +65,10 @@ internal class KClassState(val classReference: IrClass, override val irClass: Ir
         if (_constructors != null) return _constructors!!
         _constructors = classReference.declarations
             .filterIsInstance<IrConstructor>()
-            .map { KFunctionProxy(KFunctionState(it, callInterceptor.irBuiltIns), callInterceptor) }
+            .map {
+                val irClass = callInterceptor.irBuiltIns.kFunctionN(it.valueParameters.size)
+                KFunctionProxy(KFunctionState(it, irClass, callInterceptor.environment), callInterceptor)
+            }
         return _constructors!!
     }
 

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/state/reflection/KFunctionState.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.ir.interpreter.state.reflection
 
-import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrFunction
@@ -14,10 +13,6 @@ import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
 import org.jetbrains.kotlin.ir.expressions.putArgument
 import org.jetbrains.kotlin.ir.interpreter.*
-import org.jetbrains.kotlin.ir.interpreter.CallInterceptor
-import org.jetbrains.kotlin.ir.interpreter.TEMP_FUNCTION_FOR_INTERPRETER
-import org.jetbrains.kotlin.ir.interpreter.createTempClass
-import org.jetbrains.kotlin.ir.interpreter.createTempFunction
 import org.jetbrains.kotlin.ir.interpreter.proxy.reflection.KParameterProxy
 import org.jetbrains.kotlin.ir.interpreter.proxy.reflection.KTypeParameterProxy
 import org.jetbrains.kotlin.ir.interpreter.proxy.reflection.KTypeProxy
@@ -34,72 +29,91 @@ import kotlin.reflect.KType
 import kotlin.reflect.KTypeParameter
 
 internal class KFunctionState(
-    val irFunction: IrFunction, override val irClass: IrClass, override val fields: MutableList<Variable>,
+    val irFunction: IrFunction,
+    override val irClass: IrClass,
+    environment: IrInterpreterEnvironment,
+    override val fields: MutableList<Variable> = mutableListOf()
 ) : ReflectionState(), StateWithClosure {
     override val upValues: MutableList<Variable> = mutableListOf()
     private var _parameters: List<KParameter>? = null
     private var _returnType: KType? = null
     private var _typeParameters: List<KTypeParameter>? = null
 
-    private val functionClass: IrClass
+    val invokeSymbol: IrFunctionSymbol = environment.cachedLambdasAndReferences
-    val invokeSymbol: IrFunctionSymbol
+        .getOrDefault(
+            irFunction.symbol,
+            createInvokeFunction(
+                irFunction,
+                irClass,
+                irFunction.dispatchReceiverParameter?.let { getField(it.symbol) } != null,
+                irFunction.extensionReceiverParameter?.let { getField(it.symbol) } != null
+            ).symbol
+        )
 
-    init {
+    companion object {
-        val invokeFunction = irClass.declarations.filterIsInstance<IrSimpleFunction>().single { it.name == OperatorNameConventions.INVOKE }
+        fun createInvokeFunction(
-        // TODO do we need new class here? if yes, do we need different names for temp classes?
+            irFunction: IrFunction, irClass: IrClass, hasDispatchReceiver: Boolean, hasExtensionReceiver: Boolean
-        functionClass = createTempClass(Name.identifier("Function\$0")).apply { parent = irFunction.parent }
+        ): IrSimpleFunction {
+            val invokeFunction = irClass.declarations
+                .filterIsInstance<IrSimpleFunction>()
+                .single { it.name == OperatorNameConventions.INVOKE }
+            // TODO do we need new class here? if yes, do we need different names for temp classes?
+            val functionClass = createTempClass(Name.identifier("Function\$0")).apply { parent = irFunction.parent }
 
-        functionClass.superTypes += irClass.defaultType
+            functionClass.superTypes += irClass.defaultType
-        functionClass.declarations += createTempFunction(
+            val newFunctionToInvoke = createTempFunction(
-            OperatorNameConventions.INVOKE, irFunction.returnType, TEMP_FUNCTION_FOR_INTERPRETER
+                OperatorNameConventions.INVOKE, irFunction.returnType, TEMP_FUNCTION_FOR_INTERPRETER
-        ).apply impl@{
+            ).apply impl@{
-            parent = functionClass
+                parent = functionClass
-            overriddenSymbols = listOf(invokeFunction.symbol)
+                overriddenSymbols = listOf(invokeFunction.symbol)
 
-            dispatchReceiverParameter = invokeFunction.dispatchReceiverParameter?.deepCopyWithSymbols(initialParent = this)
+                dispatchReceiverParameter = invokeFunction.dispatchReceiverParameter?.deepCopyWithSymbols(initialParent = this)
-            valueParameters = mutableListOf()
+                valueParameters = mutableListOf()
 
-            val call = when (irFunction) {
+                val call = when (irFunction) {
-                is IrSimpleFunction -> irFunction.createCall()
+                    is IrSimpleFunction -> irFunction.createCall()
-                is IrConstructor -> irFunction.createConstructorCall()
+                    is IrConstructor -> irFunction.createConstructorCall()
-                else -> TODO("Unsupported symbol $symbol for invoke")
+                    else -> TODO("Unsupported symbol $symbol for invoke")
-            }.apply {
+                }.apply {
-                val dispatchParameter = irFunction.dispatchReceiverParameter
+                    val dispatchParameter = irFunction.dispatchReceiverParameter
-                val extensionParameter = irFunction.extensionReceiverParameter
+                    val extensionParameter = irFunction.extensionReceiverParameter
 
-                if (dispatchParameter != null) {
+                    if (dispatchParameter != null) {
-                    dispatchReceiver = dispatchParameter.createGetValue()
+                        dispatchReceiver = dispatchParameter.createGetValue()
-                    if (getField(dispatchParameter.symbol) == null) (this@impl.valueParameters as MutableList) += dispatchParameter
+                        if (!hasDispatchReceiver) (this@impl.valueParameters as MutableList) += dispatchParameter
-                }
+                    }
-                if (extensionParameter != null) {
+                    if (extensionParameter != null) {
-                    extensionReceiver = extensionParameter.createGetValue()
+                        extensionReceiver = extensionParameter.createGetValue()
-                    if (getField(extensionParameter.symbol) == null) (this@impl.valueParameters as MutableList) += extensionParameter
+                        if (!hasExtensionReceiver) (this@impl.valueParameters as MutableList) += extensionParameter
-                }
+                    }
-                irFunction.valueParameters.forEach {
+                    irFunction.valueParameters.forEach {
-                    putArgument(it, it.createGetValue())
+                        putArgument(it, it.createGetValue())
-                    (this@impl.valueParameters as MutableList) += it
+                        (this@impl.valueParameters as MutableList) += it
+                    }
                 }
+
+                body = listOf(this.createReturn(call)).wrapWithBlockBody()
             }
-
+            functionClass.declarations += newFunctionToInvoke
-            body = listOf(this.createReturn(call)).wrapWithBlockBody()
+            return newFunctionToInvoke
-            invokeSymbol = this.symbol
         }
     }
 
-    constructor(irFunction: IrFunction, irClass: IrClass) : this(irFunction, irClass, mutableListOf())
+    constructor(
-
+        functionReference: IrFunctionReference,
-    constructor(functionReference: IrFunctionReference, dispatchReceiver: Variable?, extensionReceiver: Variable?) : this(
+        environment: IrInterpreterEnvironment,
+        dispatchReceiver: Variable?,
+        extensionReceiver: Variable?
+    ) : this(
         functionReference.symbol.owner,
         functionReference.type.classOrNull!!.owner,
+        environment,
         listOfNotNull(dispatchReceiver, extensionReceiver).toMutableList()
     ) {
         // receivers are used in comparison of two functions in KFunctionProxy
         upValues += fields
     }
 
-    constructor(irFunction: IrFunction, irBuiltIns: IrBuiltIns) :
-            this(irFunction, irBuiltIns.kFunctionN(irFunction.valueParameters.size), mutableListOf())
-
     override fun getIrFunctionByIrCall(expression: IrCall): IrFunction? {
         if (expression.symbol.owner.name == OperatorNameConventions.INVOKE) return invokeSymbol.owner
         return super.getIrFunctionByIrCall(expression)

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/state/reflection/KPropertyState.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrProperty
 import org.jetbrains.kotlin.ir.expressions.IrPropertyReference
 import org.jetbrains.kotlin.ir.interpreter.CallInterceptor
+import org.jetbrains.kotlin.ir.interpreter.IrInterpreterEnvironment
 import org.jetbrains.kotlin.ir.interpreter.proxy.reflection.KParameterProxy
 import org.jetbrains.kotlin.ir.interpreter.proxy.reflection.KTypeProxy
 import org.jetbrains.kotlin.ir.interpreter.state.State
@@ -23,10 +24,10 @@ internal class KPropertyState(val property: IrProperty, override val irClass: Ir
     private var _parameters: List<KParameter>? = null
     private var _returnType: KType? = null
 
-    fun convertGetterToKFunctionState(): KFunctionState {
+    fun convertGetterToKFunctionState(environment: IrInterpreterEnvironment): KFunctionState {
         val getterClass = irClass.getIrClassOfReflection("getter")
         val functionType = getterClass.superTypes.single { it.classOrNull?.owner?.name?.asString() == "Function1" }
-        return KFunctionState(property.getter!!, functionType.classOrNull!!.owner)
+        return KFunctionState(property.getter!!, functionType.classOrNull!!.owner, environment)
     }
 
     fun getParameters(callInterceptor: CallInterceptor): List<KParameter> {