FIR: Support coercion-to-unit semantics for callable references

compiler/fir/analysis-tests/testData/resolveWithStdlib/callableReferences/coercionToUnit.kt

+public fun <T> T.myAlso(block: (T) -> Unit): T = TODO()
+
+class B {
+    fun add(x: String): Boolean = true
+}
+
+fun main(b: B) {
+    "".myAlso(b::add)
+}

compiler/fir/analysis-tests/testData/resolveWithStdlib/callableReferences/coercionToUnit.txt

+FILE: coercionToUnit.kt
+    public final fun <T> R|T|.myAlso(block: R|(T) -> kotlin/Unit|): R|T| {
+        ^myAlso R|kotlin/TODO|()
+    }
+    public final class B : R|kotlin/Any| {
+        public constructor(): R|B| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final fun add(x: R|kotlin/String|): R|kotlin/Boolean| {
+            ^add Boolean(true)
+        }
+
+    }
+    public final fun main(b: R|B|): R|kotlin/Unit| {
+        String().R|/myAlso|<R|kotlin/String|>(R|<local>/b|::R|/B.add|)
+    }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithStdlibTestGenerated.java

@@ -300,6 +300,11 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
             runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/callableReferences/beyoundCalls.kt");
         }
 
+        @TestMetadata("coercionToUnit.kt")
+        public void testCoercionToUnit() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/callableReferences/coercionToUnit.kt");
+        }
+
         @TestMetadata("companions.kt")
         public void testCompanions() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/callableReferences/companions.kt");

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolverParts.kt

@@ -11,11 +11,13 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
+import org.jetbrains.kotlin.fir.inferenceContext
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.inference.ResolvedCallableReferenceAtom
 import org.jetbrains.kotlin.fir.resolve.inference.csBuilder
