Hack-in setting extension function type annotation in IR

compiler/fir/analysis-tests/testData/resolve/expresssions/invoke/explicitReceiver2.kt

@@ -1,6 +1,6 @@
 
 class Bar {
-    operator fun invoke(): Foo { return this }
+    operator fun invoke(): Foo { return this } // (1)
 
 }
 
@@ -11,9 +11,9 @@ fun x() {
 class Foo {
 
 
-    operator fun Bar.invoke(): Foo { return this }
+    operator fun Bar.invoke(): Foo { return this } // (2)
 
     val x: Bar = Bar()
 
-    fun bar() = x() // Should resolve to invoke
+    fun bar() = x() // Should resolve to invoke (1)
 }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrTypeConverter.kt

@@ -5,18 +5,23 @@
 
 package org.jetbrains.kotlin.fir.backend
 
+import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.*
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.IrTypeArgument
+import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.IrStarProjectionImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 
 class Fir2IrTypeConverter(
     private val components: Fir2IrComponents
@@ -52,6 +57,14 @@ class Fir2IrTypeConverter(
         StandardClassIds.Char to irBuiltIns.charType
     )
 
+    internal val extensionFunctionTypeAnnotationConstructorCall by lazy {
+        val symbol =
+            session.firSymbolProvider.getClassLikeSymbolByFqName(CompilerConeAttributes.ExtensionFunctionType.ANNOTATION_CLASS_ID)!!
+                .toSymbol(session, classifierStorage, ConversionTypeContext.DEFAULT) as IrClassSymbol
+        val ctor = symbol.owner.declarations.firstIsInstance<IrConstructor>()
+        IrConstructorCallImpl(0, 0, symbol.defaultType, ctor.symbol, 0, 0, 0)
+    }
+
     fun FirTypeRef.toIrType(typeContext: ConversionTypeContext = ConversionTypeContext.DEFAULT): IrType {
         return when (this) {
             !is FirResolvedTypeRef -> createErrorType()
@@ -77,11 +90,12 @@ class Fir2IrTypeConverter(
                     val firSymbol = this.lookupTag.toSymbol(session) ?: return createErrorType()
                     firSymbol.toSymbol(session, classifierStorage, typeContext)
                 }
+
                 // TODO: annotations
                 IrSimpleTypeImpl(
                     irSymbol, !typeContext.definitelyNotNull && this.isMarkedNullable,
                     typeArguments.map { it.toIrTypeArgument() },
-                    emptyList()
+                    listOfNotNull(this.attributes.extensionFunctionType?.let { extensionFunctionTypeAnnotationConstructorCall })
                 )
             }
             is ConeFlexibleType -> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerGroup.kt

@@ -130,5 +130,5 @@ private constructor(
 }
 
 enum class InvokeResolvePriority {
-    NONE, COMMON_INVOKE, INVOKE_EXTENSION;
+    NONE, INVOKE_RECEIVER, COMMON_INVOKE, INVOKE_EXTENSION;
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerLevelHandler.kt

@@ -36,22 +36,37 @@ internal class CandidateFactoriesAndCollectors(
 
 typealias EnqueueTasksForInvokeReceiverCandidates = () -> Unit
 
-internal class TowerLevelHandler {
+internal open class TowerLevelHandler(
+    val collector: CandidateCollector,
+    val candidateFactoriesAndCollectors: CandidateFactoriesAndCollectors,
+    val candidateFactory: CandidateFactory
+) {
 
     // Try to avoid adding additional state here
     private var processResult = ProcessorAction.NONE
 
+    open fun onSuccessfulLevel(towerGroup: TowerGroup) {}
+
+    inline fun handleLevel(
+        info: CallInfo,
+        explicitReceiverKind: ExplicitReceiverKind,
+        group: TowerGroup,
+        towerLevel: SessionBasedTowerLevel,
+        onEmptyLevel: (TowerGroup) -> Unit
+    ) {
+        if (handleLevel(info, explicitReceiverKind, group, towerLevel) == ProcessorAction.NONE) {
+            onEmptyLevel(group)
+        }
+        if (collector.isSuccess()) onSuccessfulLevel(group)
+    }
+
     fun handleLevel(
         info: CallInfo,
         explicitReceiverKind: ExplicitReceiverKind,
         group: TowerGroup,
-        candidateFactoriesAndCollectors: CandidateFactoriesAndCollectors,
-        towerLevel: SessionBasedTowerLevel,
-        invokeResolveMode: InvokeResolveMode?,
-        candidateFactory: CandidateFactory,
-        enqueueResolverTasksForInvokeReceiverCandidates: EnqueueTasksForInvokeReceiverCandidates
+        towerLevel: SessionBasedTowerLevel
     ): ProcessorAction {
-        val resultCollector = candidateFactoriesAndCollectors.resultCollector
+        val resultCollector = collector
         val processor =
             TowerScopeLevelProcessor(
                 info.explicitReceiver,
@@ -60,6 +75,7 @@ internal class TowerLevelHandler {
                 candidateFactory,
                 group
             )
+
         when (info.callKind) {
             CallKind.VariableAccess -> {
                 towerLevel.processProperties(info.name, processor)
@@ -69,34 +85,7 @@ internal class TowerLevelHandler {
                 }
             }
             CallKind.Function -> {
-                val invokeBuiltinExtensionMode =
-                    invokeResolveMode == InvokeResolveMode.RECEIVER_FOR_INVOKE_BUILTIN_EXTENSION
-
-                if (!invokeBuiltinExtensionMode) {
-                    towerLevel.processFunctions(info.name, processor)
-                }
-
-                if (invokeResolveMode == InvokeResolveMode.IMPLICIT_CALL_ON_GIVEN_RECEIVER ||
-                    resultCollector.isSuccess()
-                ) {
-                    return processResult
-                }
-
-                val invokeReceiverProcessor = TowerScopeLevelProcessor(
-                    info.explicitReceiver,
-                    explicitReceiverKind,
-                    candidateFactoriesAndCollectors.invokeReceiverCollector!!,
-                    if (invokeBuiltinExtensionMode) candidateFactoriesAndCollectors.invokeBuiltinExtensionReceiverCandidateFactory!!
-                    else candidateFactoriesAndCollectors.invokeReceiverCandidateFactory!!,
-                    group
-                )
-                candidateFactoriesAndCollectors.invokeReceiverCollector.newDataSet()
-                towerLevel.processProperties(info.name, invokeReceiverProcessor)
-                towerLevel.processObjectsAsVariables(info.name, invokeReceiverProcessor)
-
-                if (candidateFactoriesAndCollectors.invokeReceiverCollector.isSuccess()) {
-                    enqueueResolverTasksForInvokeReceiverCandidates()
-                }
+                towerLevel.processFunctions(info.name, processor)
             }
             CallKind.CallableReference -> {
                 val stubReceiver = info.stubReceiver

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -122,6 +122,13 @@ class FirCallCompletionResultsWriterTransformer(
         return result.compose()
     }
 
+    override fun transformCheckedSafeCallSubject(
+        checkedSafeCallSubject: FirCheckedSafeCallSubject,
+        data: ExpectedArgumentType?
+    ): CompositeTransformResult<FirStatement> {
+        return checkedSafeCallSubject.transform(integerApproximator, data?.getExpectedType(checkedSafeCallSubject))
+    }
+
     override fun transformFunctionCall(functionCall: FirFunctionCall, data: ExpectedArgumentType?): CompositeTransformResult<FirStatement> {
         val calleeReference = functionCall.calleeReference as? FirNamedReferenceWithCandidate
             ?: return functionCall.compose()

compiler/tests-spec/testData/diagnostics/linked/overload-resolution/choosing-the-most-specific-candidate-from-the-overload-candidate-set/algorithm-of-msc-selection/p-3/pos/1.4.fir.kt

@@ -60,8 +60,8 @@ class Case4() : I {
         I.invoke("")
         <!DEBUG_INFO_CALL("fqName: Case4.Companion.invoke; typeCall: operator function")!>invoke("")<!>
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any")!>invoke("")<!>
-        <!DEBUG_INFO_CALL("fqName: fqName is unknown; typeCall: unresolved")!><!UNRESOLVED_REFERENCE!>I<!>("")<!>
-        <!DEBUG_INFO_EXPRESSION_TYPE("ERROR CLASS: Unresolved name: I")!><!UNRESOLVED_REFERENCE!>I<!>("")<!>
+        <!DEBUG_INFO_CALL("fqName: I.Companion.invoke; typeCall: variable&invoke")!>I("")<!>
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>I("")<!>
         <!DEBUG_INFO_CALL("fqName: Case4.Companion.invoke; typeCall: variable&invoke")!>Case4("")<!>
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any")!>Case4("")<!>
     }

compiler/tests-spec/testData/diagnostics/linked/overload-resolution/resolving-callable-references/bidirectional-resolution-for-callable-calls/p-3/neg/1.1.fir.kt

@@ -55,7 +55,7 @@ class Case3() : I {
 
     fun case() {
         I.<!AMBIGUITY!>invoke<!>(<!UNRESOLVED_REFERENCE!>::x<!>)
-        <!UNRESOLVED_REFERENCE!>I<!>(<!UNRESOLVED_REFERENCE!>::x<!>)
+        <!AMBIGUITY!>I<!>(<!UNRESOLVED_REFERENCE!>::x<!>)
         <!AMBIGUITY!>Case3<!>(<!UNRESOLVED_REFERENCE!>::x<!>)
     }
 }

compiler/tests-spec/testData/diagnostics/linked/overload-resolution/resolving-callable-references/bidirectional-resolution-for-callable-calls/p-3/pos/1.4.fir.kt

@@ -66,13 +66,13 @@ class Case3() : I {
         I.invoke(::y)
         I.invoke(::y)
 
-        <!UNRESOLVED_REFERENCE!>I<!>(<!DEBUG_INFO_EXPRESSION_TYPE("Type is unknown"), UNRESOLVED_REFERENCE!>::x<!>)
-        <!DEBUG_INFO_EXPRESSION_TYPE("ERROR CLASS: Unresolved name: I")!><!UNRESOLVED_REFERENCE!>I<!>(<!UNRESOLVED_REFERENCE!>::x<!>)<!>
-        <!DEBUG_INFO_CALL("fqName: fqName is unknown; typeCall: unresolved")!><!UNRESOLVED_REFERENCE!>I<!>(<!UNRESOLVED_REFERENCE!>::x<!>)<!>
+        I(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.reflect.KProperty0<kotlin.String>")!>::x<!>)
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Unit")!>I(::x)<!>
+        <!DEBUG_INFO_CALL("fqName: I.Companion.invoke; typeCall: variable&invoke")!>I(::x)<!>
 
-        <!UNRESOLVED_REFERENCE!>I<!>(<!DEBUG_INFO_EXPRESSION_TYPE("Type is unknown"), UNRESOLVED_REFERENCE!>::y<!>)
-        <!DEBUG_INFO_EXPRESSION_TYPE("ERROR CLASS: Unresolved name: I")!><!UNRESOLVED_REFERENCE!>I<!>(<!UNRESOLVED_REFERENCE!>::y<!>)<!>
-        <!DEBUG_INFO_CALL("fqName: fqName is unknown; typeCall: unresolved")!><!UNRESOLVED_REFERENCE!>I<!>(<!UNRESOLVED_REFERENCE!>::y<!>)<!>
+        I(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.reflect.KFunction0<kotlin.String>")!>::y<!>)
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Unit")!>I(::y)<!>
+        <!DEBUG_INFO_CALL("fqName: I.Companion.invoke; typeCall: variable&invoke")!>I(::y)<!>
 
         Case3(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.reflect.KProperty0<kotlin.String>")!>::x<!>)
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any")!>Case3(::x)<!>