[IR] Commonized SAM equals/hashCode builder

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/Ir.kt

@@ -308,6 +308,8 @@ abstract class Symbols<out T : CommonBackendContext>(val context: T, irBuiltIns:
 
     abstract val returnIfSuspended: IrSimpleFunctionSymbol
 
+    abstract val functionAdapter: IrClassSymbol
+
     open val unsafeCoerceIntrinsic: IrSimpleFunctionSymbol? = null
 
     companion object {

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SingleAbstractMethodLowering.kt

@@ -21,19 +21,18 @@ import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrGetValue
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
 import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
-import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.types.classifierOrFail
-import org.jetbrains.kotlin.ir.types.isNullable
-import org.jetbrains.kotlin.ir.types.makeNullable
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.ir.util.functions
 
 abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
     // SAM wrappers are cached, either in the file class (if it exists), or in a top-level enclosing class.
@@ -67,6 +66,8 @@ abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) :
 
     abstract fun getSuperTypeForWrapper(typeOperand: IrType): IrType
 
+    abstract val IrType.needEqualsHashCodeMethods: Boolean
+
     open val inInlineFunctionScope get() = allScopes.any { scope -> (scope.irElement as? IrFunction)?.isInline ?: false }
 
     override fun lower(irFile: IrFile) {
@@ -221,14 +222,117 @@ abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) :
             }
         }
 
-        generateEqualsHashCode(subclass, superType, field)
+        if (superType.needEqualsHashCodeMethods)
+            generateEqualsHashCode(subclass, superType, field)
 
         subclass.addFakeOverridesViaIncorrectHeuristic()
 
         return subclass
     }
 
-    protected open fun getAdditionalSupertypes(supertype: IrType): List<IrType> = emptyList()
+    private fun generateEqualsHashCode(klass: IrClass, superType: IrType, functionDelegateField: IrField) =
+        SamEqualsHashCodeMethodsGenerator(context, klass, superType) { receiver ->
+            irGetField(receiver, functionDelegateField)
+        }.generate()
 
