Support equals/hashCode for fun interfaces in JVM and JVM_IR

#KT-33455 Fixed
2020-03-23 11:23:49 +01:00
parent de461dd9a5
commit 9fa8e009c6
25 changed files with 1036 additions and 23 deletions
@@ -163,7 +163,7 @@ abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) :
            setSourceRange(createFor)
        }.apply {
            createImplicitParameterDeclarationWithWrappedDescriptor()
-            superTypes += superType
+            superTypes = listOf(superType) + getAdditionalSupertypes(superType)
            parent = enclosingContainer!!
        }

@@ -202,7 +202,7 @@ abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) :
            isSuspend = superMethod.isSuspend
            setSourceRange(createFor)
        }.apply {
-            overriddenSymbols += superMethod.symbol
+            overriddenSymbols = listOf(superMethod.symbol)
            dispatchReceiverParameter = subclass.thisReceiver!!.copyTo(this)
            valueParameters = superMethod.valueParameters.map { it.copyTo(this) }
            body = context.createIrBuilder(symbol).irBlockBody {
@@ -213,8 +213,14 @@ abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) :
            }
        }

+        generateEqualsHashCode(subclass, superType, field)
+
        subclass.addFakeOverrides()

        return subclass
    }
+
+    protected open fun getAdditionalSupertypes(supertype: IrType): List<IrType> = emptyList()
+
+    protected open fun generateEqualsHashCode(klass: IrClass, supertype: IrType, functionDelegateField: IrField) {}
 }
@@ -499,6 +499,10 @@ class JvmSymbols(
        }
    }

+    val functionAdapter: IrClassSymbol = createClass(FqName("kotlin.jvm.internal.FunctionAdapter"), ClassKind.INTERFACE) { klass ->
+        klass.addFunction("getFunctionDelegate", irBuiltIns.functionClass.starProjectedType, Modality.ABSTRACT)
+    }
+
    val getOrCreateKotlinPackage: IrSimpleFunctionSymbol =
        reflection.functionByName("getOrCreateKotlinPackage")

@@ -152,6 +152,11 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
            }
        }

