Implement proxy for custom interfaces

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

@@ -16,10 +16,9 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrErrorExpressionImpl
 import org.jetbrains.kotlin.ir.interpreter.builtins.*
 import org.jetbrains.kotlin.ir.interpreter.exceptions.*
 import org.jetbrains.kotlin.ir.interpreter.intrinsics.IntrinsicEvaluator
-import org.jetbrains.kotlin.ir.interpreter.proxy.CommonProxy
+import org.jetbrains.kotlin.ir.interpreter.proxy.CommonProxy.Companion.asProxy
-import org.jetbrains.kotlin.ir.interpreter.proxy.LambdaProxy.Companion.asProxy
 import org.jetbrains.kotlin.ir.interpreter.proxy.Proxy
-import org.jetbrains.kotlin.ir.interpreter.proxy.unwrap
+import org.jetbrains.kotlin.ir.interpreter.proxy.wrap
 import org.jetbrains.kotlin.ir.interpreter.stack.StackImpl
 import org.jetbrains.kotlin.ir.interpreter.stack.Variable
 import org.jetbrains.kotlin.ir.interpreter.state.*
@@ -93,42 +92,17 @@ class IrInterpreter(private val irBuiltIns: IrBuiltIns, private val bodyMap: Map
         }
     }
 
-    internal fun IrFunction.interpret(valueArguments: List<Variable>): Any? {
+    internal fun IrFunction.interpret(valueArguments: List<Variable>, expectedResultClass: Class<*> = Any::class.java): Any? {
         val returnLabel = stack.newFrame(initPool = valueArguments) {
             this@interpret.interpret()
         }
         return when (returnLabel.returnLabel) {
-            ReturnLabel.REGULAR -> stack.popReturnValue().unwrap(this@IrInterpreter)
+            ReturnLabel.REGULAR -> stack.popReturnValue().wrap(this@IrInterpreter, expectedResultClass)
             ReturnLabel.EXCEPTION -> throw stack.popReturnValue() as ExceptionState
             else -> TODO("$returnLabel not supported as result of interpretation")
         }
     }
 
-    internal fun State.wrapAsProxyIfNeeded(calledFromBuiltIns: Boolean = false): Any? {
-        return when (this) {
-            is ExceptionState -> this
-            is Wrapper -> this.value
-            is Primitive<*> -> this.value
-            is Common -> CommonProxy(this, this@IrInterpreter, calledFromBuiltIns)
-            is Lambda -> this.asProxy(this@IrInterpreter)
-            else -> throw AssertionError("${this::class} is unsupported as argument for wrap function")
-        }
-    }
-
-    internal fun IrFunction.getArgsForMethodInvocation(args: List<Variable>): List<Any?> {
-        val argsValues = args.map { it.state.wrapAsProxyIfNeeded() }.toMutableList()
-
-        // TODO if vararg isn't last parameter
-        // must convert vararg array into separated elements for correct invoke
-        if (this.valueParameters.lastOrNull()?.varargElementType != null) {
-            val varargValue = argsValues.last()
-            argsValues.removeAt(argsValues.size - 1)
-            argsValues.addAll(varargValue as Array<out Any?>)
-        }
-
-        return argsValues
-    }
-
     private fun IrElement.interpret(): ExecutionResult {
         try {
             incrementAndCheckCommands()
@@ -192,7 +166,7 @@ class IrInterpreter(private val irBuiltIns: IrBuiltIns, private val bodyMap: Map
 
     private fun MethodHandle?.invokeMethod(irFunction: IrFunction): ExecutionResult {
         this ?: return handleIntrinsicMethods(irFunction)
-        val argsForMethodInvocation = irFunction.getArgsForMethodInvocation(stack.getAll())
+        val argsForMethodInvocation = irFunction.getArgsForMethodInvocation(this@IrInterpreter, this.type(), stack.getAll())
         val result = withExceptionHandler { this.invokeWithArguments(argsForMethodInvocation) }
         stack.pushReturnValue(result.toState(result.getType(irFunction.returnType)))
 
@@ -212,7 +186,7 @@ class IrInterpreter(private val irBuiltIns: IrBuiltIns, private val bodyMap: Map
 
         val receiverType = irFunction.dispatchReceiverParameter?.type
         val argsType = listOfNotNull(receiverType) + irFunction.valueParameters.map { it.type }
-        val argsValues = args.map { it.wrapAsProxyIfNeeded(methodName !in setOf("plus", IrBuiltIns.OperatorNames.EQEQ)) }
+        val argsValues = args.map { it.wrap(this, calledFromBuiltIns = methodName !in setOf("plus", IrBuiltIns.OperatorNames.EQEQ)) }
 
         fun IrType.getOnlyName(): String {
             return when {
@@ -734,7 +708,7 @@ class IrInterpreter(private val irBuiltIns: IrBuiltIns, private val bodyMap: Map
                 is Primitive<*> -> arrayToList(result.value)
                 is Common -> when {
                     result.irClass.defaultType.isUnsignedArray() -> arrayToList((result.fields.single().state as Primitive<*>).value)
-                    else -> listOf(CommonProxy(result, this))
+                    else -> listOf(result.asProxy(this))
                 }
                 else -> listOf(result)
             }

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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.ir.interpreter.exceptions.throwAsUserException
 import org.jetbrains.kotlin.ir.interpreter.stack.Variable
 import org.jetbrains.kotlin.ir.interpreter.state.*
 import org.jetbrains.kotlin.ir.interpreter.proxy.Proxy
+import org.jetbrains.kotlin.ir.interpreter.proxy.wrap
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
@@ -25,6 +26,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.capitalizeAsciiOnly
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import java.lang.invoke.MethodType
 
 internal fun IrFunction.getDispatchReceiver(): IrValueParameterSymbol? = this.dispatchReceiverParameter?.symbol
 
@@ -56,6 +58,9 @@ internal fun State.toIrExpression(expression: IrExpression): IrExpression {
     }
 }
 
+/**
+ * Convert object from outer world to state
+ */
 internal fun Any?.toState(irType: IrType): State {
     return when (this) {
         is Proxy -> this.state
@@ -225,4 +230,20 @@ inline fun withExceptionHandler(block: () -> Any?): Any? {
     } catch (e: Throwable) {
         e.throwAsUserException()
     }
+}
+
+internal fun IrFunction.getArgsForMethodInvocation(interpreter: IrInterpreter, methodType: MethodType, args: List<Variable>): List<Any?> {
+    val argsValues = args
+        .mapIndexed { index, variable -> variable.state.wrap(interpreter, methodType.parameterType(index)) }
+        .toMutableList()
+
+    // TODO if vararg isn't last parameter
+    // must convert vararg array into separated elements for correct invoke
+    if (this.valueParameters.lastOrNull()?.varargElementType != null) {
+        val varargValue = argsValues.last()
+        argsValues.removeAt(argsValues.size - 1)
+        argsValues.addAll(varargValue as Array<out Any?>)
+    }
+
+    return argsValues
 }

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/CommonProxy.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.ir.interpreter.getDispatchReceiver
 import org.jetbrains.kotlin.ir.interpreter.internalName
 import org.jetbrains.kotlin.ir.interpreter.stack.Variable
 import org.jetbrains.kotlin.ir.interpreter.state.Common
+import org.jetbrains.kotlin.ir.interpreter.toState
 import org.jetbrains.kotlin.ir.util.isFakeOverriddenFromAny
 
 /**
@@ -18,12 +19,12 @@ import org.jetbrains.kotlin.ir.util.isFakeOverriddenFromAny
  *         return super.toString()
  *     }
  */
-internal class CommonProxy(override val state: Common, override val interpreter: IrInterpreter, private val calledFromBuiltIns: Boolean = false): Proxy {
+internal class CommonProxy private constructor(
+    override val state: Common, override val interpreter: IrInterpreter, private val calledFromBuiltIns: Boolean = false
+) : Proxy {
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
-        if (javaClass != other?.javaClass) return false
+        if (other !is Proxy) return false
-
-        other as CommonProxy
 
         val valueArguments = mutableListOf<Variable>()
         val equalsFun = state.getEqualsFunction()
@@ -54,4 +55,34 @@ internal class CommonProxy(override val state: Common, override val interpreter:
         toStringFun.getDispatchReceiver()!!.let { valueArguments.add(Variable(it, state)) }
         return with(interpreter) { toStringFun.interpret(valueArguments) } as String
     }
+
+    companion object {
+        internal fun Common.asProxy(interpreter: IrInterpreter, extendFrom: Class<*>? = null, calledFromBuiltIns: Boolean = false): Any {
+            val commonProxy = CommonProxy(this, interpreter, calledFromBuiltIns)
+
+            val interfaces = when (extendFrom) {
+                null, Object::class.java -> arrayOf(Proxy::class.java)
+                else -> arrayOf(extendFrom, Proxy::class.java)
+            }
+            return java.lang.reflect.Proxy.newProxyInstance(ClassLoader.getSystemClassLoader(), interfaces)
+            { /*proxy*/_, method, args ->
+                when {
+                    method.declaringClass == Proxy::class.java && method.name == "getState" -> commonProxy.state
+                    method.declaringClass == Proxy::class.java && method.name == "getInterpreter" -> commonProxy.interpreter
+                    method.name == "equals" && method.parameterTypes.single().isObject() -> commonProxy.equals(args.single())
+                    method.name == "hashCode" && method.parameterTypes.isEmpty() -> commonProxy.hashCode()
+                    method.name == "toString" && method.parameterTypes.isEmpty() -> commonProxy.toString()
+                    else -> {
+                        val irFunction = commonProxy.state.getIrFunction(method)
+                            ?: throw AssertionError("Cannot find method $method in ${commonProxy.state}")
+                        val valueArguments = mutableListOf<Variable>()
+                        valueArguments += Variable(irFunction.getDispatchReceiver()!!, commonProxy.state)
+                        valueArguments += irFunction.valueParameters
+                            .mapIndexed { index, parameter -> Variable(parameter.symbol, args[index].toState(parameter.type)) }
+                        with(interpreter) { irFunction.interpret(valueArguments, method.returnType) }
+                    }
+                }
+            }
+        }
+    }
 }

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/LambdaProxy.kt

@@ -11,12 +11,12 @@ import org.jetbrains.kotlin.ir.interpreter.stack.Variable
 import org.jetbrains.kotlin.ir.interpreter.state.Lambda
 import org.jetbrains.kotlin.ir.interpreter.toState
 
-internal class LambdaProxy(override val state: Lambda, override val interpreter: IrInterpreter): Proxy {
+internal class LambdaProxy private constructor(
+    override val state: Lambda, override val interpreter: IrInterpreter
+) : Proxy {
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
-        if (javaClass != other?.javaClass) return false
+        if (other !is Proxy) return false
-
-        other as LambdaProxy
 
         return state == other.state
     }
@@ -41,17 +41,19 @@ internal class LambdaProxy(override val state: Lambda, override val interpreter:
                 this.isKFunction -> Class.forName("kotlin.reflect.KFunction")
                 else -> throw InternalError("Cannot handle lambda $this as proxy")
             }
-            return java.lang.reflect.Proxy.newProxyInstance(ClassLoader.getSystemClassLoader(), arrayOf(functionClass))
+            return java.lang.reflect.Proxy.newProxyInstance(ClassLoader.getSystemClassLoader(), arrayOf(functionClass, Proxy::class.java))
             { proxy, method, args ->
                 when {
-                    method.name == "invoke" && method.parameterTypes.single().isObject() -> {
+                    method.declaringClass == Proxy::class.java && method.name == "getState" -> lambdaProxy.state
-                        val valueArguments = this.irFunction.valueParameters
+                    method.declaringClass == Proxy::class.java && method.name == "getInterpreter" -> lambdaProxy.interpreter
-                            .mapIndexed { index, parameter -> Variable(parameter.symbol, args[index].toState(parameter.type)) }
-                        with(interpreter) { lambdaProxy.state.irFunction.interpret(valueArguments) }
-                    }
                     method.name == "equals" && method.parameterTypes.single().isObject() -> lambdaProxy.equals(args.single())
                     method.name == "hashCode" && method.parameterTypes.isEmpty() -> lambdaProxy.hashCode()
                     method.name == "toString" && method.parameterTypes.isEmpty() -> lambdaProxy.toString()
+                    method.name == "invoke" && method.parameterTypes.single().isObject() -> {
+                        val valueArguments = this.irFunction.valueParameters
+                            .mapIndexed { index, parameter -> Variable(parameter.symbol, args[index].toState(parameter.type)) }
+                        with(interpreter) { lambdaProxy.state.irFunction.interpret(valueArguments, method.returnType) }
+                    }
                     else -> throw InternalError("Cannot invoke method ${method.name} from lambda $this")
                 }
             }

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/Proxy.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.ir.interpreter.proxy
 
 import org.jetbrains.kotlin.ir.interpreter.IrInterpreter
+import org.jetbrains.kotlin.ir.interpreter.proxy.CommonProxy.Companion.asProxy
 import org.jetbrains.kotlin.ir.interpreter.proxy.LambdaProxy.Companion.asProxy
 import org.jetbrains.kotlin.ir.interpreter.state.*
 
@@ -18,14 +19,17 @@ internal interface Proxy {
     override fun toString(): String
 }
 
-internal fun State.unwrap(interpreter: IrInterpreter): Any? {
+/**
+ * Prepare state object to be passed in outer world
+ */
+internal fun State.wrap(interpreter: IrInterpreter, extendFrom: Class<*>? = null, calledFromBuiltIns: Boolean = false): Any? {
     return when (this) {
-        is Proxy -> this.state
+        is ExceptionState -> this
-        is Primitive<*> -> this.value
-        is Common -> CommonProxy(this, interpreter)
-        is Lambda -> this.asProxy(interpreter)
         is Wrapper -> this.value
-        else -> throw AssertionError("Unsupported state for proxy unwrap: ${this::class.java}")
+        is Primitive<*> -> this.value
+        is Common -> this.asProxy(interpreter, extendFrom, calledFromBuiltIns)
+        is Lambda -> this.asProxy(interpreter)
+        else -> throw AssertionError("${this::class} is unsupported as argument for wrap function")
     }
 }
 

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

@@ -6,8 +6,13 @@
 package org.jetbrains.kotlin.ir.interpreter.state
 
 import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrProperty
+import org.jetbrains.kotlin.ir.interpreter.getLastOverridden
 import org.jetbrains.kotlin.ir.interpreter.stack.Variable
 import org.jetbrains.kotlin.ir.util.fqNameForIrSerialization
+import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
+import org.jetbrains.kotlin.ir.util.nameForIrSerialization
 
 internal class Common private constructor(override val irClass: IrClass, override val fields: MutableList<Variable>) : Complex {
     override var superWrapperClass: Wrapper? = null
@@ -16,12 +21,32 @@ internal class Common private constructor(override val irClass: IrClass, overrid
 
     constructor(irClass: IrClass) : this(irClass, mutableListOf())
 
-    override fun toString(): String {
-        return "Common(obj='${irClass.fqNameForIrSerialization}', values=$fields)"
-    }
-
     fun copyFieldsFrom(state: Complex) {
         this.fields.addAll(state.fields)
         superWrapperClass = state.superWrapperClass ?: state as? Wrapper
     }
+
+    // This method is used to get correct java method name
+    private fun getKotlinName(declaringClassName: String, methodName: String): String {
+        return when {
+            // TODO see specialBuiltinMembers.kt
+            //"kotlin.collections.Map.<get-entries>" -> "entrySet"
+            //"kotlin.collections.Map.<get-keys>" -> "keySet"
+            declaringClassName == "java.lang.CharSequence" && methodName == "charAt" -> "get"
+            //"kotlin.collections.MutableList.removeAt" -> "remove"
+            else -> methodName
+        }
+    }
+
+    fun getIrFunction(method: java.lang.reflect.Method): IrFunction? {
+        val methodName = getKotlinName(method.declaringClass.name, method.name)
+        return when (val declaration = irClass.declarations.singleOrNull { it.nameForIrSerialization.asString() == methodName }) {
+            is IrProperty -> declaration.getter
+            else -> declaration as? IrFunction
+        }
+    }
+
+    override fun toString(): String {
+        return "Common(obj='${irClass.fqNameForIrSerialization}', values=$fields)"
+    }
 }