Simplify the interpretation logic of methods defined in super Wrapper

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

@@ -42,8 +42,8 @@ internal interface CallInterceptor {
     val interpreter: IrInterpreter
 
     fun interceptProxy(irFunction: IrFunction, valueArguments: List<Variable>, expectedResultClass: Class<*> = Any::class.java): Any?
-    fun interceptCall(call: IrCall, irFunction: IrFunction, receiver: State?, args: List<State>, defaultAction: () -> Unit)
+    fun interceptCall(call: IrCall, irFunction: IrFunction, args: List<State>, defaultAction: () -> Unit)
-    fun interceptConstructor(constructorCall: IrFunctionAccessExpression, receiver: State, args: List<State>, defaultAction: () -> Unit)
+    fun interceptConstructor(constructorCall: IrFunctionAccessExpression, args: List<State>, defaultAction: () -> Unit)
     fun interceptGetObjectValue(expression: IrGetObjectValue, defaultAction: () -> Unit)
     fun interceptEnumEntry(enumEntry: IrEnumEntry, defaultAction: () -> Unit)
     fun interceptJavaStaticField(expression: IrGetField)
@@ -65,8 +65,9 @@ internal class DefaultCallInterceptor(override val interpreter: IrInterpreter) :
         }.wrap(this@DefaultCallInterceptor, remainArraysAsIs = true, extendFrom = expectedResultClass)
     }
 