+        private val needToGenerateSamEqualsHashCodeMethods =
+            samSuperType != null &&
+                    samSuperType.getClass()?.origin != IrDeclarationOrigin.IR_EXTERNAL_JAVA_DECLARATION_STUB &&
+                    (adaptedReferenceOriginalTarget != null || !isLambda)
+
        private val superType =
            samSuperType ?: when {
                adaptedReferenceOriginalTarget != null -> context.ir.symbols.adaptedFunctionReference
@@ -169,10 +174,18 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
            name = SpecialNames.NO_NAME_PROVIDED
        }.apply {
            parent = currentDeclarationParent ?: error("No current declaration parent at ${irFunctionReference.dump()}")
-            superTypes += superType
-            if (samSuperType == null)
-                superTypes += functionSuperClass.typeWith(parameterTypes)
-            if (irFunctionReference.isSuspend) superTypes += context.ir.symbols.suspendFunctionInterface.defaultType
+            superTypes = listOfNotNull(
+                superType,
+                if (samSuperType == null)
+                    functionSuperClass.typeWith(parameterTypes)
+                else null,
+                if (irFunctionReference.isSuspend)
+                    context.ir.symbols.suspendFunctionInterface.defaultType
+                else null,
+                if (needToGenerateSamEqualsHashCodeMethods)
+                    context.ir.symbols.functionAdapter.defaultType
+                else null,
+            )
            createImplicitParameterDeclarationWithWrappedDescriptor()
            copyAttributes(irFunctionReference)
            if (isLambda) {
@@ -195,11 +208,11 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
        fun build(): IrExpression = context.createJvmIrBuilder(currentScope!!.scope.scopeOwnerSymbol).run {
            irBlock(irFunctionReference.startOffset, irFunctionReference.endOffset) {
                val constructor = createConstructor()
-                createInvokeMethod(
+                val boundReceiverVar =
                    if (samSuperType != null && boundReceiver != null) {
                        irTemporary(boundReceiver.second)
                    } else null
-                )
+                createInvokeMethod(boundReceiverVar)

                if (!isLambda && samSuperType == null && !useOptimizedSuperClass) {
                    createLegacyMethodOverride(irSymbols.functionReferenceGetSignature.owner) {
@@ -213,6 +226,10 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
                    }
                }

+                if (needToGenerateSamEqualsHashCodeMethods) {
+                    generateSamEqualsHashCodeMethods(boundReceiverVar)
+                }
+
                +functionReferenceClass
                +irCall(constructor.symbol).apply {
                    if (valueArgumentsCount > 0) putValueArgument(0, boundReceiver!!.second)
@@ -220,6 +237,40 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
            }
        }

+        private fun JvmIrBuilder.generateSamEqualsHashCodeMethods(boundReceiverVar: IrVariable?) {
+            checkNotNull(samSuperType) { "equals/hashCode can only be generated for fun interface wrappers: ${callee.render()}" }
+
+            if (!useOptimizedSuperClass) {
+                // This is the case of a fun interface wrapper over a (maybe adapted) function reference,
+                // with `-Xno-optimized-callable-referenced` enabled. We can't use constructors of FunctionReferenceImpl,
+                // so we'd need to basically generate a full class for a reference inheriting from FunctionReference,
+                // effectively disabling the optimization of fun interface wrappers over references.
+                // This scenario is probably not very popular because it involves using equals/hashCode on function references
+                // and enabling the mentioned internal compiler argument.
+                // Right now we generate them as abstract so that any call would result in AbstractMethodError.
+                // TODO: generate getFunctionDelegate, equals and hashCode properly in this case
+                functionReferenceClass.addFunction("equals", backendContext.irBuiltIns.booleanType, Modality.ABSTRACT).apply {
+                    addValueParameter("other", backendContext.irBuiltIns.anyNType)
+                }
+                functionReferenceClass.addFunction("hashCode", backendContext.irBuiltIns.intType, Modality.ABSTRACT)
+                return
+            }
+
+            SamEqualsHashCodeMethodsGenerator(backendContext, functionReferenceClass, samSuperType) { receiver ->
+                val internalClass = when {
+                    adaptedReferenceOriginalTarget != null -> backendContext.ir.symbols.adaptedFunctionReference
+                    else -> backendContext.ir.symbols.functionReferenceImpl
+                }
+                val constructor = internalClass.owner.constructors.single {
+                    // arity, [receiver], owner, name, signature, flags
+                    it.valueParameters.size == 1 + (if (boundReceiver != null) 1 else 0) + 4
+                }
+                irCallConstructor(constructor.symbol, emptyList()).apply {
+                    generateConstructorCallArguments(this) { irGet(boundReceiverVar!!) }
+                }
+            }.generate()
+        }
+
        private fun createConstructor(): IrConstructor =
            functionReferenceClass.addConstructor {
                origin = JvmLoweredDeclarationOrigin.GENERATED_MEMBER_IN_CALLABLE_REFERENCE
@@ -6,12 +6,24 @@
 package org.jetbrains.kotlin.backend.jvm.lower

 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.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.descriptors.IrBuiltIns
+import org.jetbrains.kotlin.ir.expressions.IrExpression
 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(
@@ -21,8 +33,115 @@ internal val singleAbstractMethodPhase = makeIrFilePhase(
 )

 private class JvmSingleAbstractMethodLowering(context: JvmBackendContext) : SingleAbstractMethodLowering(context) {
+    private val jvmContext: JvmBackendContext get() = context as JvmBackendContext
+
    override val privateGeneratedWrapperVisibility: Visibility
        get() = JavaVisibilities.PACKAGE_VISIBILITY

-    override fun getSuperTypeForWrapper(typeOperand: IrType) = typeOperand.erasedUpperBound.defaultType
-}
+    override fun getSuperTypeForWrapper(typeOperand: IrType): IrType =
+        typeOperand.erasedUpperBound.defaultType
+
+    private val IrType.isKotlinFunInterface: Boolean
+        get() = getClass()?.origin != IrDeclarationOrigin.IR_EXTERNAL_JAVA_DECLARATION_STUB
+
+    override fun getAdditionalSupertypes(supertype: IrType): List<IrType> =
+        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!!)
+                        }
+                    }
+                )
+            }
+        }
+    }
+}