From 668bb4fd713925a072ba3562984bccd5f8e24cc5 Mon Sep 17 00:00:00 2001
From: Ivan Kylchik <ivan.kylchik@jetbrains.com>
Date: Wed, 23 Jun 2021 14:23:39 +0300
Subject: [PATCH] Fix equals check in interpreter in case when object is
 wrapped as Proxy

---
 .../kotlin/ir/interpreter/proxy/CommonProxy.kt      |  7 ++++---
 .../jetbrains/kotlin/ir/interpreter/proxy/Proxy.kt  |  8 ++++++++
 .../testData/ir/interpreter/proxy/customEquals.kt   | 13 +++++++++++++
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/CommonProxy.kt b/compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/CommonProxy.kt
index f92ced1b818..17cf057cbc0 100644
--- a/compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/CommonProxy.kt
+++ b/compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/CommonProxy.kt
@@ -11,11 +11,12 @@ import org.jetbrains.kotlin.ir.interpreter.CallInterceptor
 import org.jetbrains.kotlin.ir.interpreter.getDispatchReceiver
 import org.jetbrains.kotlin.ir.interpreter.stack.Variable
 import org.jetbrains.kotlin.ir.interpreter.state.Common
+import org.jetbrains.kotlin.ir.interpreter.state.State
 import org.jetbrains.kotlin.ir.interpreter.toState
 import org.jetbrains.kotlin.ir.util.isFakeOverriddenFromAny
 
 internal class CommonProxy private constructor(override val state: Common, override val callInterceptor: CallInterceptor) : Proxy {
-    private fun defaultEquals(other: Proxy): Boolean = this.state === other.state
+    private fun defaultEquals(other: Any?): Boolean = if (other is Proxy) this.state === other.state else false
     private fun defaultHashCode(): Int = System.identityHashCode(state)
     private fun defaultToString(): String = "${state.irClass.internalName()}@" + hashCode().toString(16).padStart(8, '0')
 
@@ -32,14 +33,14 @@ internal class CommonProxy private constructor(override val state: Common, overr
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
-        if (other !is Proxy) return false
+        if (other == null) return false
 
         val valueArguments = mutableListOf<Variable>()
         val equalsFun = state.getEqualsFunction()
         if (equalsFun.isFakeOverriddenFromAny() || equalsFun.wasAlreadyCalled()) return defaultEquals(other)
 
         equalsFun.getDispatchReceiver()!!.let { valueArguments.add(Variable(it, state)) }
-        valueArguments.add(Variable(equalsFun.valueParameters.single().symbol, other.state))
+        valueArguments.add(Variable(equalsFun.valueParameters.single().symbol, if (other is Proxy) other.state else other as State))
 
         return callInterceptor.interceptProxy(equalsFun, valueArguments) as Boolean
     }
diff --git a/compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/Proxy.kt b/compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/Proxy.kt
index 16025376a16..39a06b0f554 100644
--- a/compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/Proxy.kt
+++ b/compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/proxy/Proxy.kt
@@ -45,6 +45,14 @@ internal fun State.wrap(callInterceptor: CallInterceptor, remainArraysAsIs: Bool
  */
 internal fun List<State>.wrap(callInterceptor: CallInterceptor, irFunction: IrFunction, methodType: MethodType? = null): List<Any?> {
     val name = irFunction.fqNameWhenAvailable?.asString() ?: ""
+    if (name == "kotlin.internal.ir.EQEQ" && this.any { it is Common }) {
+        // in case of custom `equals` it is important not to lose information obout type
+        // so all states remain as is, only common will be converted to Proxy
+        return mapIndexed { index, state ->
+            if (state !is Common) return@mapIndexed state
+            state.wrap(callInterceptor, remainArraysAsIs = true, methodType?.parameterType(index))
+        }
+    }
     return this.mapIndexed { index, state ->
         // don't get arrays from Primitive in case of "set" and "Pair.<init>"; information about type will be lost
         val unwrapArrays = (name == "kotlin.Array.set" && index != 0) || name == "kotlin.Pair.<init>" || name == "kotlin.internal.ir.CHECK_NOT_NULL"
diff --git a/compiler/testData/ir/interpreter/proxy/customEquals.kt b/compiler/testData/ir/interpreter/proxy/customEquals.kt
index 62a0a0c1d37..b1b73b87bc6 100644
--- a/compiler/testData/ir/interpreter/proxy/customEquals.kt
+++ b/compiler/testData/ir/interpreter/proxy/customEquals.kt
@@ -1,5 +1,18 @@
 import kotlin.collections.*
 
+@CompileTimeCalculation
+class A(val a: Int) {
+    override fun equals(other: Any?): Boolean {
+        return other is Int && other == a
+    }
+}
+const val customEquals1 = <!EVALUATED: `true`!>A(1) == 1<!>
+const val customEquals2 = <!EVALUATED: `false`!>A(1) == 123<!>
+const val customEquals3 = <!EVALUATED: `false`!>1 == A(1)<!>
+const val customEquals4 = <!EVALUATED: `false`!>123 == A(1)<!>
+const val customEquals5 = <!EVALUATED: `false`!>null == A(1)<!>
+const val customEquals6 = <!EVALUATED: `false`!>A(1) == null<!>
+
 @CompileTimeCalculation
 class B(val b: Int) {
     override fun equals(other: Any?): Boolean {