+import org.jetbrains.kotlin.fir.resolve.inference.extractInputOutputTypesFromCallableReferenceExpectedType
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.SyntheticSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
@@ -231,10 +233,11 @@ internal object CheckCallableReferenceExpectedType : CheckerStage() {
         val returnTypeRef = candidate.bodyResolveComponents.returnTypeCalculator.tryCalculateReturnType(fir)
 
         val resultingType: ConeKotlinType = when (fir) {
-            is FirFunction -> createKFunctionType(
+            is FirFunction -> callInfo.session.createKFunctionType(
                 fir, resultingReceiverType, returnTypeRef,
                 expectedParameterNumberWithReceiver = expectedType?.let { it.typeArguments.size - 1 },
-                isSuspend = (fir as? FirSimpleFunction)?.isSuspend == true
+                isSuspend = (fir as? FirSimpleFunction)?.isSuspend == true,
+                expectedReturnType = extractInputOutputTypesFromCallableReferenceExpectedType(expectedType, callInfo.session)?.outputType
             )
             is FirVariable<*> -> createKPropertyType(fir, resultingReceiverType, returnTypeRef)
             else -> ConeKotlinErrorType("Unknown callable kind: ${fir::class}")
@@ -286,13 +289,15 @@ private fun createKPropertyType(
     )
 }
 
-private fun createKFunctionType(
+private fun FirSession.createKFunctionType(
     function: FirFunction<*>,
     receiverType: ConeKotlinType?,
     returnTypeRef: FirResolvedTypeRef,
     expectedParameterNumberWithReceiver: Int?,
-    isSuspend: Boolean
+    isSuspend: Boolean,
+    expectedReturnType: ConeKotlinType?
 ): ConeKotlinType {
+    // The similar adaptations: defaults and coercion-to-unit happen at org.jetbrains.kotlin.resolve.calls.components.CallableReferencesCandidateFactory.getCallableReferenceAdaptation
     val parameterTypes = mutableListOf<ConeKotlinType>()
     val expectedParameterNumber = when {
         expectedParameterNumberWithReceiver == null -> null
@@ -306,9 +311,15 @@ private fun createKFunctionType(
         }
     }
 
+    val returnType =
+        if (expectedReturnType != null && inferenceContext.run { expectedReturnType.isUnit() })
+            expectedReturnType
+        else
+            returnTypeRef.coneTypeSafe() ?: ConeKotlinErrorType("No type for return type of $function")
+
     return createFunctionalType(
         parameterTypes, receiverType = receiverType,
-        rawReturnType = returnTypeRef.coneTypeSafe() ?: ConeKotlinErrorType("No type for return type of $function"),
+        rawReturnType = returnType,
         isKFunctionType = true,
         isSuspend = isSuspend
     )

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedAtoms.kt

@@ -99,9 +99,9 @@ class ResolvedCallableReferenceAtom(
 
 //  -------------------------- Utils --------------------------
 
-private data class InputOutputTypes(val inputTypes: List<ConeKotlinType>, val outputType: ConeKotlinType)
+internal data class InputOutputTypes(val inputTypes: List<ConeKotlinType>, val outputType: ConeKotlinType)
 
-private fun extractInputOutputTypesFromCallableReferenceExpectedType(
+internal fun extractInputOutputTypesFromCallableReferenceExpectedType(
     expectedType: ConeKotlinType?,
     session: FirSession
 ): InputOutputTypes? {
@@ -150,4 +150,4 @@ private fun extractInputOutputTypesFromFunctionType(
     )
 
     return InputOutputTypes(parameters, outputType)
-}
\ No newline at end of file
+}

compiler/testData/codegen/box/callableReference/adaptedReferences/noAdaptedReferencesIfNoOptimizedReferencesEnabled.kt

 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 // KOTLIN_CONFIGURATION_FLAGS: +JVM.NO_OPTIMIZED_CALLABLE_REFERENCES
 

compiler/testData/codegen/box/callableReference/adaptedReferences/noNameClashForReferencesToSameFunction.kt

-// IGNORE_BACKEND_FIR: JVM_IR
-
 var result = ""
 
 class C(val token: String) {

compiler/testData/codegen/box/callableReference/bound/coercionToUnit.kt

-// IGNORE_BACKEND_FIR: JVM_IR
 // KJS_WITH_FULL_RUNTIME
 // !LANGUAGE: +NewInference
 // WITH_RUNTIME

compiler/testData/codegen/box/callableReference/function/coercionToUnit.kt

 // !LANGUAGE: +NewInference
-// IGNORE_BACKEND_FIR: JVM_IR
 
 fun foo(s: String): Boolean {
     if (s != "kotlin") throw AssertionError(s)

compiler/testData/diagnostics/tests/callableReference/function/ambiguityTopLevelVsTopLevel.fir.kt

@@ -5,5 +5,5 @@ fun foo(x: Any, y: Int) = y
 fun main() {
     <!UNRESOLVED_REFERENCE!>::foo<!>
     
-    val fooRef: (Int, Any) -> Unit = <!UNRESOLVED_REFERENCE!>::foo<!>
+    val fooRef: (Int, Any) -> Unit = ::foo
 }

compiler/testData/diagnostics/tests/suspendConversion/suspendConversionWithFunInterfaces.fir.kt

@@ -17,5 +17,5 @@ fun test() {
     <!INAPPLICABLE_CANDIDATE!>foo1<!>(::bar1)
     <!INAPPLICABLE_CANDIDATE!>foo1<!>(::bar2)
 
-    <!INAPPLICABLE_CANDIDATE!>foo1<!>(<!UNRESOLVED_REFERENCE!>::bar3<!>) // Should be ambiguity
+    foo1(::bar3) // Should be ambiguity
 }

compiler/testData/ir/irText/expressions/callableReferences/adaptedWithCoercionToUnit.fir.txt

@@ -19,16 +19,16 @@ FILE fqName:<root> fileName:/adaptedWithCoercionToUnit.kt
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun fnv (vararg xs: kotlin.Int): kotlin.Int declared in <root>'
         CONST Int type=kotlin.Int value=1
-  FUN name:test0 visibility:public modality:FINAL <> () returnType:IrErrorType
+  FUN name:test0 visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun test0 (): IrErrorType declared in <root>'
-        ERROR_CALL 'Unresolved reference: <Inapplicable(INAPPLICABLE): [/useUnit0]>#' type=IrErrorType
-          FUNCTION_REFERENCE 'public final fun fn0 (): kotlin.Int declared in <root>' type=kotlin.reflect.KFunction0<kotlin.Int> origin=null reflectionTarget=<same>
-  FUN name:test1 visibility:public modality:FINAL <> () returnType:IrErrorType
+      RETURN type=kotlin.Nothing from='public final fun test0 (): kotlin.Unit declared in <root>'
+        CALL 'public final fun useUnit0 (fn: kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+          fn: FUNCTION_REFERENCE 'public final fun fn0 (): kotlin.Int declared in <root>' type=kotlin.reflect.KFunction0<kotlin.Unit> origin=null reflectionTarget=<same>
+  FUN name:test1 visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun test1 (): IrErrorType declared in <root>'
-        ERROR_CALL 'Unresolved reference: <Inapplicable(INAPPLICABLE): [/useUnit1]>#' type=IrErrorType
-          FUNCTION_REFERENCE 'public final fun fn1 (x: kotlin.Int): kotlin.Int declared in <root>' type=kotlin.reflect.KFunction1<kotlin.Int, kotlin.Int> origin=null reflectionTarget=<same>
+      RETURN type=kotlin.Nothing from='public final fun test1 (): kotlin.Unit declared in <root>'
+        CALL 'public final fun useUnit1 (fn: kotlin.Function1<kotlin.Int, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+          fn: FUNCTION_REFERENCE 'public final fun fn1 (x: kotlin.Int): kotlin.Int declared in <root>' type=kotlin.reflect.KFunction1<kotlin.Int, kotlin.Unit> origin=null reflectionTarget=<same>
   FUN name:testV0 visibility:public modality:FINAL <> () returnType:IrErrorType
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun testV0 (): IrErrorType declared in <root>'

compiler/testData/ir/irText/expressions/callableReferences/withAdaptedArguments.fir.txt

@@ -73,11 +73,11 @@ FILE fqName:<root> fileName:/withAdaptedArguments.kt
       RETURN type=kotlin.Nothing from='public final fun testVararg (): IrErrorType declared in <root>'
         ERROR_CALL 'Unresolved reference: <Inapplicable(INAPPLICABLE): [/use]>#' type=IrErrorType
           FUNCTION_REFERENCE 'public final fun fnWithVarargs (vararg xs: kotlin.Int): kotlin.String declared in <root>' type=kotlin.reflect.KFunction1<kotlin.IntArray, kotlin.String> origin=null reflectionTarget=<same>
-  FUN name:testCoercionToUnit visibility:public modality:FINAL <> () returnType:IrErrorType
+  FUN name:testCoercionToUnit visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun testCoercionToUnit (): IrErrorType declared in <root>'
-        ERROR_CALL 'Unresolved reference: <Inapplicable(INAPPLICABLE): [/coerceToUnit]>#' type=IrErrorType
-          FUNCTION_REFERENCE 'public final fun fnWithDefault (a: kotlin.Int, b: kotlin.Int): kotlin.String declared in <root>' type=kotlin.reflect.KFunction1<kotlin.Int, kotlin.String> origin=null reflectionTarget=<same>
+      RETURN type=kotlin.Nothing from='public final fun testCoercionToUnit (): kotlin.Unit declared in <root>'
+        CALL 'public final fun coerceToUnit (fn: kotlin.Function1<kotlin.Int, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+          fn: FUNCTION_REFERENCE 'public final fun fnWithDefault (a: kotlin.Int, b: kotlin.Int): kotlin.String declared in <root>' type=kotlin.reflect.KFunction1<kotlin.Int, kotlin.Unit> origin=null reflectionTarget=<same>
   FUN name:testImportedObjectMember visibility:public modality:FINAL <> () returnType:IrErrorType
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun testImportedObjectMember (): IrErrorType declared in <root>'