-    protected open fun generateEqualsHashCode(klass: IrClass, supertype: IrType, functionDelegateField: IrField) {}
+    private fun getAdditionalSupertypes(supertype: IrType) =
+        if (supertype.needEqualsHashCodeMethods)
+            listOf(context.ir.symbols.functionAdapter.typeWith())
+        else emptyList()
+}
+
+/**
+ * Generates equals and hashCode for SAM and fun interface wrappers, as well as an implementation of getFunctionDelegate
+ * (inherited from kotlin.jvm.internal.FunctionAdapter), needed to properly implement them.
+ * This class is used in two places:
+ * - FunctionReferenceLowering, which is the case of SAM conversion of a (maybe adapted) function reference, e.g. `FunInterface(foo::bar)`.
+ *   Note that we don't generate equals/hashCode for SAM conversion of lambdas, e.g. `FunInterface {}`, even though lambdas are represented
+ *   as a local function + reference to it. The reason for this is that all lambdas are unique, so after SAM conversion they are still
+ *   never equal to each other. See [FunctionReferenceLowering.FunctionReferenceBuilder.needToGenerateSamEqualsHashCodeMethods].
+ * - SingleAbstractMethodLowering, which is the case of SAM conversion of any value of a functional type,
+ *   e.g. `val f = {}; FunInterface(f)`.
+ */
+class SamEqualsHashCodeMethodsGenerator(
+    private val context: CommonBackendContext,
+    private val klass: IrClass,
+    private val samSuperType: IrType,
+    private val obtainFunctionDelegate: IrBuilderWithScope.(receiver: IrExpression) -> IrExpression,
+) {
+    private val functionAdapterClass = context.ir.symbols.functionAdapter.owner
+
+    private val builtIns: IrBuiltIns get() = context.irBuiltIns
+    private val getFunctionDelegate = functionAdapterClass.functions.single { it.name.asString() == "getFunctionDelegate" }
+
+    fun generate() {
+        generateGetFunctionDelegate()
+        generateEquals()
+        generateHashCode()
+    }
+
+    private fun generateGetFunctionDelegate() {
+        klass.addFunction(getFunctionDelegate.name.asString(), getFunctionDelegate.returnType).apply {
+            overriddenSymbols = listOf(getFunctionDelegate.symbol)
+            body = context.createIrBuilder(symbol).run {
+                irExprBody(obtainFunctionDelegate(irGet(dispatchReceiverParameter!!)))
+            }
+        }
+    }
+
+    private fun generateEquals() {
+        klass.addFunction("equals", builtIns.booleanType).apply {
+            overriddenSymbols = klass.superTypes.mapNotNull {
+                it.getClass()?.functions?.singleOrNull {
+                    it.name.asString() == "equals" &&
+                            it.extensionReceiverParameter == null &&
+                            it.valueParameters.singleOrNull()?.type == builtIns.anyNType
+                }?.symbol
+            }
+
+            val other = addValueParameter("other", builtIns.anyNType)
+            body = context.createIrBuilder(symbol).run {
+                irExprBody(
+                    irIfThenElse(
+                        builtIns.booleanType,
+                        irIs(irGet(other), samSuperType),
+                        irIfThenElse(
+                            builtIns.booleanType,
+                            irIs(irGet(other), functionAdapterClass.typeWith()),
+                            irEquals(
+                                irCall(getFunctionDelegate).also {
+                                    it.dispatchReceiver = irGet(dispatchReceiverParameter!!)
+                                },
+                                irCall(getFunctionDelegate).also {
+                                    it.dispatchReceiver = irImplicitCast(irGet(other), functionAdapterClass.typeWith())
+                                }
+                            ),
+                            irFalse()
+                        ),
+                        irFalse()
+                    )
+                )
+            }
+        }
+    }
+
+    private fun generateHashCode() {
+        klass.addFunction("hashCode", builtIns.intType).apply {
+
+            fun isHashCode(function: IrSimpleFunction) =
+                function.name.asString() == "hashCode" && function.extensionReceiverParameter == null && function.valueParameters.isEmpty()
+
+            overriddenSymbols = klass.superTypes.mapNotNull {
+                it.getClass()?.functions?.singleOrNull(::isHashCode)?.symbol
+            }
+            val hashCode = builtIns.anyClass.owner.functions.single(::isHashCode).symbol
+            body = context.createIrBuilder(symbol).run {
+                irExprBody(
+                    irCall(hashCode).also {
+                        it.dispatchReceiver = irCall(getFunctionDelegate).also {
+                            it.dispatchReceiver = irGet(dispatchReceiverParameter!!)
+                        }
+                    }
+                )
+            }
+        }
+    }
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt

@@ -215,6 +215,9 @@ class JsIrBackendContext(
             override val coroutineGetContext = symbolTable.referenceSimpleFunction(getJsInternalFunction(GET_COROUTINE_CONTEXT_NAME))
 
             override val returnIfSuspended = symbolTable.referenceSimpleFunction(getJsInternalFunction("returnIfSuspended"))
+
+            override val functionAdapter: IrClassSymbol
+                get() = TODO("Not implemented")
         }
 
         override fun unfoldInlineClassType(irType: IrType): IrType? {

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/JsSingleAbstractMethodLowering.kt

@@ -28,6 +28,8 @@ class JsSingleAbstractMethodLowering(context: JsIrBackendContext) : SingleAbstra
         return Visibilities.PRIVATE
     }
 
+    override val IrType.needEqualsHashCodeMethods get() = false
+
     private var enclosingBodyContainer: IrDeclaration? = null
 
     override fun lower(irFile: IrFile) {

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmSymbols.kt

@@ -502,7 +502,7 @@ class JvmSymbols(
         }
     }
 
-    val functionAdapter: IrClassSymbol = createClass(FqName("kotlin.jvm.internal.FunctionAdapter"), ClassKind.INTERFACE) { klass ->
+    override val functionAdapter: IrClassSymbol = createClass(FqName("kotlin.jvm.internal.FunctionAdapter"), ClassKind.INTERFACE) { klass ->
         klass.addFunction("getFunctionDelegate", irBuiltIns.functionClass.starProjectedType, Modality.ABSTRACT)
     }
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/FunctionReferenceLowering.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.backend.jvm.lower
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.ir.*
+import org.jetbrains.kotlin.backend.common.lower.SamEqualsHashCodeMethodsGenerator
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/JvmSingleAbstractMethodLowering.kt

@@ -7,27 +7,15 @@ package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.ScopeWithIr
 import org.jetbrains.kotlin.backend.common.lower.SingleAbstractMethodLowering
-import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.ir.erasedUpperBound
 import org.jetbrains.kotlin.descriptors.Visibilities
-import org.jetbrains.kotlin.descriptors.Visibility
-import org.jetbrains.kotlin.ir.builders.*
-import org.jetbrains.kotlin.ir.builders.declarations.addFunction
-import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
-import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
-import org.jetbrains.kotlin.ir.declarations.IrField
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
-import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.getClass
-import org.jetbrains.kotlin.ir.types.typeWith
 import org.jetbrains.kotlin.ir.util.defaultType
-import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.load.java.JavaVisibilities
 
 internal val singleAbstractMethodPhase = makeIrFilePhase(
@@ -37,8 +25,6 @@ internal val singleAbstractMethodPhase = makeIrFilePhase(
 )
 
 private class JvmSingleAbstractMethodLowering(context: JvmBackendContext) : SingleAbstractMethodLowering(context) {
-    private val jvmContext: JvmBackendContext get() = context as JvmBackendContext
-
     override fun getWrapperVisibility(expression: IrTypeOperatorCall, scopes: List<ScopeWithIr>) =
         if (inInlineFunctionScope) Visibilities.PUBLIC else JavaVisibilities.PACKAGE_VISIBILITY
 
@@ -48,104 +34,5 @@ private class JvmSingleAbstractMethodLowering(context: JvmBackendContext) : Sing
     private val IrType.isKotlinFunInterface: Boolean
         get() = getClass()?.origin != IrDeclarationOrigin.IR_EXTERNAL_JAVA_DECLARATION_STUB
 
-    override fun getAdditionalSupertypes(supertype: IrType): List<IrType> =
+    override val IrType.needEqualsHashCodeMethods get() = isKotlinFunInterface
-        if (supertype.isKotlinFunInterface) listOf(jvmContext.ir.symbols.functionAdapter.owner.typeWith())
-        else emptyList()
-
-    override fun generateEqualsHashCode(klass: IrClass, supertype: IrType, functionDelegateField: IrField) {
-        if (!supertype.isKotlinFunInterface) return
-
-        SamEqualsHashCodeMethodsGenerator(jvmContext, klass, supertype) { receiver ->
-            irGetField(receiver, functionDelegateField)
-        }.generate()
-    }
-}
-
-/**
- * Generates equals and hashCode for SAM and fun interface wrappers, as well as an implementation of getFunctionDelegate
- * (inherited from kotlin.jvm.internal.FunctionAdapter), needed to properly implement them.
- * This class is used in two places:
- * - FunctionReferenceLowering, which is the case of SAM conversion of a (maybe adapted) function reference, e.g. `FunInterface(foo::bar)`.
- *   Note that we don't generate equals/hashCode for SAM conversion of lambdas, e.g. `FunInterface {}`, even though lambdas are represented
- *   as a local function + reference to it. The reason for this is that all lambdas are unique, so after SAM conversion they are still
- *   never equal to each other. See [FunctionReferenceLowering.FunctionReferenceBuilder.needToGenerateSamEqualsHashCodeMethods].
- * - JvmSingleAbstractMethodLowering, which is the case of SAM conversion of any value of a functional type,
- *   e.g. `val f = {}; FunInterface(f)`.
- */
-internal class SamEqualsHashCodeMethodsGenerator(
-    private val context: JvmBackendContext,
-    private val klass: IrClass,
-    private val samSuperType: IrType,
-    private val obtainFunctionDelegate: IrBuilderWithScope.(receiver: IrExpression) -> IrExpression,
-) {
-    private val builtIns: IrBuiltIns get() = context.irBuiltIns
-    private val functionAdapterClass = context.ir.symbols.functionAdapter.owner
-    private val getFunctionDelegate = functionAdapterClass.functions.single { it.name.asString() == "getFunctionDelegate" }
-
-    fun generate() {
-        generateGetFunctionDelegate()
-        generateEquals()
-        generateHashCode()
-    }
-
-    private fun generateGetFunctionDelegate() {
-        klass.addFunction(getFunctionDelegate.name.asString(), getFunctionDelegate.returnType).apply {
-            overriddenSymbols = listOf(getFunctionDelegate.symbol)
-            body = context.createIrBuilder(symbol).run {
-                irExprBody(obtainFunctionDelegate(irGet(dispatchReceiverParameter!!)))
-            }
-        }
-    }
-
-    private fun generateEquals() {
-        klass.addFunction("equals", builtIns.booleanType).apply {
-            overriddenSymbols = listOf(samSuperType.getClass()!!.functions.single {
-                it.name.asString() == "equals" &&
-                        it.extensionReceiverParameter == null &&
-                        it.valueParameters.singleOrNull()?.type == builtIns.anyNType
-            }.symbol)
-
-            val other = addValueParameter("other", builtIns.anyNType)
-            body = context.createIrBuilder(symbol).run {
-                irExprBody(
-                    irIfThenElse(
-                        builtIns.booleanType,
-                        irIs(irGet(other), samSuperType),
-                        irIfThenElse(
-                            builtIns.booleanType,
-                            irIs(irGet(other), functionAdapterClass.typeWith()),
-                            irEquals(
-                                irCall(getFunctionDelegate).also {
-                                    it.dispatchReceiver = irGet(dispatchReceiverParameter!!)
-                                },
-                                irCall(getFunctionDelegate).also {
-                                    it.dispatchReceiver = irImplicitCast(irGet(other), functionAdapterClass.typeWith())
-                                }
-                            ),
-                            irFalse()
-                        ),
-                        irFalse()
-                    )
-                )
-            }
-        }
-    }
-
-    private fun generateHashCode() {
-        klass.addFunction("hashCode", builtIns.intType).apply {
-            val hashCode = klass.superTypes.first().getClass()!!.functions.single {
-                it.name.asString() == "hashCode" && it.extensionReceiverParameter == null && it.valueParameters.isEmpty()
-            }.symbol
-            overriddenSymbols = listOf(hashCode)
-            body = context.createIrBuilder(symbol).run {
-                irExprBody(
-                    irCall(hashCode).also {
-                        it.dispatchReceiver = irCall(getFunctionDelegate).also {
-                            it.dispatchReceiver = irGet(dispatchReceiverParameter!!)
-                        }
-                    }
-                )
-            }
-        }
-    }
 }

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmSymbols.kt

@@ -59,6 +59,9 @@ class WasmSymbols(
     override val returnIfSuspended
         get() = TODO()
 
+    override val functionAdapter: IrClassSymbol
+        get() = TODO()
+
     private val wasmInternalPackage = context.module.getPackage(FqName("kotlin.wasm.internal"))
 
     val equalityFunctions = mapOf(