Implement interpretation for lambdas and anonymous functions

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

@@ -10,10 +10,7 @@ import org.jetbrains.kotlin.backend.common.interpreter.stack.*
 import org.jetbrains.kotlin.builtins.DefaultBuiltIns
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
-import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
-import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
@@ -79,7 +76,7 @@ class IrInterpreter(irModule: IrModuleFragment) {
             }
 
             return when (code) {
-                Code.RETURN -> when (this) {
+                Code.RETURN -> when (this) { // TODO check label
                     is IrCall, is IrReturnableBlock -> Code.NEXT
                     else -> Code.RETURN
                 }
@@ -215,9 +212,6 @@ class IrInterpreter(irModule: IrModuleFragment) {
     private fun interpretCall(expression: IrCall, data: Frame): Code {
         val newFrame = InterpreterFrame()
 
-        interpretValueParameters(expression, data).also { if (it != Code.NEXT) return it }
-        val valueParameters = expression.symbol.descriptor.valueParameters.map { Variable(it, data.popReturnValue()) }
-
         // dispatch receiver processing
         val rawDispatchReceiver = expression.dispatchReceiver
         rawDispatchReceiver?.interpret(data)?.also { if (it != Code.NEXT) return it }
@@ -225,7 +219,10 @@ class IrInterpreter(irModule: IrModuleFragment) {
         val irFunctionReceiver = if (expression.superQualifierSymbol == null) dispatchReceiver else (dispatchReceiver as Complex).superType
         // it is important firstly to add receiver, then arguments; this order is used in builtin method call
         val irFunction = irFunctionReceiver?.getIrFunction(expression.symbol.descriptor) ?: expression.symbol.owner
-        irFunctionReceiver?.let { newFrame.addVar(Variable(irFunction.symbol.getDispatchReceiver()!!, it)) }
+        irFunctionReceiver?.let { receiverState ->
+            // dispatch receiver is null if irFunction is lambda and so there is no need to add it as variable in frame
+            irFunction.symbol.getDispatchReceiver()?.let { newFrame.addVar(Variable(it, receiverState)) }
+        }
 
         // extension receiver processing
         val rawExtensionReceiver = expression.extensionReceiver
@@ -233,7 +230,11 @@ class IrInterpreter(irModule: IrModuleFragment) {
         val extensionReceiver = rawExtensionReceiver?.let { data.popReturnValue() }
         extensionReceiver?.let { newFrame.addVar(Variable(irFunction.symbol.getExtensionReceiver()!!, it)) }
 
-        newFrame.addAll(valueParameters)
+        // if irFunction is lambda and it has receiver, then first descriptor must be taken from extension receiver
+        val receiverAsFirstArgument = if (irFunction.isLocalLambda()) listOfNotNull(irFunction.symbol.getExtensionReceiver()) else listOf()
+        val valueParametersDescriptors = receiverAsFirstArgument + irFunction.descriptor.valueParameters
+        interpretValueParameters(expression, data).also { if (it != Code.NEXT) return it }
+        newFrame.addAll(valueParametersDescriptors.map { Variable(it, data.popReturnValue()) })
 
         val isWrapper = (dispatchReceiver is Wrapper && rawExtensionReceiver == null) ||
                 (extensionReceiver is Wrapper && rawDispatchReceiver == null)
@@ -304,6 +305,11 @@ class IrInterpreter(irModule: IrModuleFragment) {
     }
 
     private fun interpretBlock(block: IrBlock, data: Frame): Code {
+        if (block.isLambdaFunction()) {
+            val irLambdaFunction = block.statements.filterIsInstance<IrFunctionReference>().first()
+            data.pushReturnValue(Lambda(irLambdaFunction.symbol.owner, irLambdaFunction.type.classOrNull!!.owner))
+            return Code.NEXT
+        }
         return interpretStatements(block.statements, data)
     }
 

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/interpreter/Utils.kt

@@ -10,12 +10,10 @@ import org.jetbrains.kotlin.backend.common.interpreter.stack.Frame
 import org.jetbrains.kotlin.backend.common.interpreter.stack.Primitive
 import org.jetbrains.kotlin.backend.common.interpreter.stack.State
 import org.jetbrains.kotlin.backend.common.interpreter.stack.Wrapper
+import org.jetbrains.kotlin.builtins.functions.FunctionInvokeDescriptor
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
-import org.jetbrains.kotlin.ir.declarations.IrAnnotationContainer
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
@@ -63,7 +61,8 @@ private fun DeclarationDescriptor.isSubtypeOf(other: DeclarationDescriptor): Boo
 }
 
 private fun DeclarationDescriptor.hasSameNameAs(other: DeclarationDescriptor): Boolean {
-    return this is VariableDescriptor && other is VariableDescriptor && this.name == other.name
+    return (this is VariableDescriptor && other is VariableDescriptor && this.name == other.name) ||
+            (this is FunctionInvokeDescriptor && other is FunctionInvokeDescriptor && this.valueParameters.size == other.valueParameters.size)
 }
 
 fun IrCall.isAbstract(): Boolean {
@@ -82,6 +81,14 @@ fun IrFunction.isFakeOverridden(): Boolean {
     return this.symbol.owner.isFakeOverride
 }
 
+fun IrContainerExpression.isLambdaFunction(): Boolean {
+    return this.origin == IrStatementOrigin.LAMBDA || this.origin == IrStatementOrigin.ANONYMOUS_FUNCTION
+}
+
+fun IrDeclaration.isLocalLambda(): Boolean {
+    return this.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+}
+
 fun State.toIrExpression(expression: IrExpression): IrExpression {
     return when (this) {
         is Primitive<*> ->

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/interpreter/stack/State.kt

@@ -29,6 +29,7 @@ import java.lang.invoke.MethodType
 
 interface State {
     val fields: MutableList<Variable>
+    val irClass: IrClass
 
     fun getState(descriptor: DeclarationDescriptor): State? {
         return fields.firstOrNull { it.descriptor.equalTo(descriptor) }?.state
@@ -41,19 +42,18 @@ interface State {
 
 class Primitive<T>(var value: T, val type: IrType) : State {
     override val fields: MutableList<Variable> = mutableListOf()
+    override val irClass: IrClass = type.classOrNull!!.owner
 
     init {
-        val properties = type.classOrNull!!.owner.declarations.filterIsInstance<IrProperty>()
+        val properties = irClass.declarations.filterIsInstance<IrProperty>()
         for (property in properties) {
-            val propertySignature = CompileTimeFunction(property.name.asString(), listOf(type.getName()))
+            val propertySignature = CompileTimeFunction(property.name.asString(), listOf(irClass.descriptor.defaultType.toString()))
             val propertyValue = unaryFunctions[propertySignature]?.invoke(value)
                 ?: throw NoSuchMethodException("For given property $propertySignature there is no entry in unary map")
             fields.add(Variable(property.descriptor, Primitive(propertyValue, property.backingField!!.type)))
         }
     }
 
-    private fun IrType.getName() = this.classOrNull!!.descriptor.defaultType.toString()
-
     override fun setState(newVar: Variable) {
         newVar.state as? Primitive<T> ?: throw IllegalArgumentException("Cannot set $newVar in current $this")
         value = newVar.state.value
@@ -65,7 +65,7 @@ class Primitive<T>(var value: T, val type: IrType) : State {
 
     override fun getIrFunction(descriptor: FunctionDescriptor): IrFunction? {
         // must add property's getter to declaration's list because they are not present in ir class for primitives
-        val declarations = type.classOrNull!!.owner.declarations.map { if (it is IrProperty) it.getter else it }
+        val declarations = irClass.declarations.map { if (it is IrProperty) it.getter else it }
         return declarations.filterIsInstance<IrFunction>()
             .filter { it.descriptor.name == descriptor.name }
             .firstOrNull { it.descriptor.valueParameters.map { it.type } == descriptor.valueParameters.map { it.type } }
@@ -92,11 +92,11 @@ class Primitive<T>(var value: T, val type: IrType) : State {
     }
 
     override fun toString(): String {
-        return "Primitive(value=$value, type=${type.getName()})"
+        return "Primitive(value=$value, type=${irClass.descriptor.defaultType})"
     }
 }
 
-open class Complex(private var type: IrClass, override val fields: MutableList<Variable>) : State {
+open class Complex(override var irClass: IrClass, override val fields: MutableList<Variable>) : State {
     var superType: Complex? = null
     var instance: Complex? = null
 
@@ -106,7 +106,7 @@ open class Complex(private var type: IrClass, override val fields: MutableList<V
     }
 
     fun getReceiver(): DeclarationDescriptor {
-        return type.thisReceiver!!.descriptor
+        return irClass.thisReceiver!!.descriptor
     }
 
     override fun setState(newVar: Variable) {
@@ -117,31 +117,33 @@ open class Complex(private var type: IrClass, override val fields: MutableList<V
     }
 
     override fun getIrFunction(descriptor: FunctionDescriptor): IrFunction? {
-        return type.declarations.filterIsInstance<IrFunction>()
+        return irClass.declarations.filterIsInstance<IrFunction>()
             .filter { it.descriptor.name == descriptor.name }
             .firstOrNull { it.descriptor.valueParameters.map { it.type } == descriptor.valueParameters.map { it.type } }
     }
 
     override fun copy(): State {
-        return Complex(type, fields).apply {
+        return Complex(irClass, fields).apply {
             this@apply.superType = this@Complex.superType
             this@apply.instance = this@Complex.instance
         }
     }
 
     override fun toString(): String {
-        return "Complex(obj='${type.fqNameForIrSerialization}', super=$superType, values=$fields)"
+        return "Complex(obj='${irClass.fqNameForIrSerialization}', super=$superType, values=$fields)"
     }
 }
 
-class Wrapper(val value: Any, private val type: IrClass) : Complex(type, mutableListOf()) {
+class Wrapper(val value: Any, override var irClass: IrClass) : Complex(irClass, mutableListOf()) {
-    private val typeFqName = type.fqNameForIrSerialization.toUnsafe()
+    private val typeFqName = irClass.fqNameForIrSerialization.toUnsafe()
-    private val receiverClass = type.defaultType.getClass(true)
+    private val receiverClass = irClass.defaultType.getClass(true)
 
     fun getMethod(irFunction: IrFunction): MethodHandle {
-        // used to get correct java method
+        // intrinsicName is used to get correct java method
         // for example, method 'get' in kotlin StringBuilder is actually 'charAt' in java StringBuilder
         val intrinsicName = irFunction.getEvaluateIntrinsicValue()
+
+        // if function is actually a getter, then use property name as method name
         val property = (irFunction as? IrFunctionImpl)?.correspondingPropertySymbol?.owner
         val methodName = intrinsicName ?: (property ?: irFunction).name.toString()
 
@@ -182,7 +184,7 @@ class Wrapper(val value: Any, private val type: IrClass) : Complex(type, mutable
             val fqName = this.getFqName()
             val owner = this.classOrNull?.owner
             return when {
-                owner?.hasAnnotation(evaluateIntrinsicAnnotation) == true -> Class.forName(owner.getEvaluateIntrinsicValue())
+                owner.hasAnnotation(evaluateIntrinsicAnnotation) -> Class.forName(owner!!.getEvaluateIntrinsicValue())
                 this.isPrimitiveType() -> getPrimitiveClass(fqName!!, asObject)
                 this.isArray() -> if (asObject) Array<Any?>::class.javaObjectType else Array<Any?>::class.java
                 //TODO primitive array
@@ -234,10 +236,31 @@ class Wrapper(val value: Any, private val type: IrClass) : Complex(type, mutable
     }
 
     override fun copy(): State {
-        return Wrapper(value, type)
+        return Wrapper(value, irClass)
     }
 
     override fun toString(): String {
         return "Wrapper(obj='$typeFqName', value=$value)"
     }
 }
+
+class Lambda(val irFunction: IrFunction, override var irClass: IrClass) : Complex(irClass, mutableListOf()) {
+    // irFunction is anonymous declaration, but irCall will contain descriptor of invoke method from Function interface
+    private val invokeDescriptor = irClass.declarations.first().descriptor
+
+    override fun setState(newVar: Variable) {
+        throw UnsupportedOperationException("Method setState is not supported in Lambda class")
+    }
+
+    override fun getIrFunction(descriptor: FunctionDescriptor): IrFunction? {
+        return if (invokeDescriptor.equalTo(descriptor)) irFunction else null
+    }
+
+    override fun copy(): State {
+        return Lambda(irFunction, irClass)
+    }
+
+    override fun toString(): String {
+        return "Lambda(${irClass.fqNameForIrSerialization})"
+    }
+}