-    override fun interceptCall(call: IrCall, irFunction: IrFunction, receiver: State?, args: List<State>, defaultAction: () -> Unit) {
+    override fun interceptCall(call: IrCall, irFunction: IrFunction, args: List<State>, defaultAction: () -> Unit) {
         val isInlineOnly = irFunction.hasAnnotation(FqName("kotlin.internal.InlineOnly"))
+        val receiver = if (irFunction.dispatchReceiverParameter != null) args[0] else null
         when {
             receiver is Wrapper && !isInlineOnly -> receiver.getMethod(irFunction).invokeMethod(irFunction, args)
             Wrapper.mustBeHandledWithWrapper(irFunction) -> Wrapper.getStaticMethod(irFunction).invokeMethod(irFunction, args)
@@ -123,14 +124,13 @@ internal class DefaultCallInterceptor(override val interpreter: IrInterpreter) :
         }
 
         if (invokedFunction.symbol is IrConstructorSymbol) return
-        this.interceptCall(call, function, dispatchReceiver, listOfNotNull(dispatchReceiver, extensionReceiver) + valueArguments) {
+        this.interceptCall(call, function, listOfNotNull(dispatchReceiver, extensionReceiver) + valueArguments) {
             callStack.addInstruction(CompoundInstruction(function))
         }
     }
 
-    override fun interceptConstructor(
+    override fun interceptConstructor(constructorCall: IrFunctionAccessExpression, args: List<State>, defaultAction: () -> Unit) {
-        constructorCall: IrFunctionAccessExpression, receiver: State, args: List<State>, defaultAction: () -> Unit
+        val receiver = callStack.getState(constructorCall.getThisReceiver())
-    ) {
         val irConstructor = constructorCall.symbol.owner
         val irClass = irConstructor.parentAsClass
         when {

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

@@ -170,18 +170,20 @@ class IrInterpreter private constructor(
     }
 
     private fun interpretCall(call: IrCall) {
+        fun IrFunction.mustBeHandledWithSuperOf(state: State?): State? {
+            if (state is Common && this.parent == state.superWrapperClass?.irClass) return state.superWrapperClass
+            return state
+        }
+
         val owner = call.symbol.owner
         // 1. load evaluated arguments from stack
-        var dispatchReceiver = owner.getDispatchReceiver()?.let { callStack.getState(it) }
+        val dispatchReceiver = owner.getDispatchReceiver()?.let { callStack.getState(it) }
         val extensionReceiver = owner.getExtensionReceiver()?.let { callStack.getState(it) }
-        val valueArguments = owner.valueParameters.map { callStack.getState(it.symbol) }
+        val valueArguments = owner.valueParameters.map { callStack.getState(it.symbol) }.toMutableList()
 
         // 2. get correct function for interpretation
         val irFunction = dispatchReceiver?.getIrFunctionByIrCall(call) ?: call.symbol.owner
-        dispatchReceiver = when (irFunction.parent) {
+        val args = listOfNotNull(irFunction.mustBeHandledWithSuperOf(dispatchReceiver), extensionReceiver) + valueArguments
-            (dispatchReceiver as? Complex)?.superWrapperClass?.irClass -> dispatchReceiver.superWrapperClass
-            else -> dispatchReceiver
-        }
 
         callStack.dropSubFrame() // drop intermediate frame that contains variables for default arg evaluation
         callStack.newFrame(irFunction)
@@ -191,13 +193,11 @@ class IrInterpreter private constructor(
         callStack.addVariable(Variable(irFunction.symbol, KTypeState(call.type, irBuiltIns.anyClass.owner)))
 
         // 3. store arguments in memory (remap args on actual names)
-        val args = mutableListOf<Variable>()
+        val variables = mutableListOf<Variable>()
-        irFunction.getDispatchReceiver()?.let { dispatchReceiver?.let { receiver -> args.add(Variable(it, receiver)) } }
+        irFunction.getDispatchReceiver()?.let { dispatchReceiver?.let { receiver -> variables.add(Variable(it, receiver)) } }
-        irFunction.getExtensionReceiver()?.let { args.add(Variable(it, extensionReceiver ?: valueArguments.first())) }
+        irFunction.getExtensionReceiver()?.let { variables.add(Variable(it, extensionReceiver ?: valueArguments.removeFirst())) }
-        // `shift` is used when extension receiver is actually a parameter in lambda
+        irFunction.valueParameters.forEach { variables.add(Variable(it.symbol, valueArguments.removeFirst())) }
-        val shift = if (irFunction.extensionReceiverParameter != null && extensionReceiver == null) 1 else 0
+        variables.forEach { callStack.addVariable(it) }
-        irFunction.valueParameters.forEachIndexed { i, param -> args.add(Variable(param.symbol, valueArguments[i + shift])) }
-        args.forEach { callStack.addVariable(it) }
 
         // 4. store reified type parameters
         irFunction.typeParameters.filter { it.isReified }
@@ -213,7 +213,7 @@ class IrInterpreter private constructor(
             generateSequence(dispatchReceiver.outerClass) { (it.state as? Complex)?.outerClass }.forEach { callStack.addVariable(it) }
         }
 
-        callInterceptor.interceptCall(call, irFunction, dispatchReceiver, listOfNotNull(dispatchReceiver, extensionReceiver) + valueArguments) {
+        callInterceptor.interceptCall(call, irFunction, args) {
             callStack.addInstruction(CompoundInstruction(irFunction))
         }
     }
@@ -289,7 +289,7 @@ class IrInterpreter private constructor(
         }
         superReceiver?.let { callStack.addVariable(Variable(it, objectState)) }
 
-        callInterceptor.interceptConstructor(constructorCall, objectState, valueArguments) {
+        callInterceptor.interceptConstructor(constructorCall, valueArguments) {
             callStack.pushState(objectState)
             callStack.addInstruction(CompoundInstruction(constructor))
         }

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/state/Complex.kt

@@ -38,14 +38,6 @@ internal interface Complex: State {
         }
     }
 
-    private fun getThisOrSuperReceiver(superIrClass: IrClass?): IrClass? {
-        return when {
-            superIrClass == null -> this.irClass
-            superIrClass.isInterface -> superIrClass
-            else -> irClass.superTypes.map { it.classOrNull?.owner }.singleOrNull { it?.isInterface == false }
-        }
-    }
-
     fun getOverridden(owner: IrSimpleFunction): IrSimpleFunction {
         if (owner.parent == superWrapperClass?.irClass) return owner
         if (!owner.isFakeOverride || owner.body != null || owner.parentAsClass.defaultType.isAny()) return owner
@@ -55,7 +47,7 @@ internal interface Complex: State {
     }
 
     override fun getIrFunctionByIrCall(expression: IrCall): IrFunction? {
-        val receiver = getThisOrSuperReceiver(expression.superQualifierSymbol?.owner) ?: return null
+        val receiver = expression.superQualifierSymbol?.owner ?: irClass
         val irFunction = getIrFunctionFromGivenClass(receiver, expression.symbol) ?: return null
         return getOverridden(irFunction as IrSimpleFunction)
     }

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/state/Wrapper.kt

@@ -34,12 +34,12 @@ internal class Wrapper(val value: Any, override val irClass: IrClass) : Complex
         when {
             javaClass == HashMap::class.java -> {
                 val nodeClass = javaClass.declaredClasses.single { it.name.contains("\$Node") }
-                val mutableMap = irClass.superTypes.mapNotNull { it.classOrNull?.owner }.single { it.isInterface }
+                val mutableMap = irClass.superTypes.mapNotNull { it.classOrNull?.owner }.single { it.name.asString() == "MutableMap" }
                 javaClassToIrClass += nodeClass to mutableMap.declarations.filterIsInstance<IrClass>().single()
             }
             javaClass == LinkedHashMap::class.java -> {
                 val entryClass = javaClass.declaredClasses.single { it.name.contains("\$Entry") }
-                val mutableMap = irClass.superTypes.mapNotNull { it.classOrNull?.owner }.single { it.isInterface }
+                val mutableMap = irClass.superTypes.mapNotNull { it.classOrNull?.owner }.single { it.name.asString() == "MutableMap" }
                 javaClassToIrClass += entryClass to mutableMap.declarations.filterIsInstance<IrClass>().single()
             }
             javaClass.canonicalName == "java.util.Collections.SingletonMap" -> {
@@ -92,7 +92,9 @@ internal class Wrapper(val value: Any, override val irClass: IrClass) : Complex
 
         // TODO remove later; used for tests only
         private val intrinsicClasses = setOf(
-            "kotlin.text.StringBuilder", "kotlin.Pair", "kotlin.collections.HashMap",
+            "kotlin.text.StringBuilder", "kotlin.Pair", "kotlin.collections.ArrayList",
+            "kotlin.collections.HashMap", "kotlin.collections.LinkedHashMap",
+            "kotlin.collections.HashSet", "kotlin.collections.LinkedHashSet",
             "kotlin.text.RegexOption", "kotlin.text.Regex", "kotlin.text.Regex.Companion", "kotlin.text.MatchGroup",
         )
 
@@ -102,6 +104,7 @@ internal class Wrapper(val value: Any, override val irClass: IrClass) : Complex
             "kotlin.collections.arrayListOf(Array)" to "kotlin.collections.CollectionsKt",
             "Char.kotlin.text.isWhitespace()" to "kotlin.text.CharsKt",
             "Array.kotlin.collections.toMutableList()" to "kotlin.collections.ArraysKt",
+            "Array.kotlin.collections.copyToArrayOfAny(Boolean)" to "kotlin.collections.CollectionsKt",
         )
 
         fun associateJavaClassWithIrClass(javaClass: Class<*>, irClass: IrClass) {
@@ -240,6 +243,13 @@ internal class Wrapper(val value: Any, override val irClass: IrClass) : Complex
                 fqName == "kotlin.collections.ListIterator" || fqName == "kotlin.collections.MutableListIterator" -> ListIterator::class.java
                 fqName == "kotlin.collections.Iterator" || fqName == "kotlin.collections.MutableIterator" -> Iterator::class.java
                 fqName == "kotlin.collections.Map.Entry" || fqName == "kotlin.collections.MutableMap.MutableEntry" -> Map.Entry::class.java
+                fqName == "kotlin.collections.ArrayList" -> ArrayList::class.java
+                fqName == "kotlin.collections.HashMap" -> HashMap::class.java
+                fqName == "kotlin.collections.HashSet" -> HashSet::class.java
+                fqName == "kotlin.collections.LinkedHashMap" -> LinkedHashMap::class.java
+                fqName == "kotlin.collections.LinkedHashSet" -> LinkedHashSet::class.java
+                fqName == "kotlin.text.StringBuilder" -> StringBuilder::class.java
+                fqName == "kotlin.text.Appendable" -> Appendable::class.java
                 fqName == null -> Any::class.java // null if this.isTypeParameter()
                 else -> Class.forName(owner.internalName())
             }

compiler/testData/ir/interpreter/proxy/superWrapper.kt

@@ -1,7 +1,7 @@
 import kotlin.collections.*
 
 @CompileTimeCalculation
-class MyArrayList<E>: ArrayList<E>() {
+open class MyArrayList<E>: ArrayList<E>() {
     var addCounter = 0
     override fun add(element: E): Boolean {
         addCounter++
@@ -10,12 +10,22 @@ class MyArrayList<E>: ArrayList<E>() {
 }
 
 @CompileTimeCalculation
-fun test(): String {
+class MyOtherArrayList<E>: MyArrayList<E>() {
-    val list = MyArrayList<Int>()
+    override fun addAll(elements: Collection<E>): Boolean {
+        return true // do nothing
+    }
+}
+
+@CompileTimeCalculation
+fun test(list: MyArrayList<Int>): String {
     list.add(1)
     list.add(2)
     list.add(3)
+
+    val otherList = arrayListOf(4, 5, 6)
+    list.addAll(otherList)
     return "Counter " + list.addCounter + "; size " + list.size
 }
 
-const val testResult = <!EVALUATED: `Counter 3; size 3`!>test()<!>
+const val testResult1 = <!EVALUATED: `Counter 3; size 6`!>test(MyArrayList<Int>())<!>
+const val testResult2 = <!EVALUATED: `Counter 3; size 3`!>test(MyOtherArrayList<Int>())<!>