NI: take into account an extension function annotation during CST calculation

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeInferenceContext.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.types
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.NoSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
@@ -17,7 +18,6 @@ import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
-import org.jetbrains.kotlin.types.checker.NewCapturedType
 import org.jetbrains.kotlin.types.model.*
 import org.jetbrains.kotlin.utils.addToStdlib.cast
 
@@ -50,10 +50,12 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
         return coneFlexibleOrSimpleType(this, lowerBound, upperBound)
     }
 
+    // TODO: implement taking into account `isExtensionFunction`
     override fun createSimpleType(
         constructor: TypeConstructorMarker,
         arguments: List<TypeArgumentMarker>,
-        nullable: Boolean
+        nullable: Boolean,
+        isExtensionFunction: Boolean
     ): SimpleTypeMarker {
         require(constructor is FirClassifierSymbol<*>)
         return when (constructor) {
@@ -96,6 +98,12 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
                 || this is ConeTypeParameterType
     }
 
+    // TODO: implement checking for extension function
+    override fun SimpleTypeMarker.isExtensionFunction(): Boolean {
+        require(this is ConeKotlinType)
+        return false
+    }
+
     override fun KotlinTypeMarker.typeDepth() = when (this) {
         is ConeSimpleKotlinType -> typeDepth()
         is ConeFlexibleType -> maxOf(lowerBound().typeDepth(), upperBound().typeDepth())

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/types/IrTypeSystemContext.kt

@@ -222,10 +222,12 @@ interface IrTypeSystemContext : TypeSystemContext, TypeSystemCommonSuperTypesCon
         return IrDynamicTypeImpl(null, emptyList(), Variance.INVARIANT)
     }
 
+    // TODO: implement taking into account `isExtensionFunction`
     override fun createSimpleType(
         constructor: TypeConstructorMarker,
         arguments: List<TypeArgumentMarker>,
-        nullable: Boolean
+        nullable: Boolean,
+        isExtensionFunction: Boolean
     ): SimpleTypeMarker = IrSimpleTypeImpl(constructor as IrClassifierSymbol, nullable, arguments.map { it as IrTypeArgument }, emptyList())
 
     private fun TypeVariance.convertVariance(): Variance {
@@ -243,6 +245,11 @@ interface IrTypeSystemContext : TypeSystemContext, TypeSystemCommonSuperTypesCon
 
     override fun KotlinTypeMarker.canHaveUndefinedNullability() = this is IrSimpleType && classifier is IrTypeParameterSymbol
 
+    override fun SimpleTypeMarker.isExtensionFunction(): Boolean {
+        require(this is IrSimpleType)
+        return this.hasAnnotation(KotlinBuiltIns.FQ_NAMES.extensionFunctionType)
+    }
+
     override fun SimpleTypeMarker.typeDepth(): Int {
         val maxInArguments = (this as IrSimpleType).arguments.asSequence().map {
             if (it is IrStarProjection) 1 else it.getType().typeDepth()

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/NewCommonSuperTypeCalculator.kt

@@ -21,7 +21,6 @@ import org.jetbrains.kotlin.types.AbstractFlexibilityChecker.hasDifferentFlexibi
 import org.jetbrains.kotlin.types.AbstractNullabilityChecker.hasPathByNotMarkedNullableNodes
 import org.jetbrains.kotlin.types.checker.SimpleClassicTypeSystemContext
 import org.jetbrains.kotlin.types.model.*
-import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 object NewCommonSuperTypeCalculator {
     // TODO: Bridge for old calls
@@ -332,7 +331,7 @@ object NewCommonSuperTypeCalculator {
 
             arguments.add(argument)
         }
-        return createSimpleType(constructor, arguments, nullable = false)
+        return createSimpleType(constructor, arguments, nullable = false, isExtensionFunction = types.all { it.isExtensionFunction() })
     }
 
     private fun TypeSystemCommonSuperTypesContext.uncaptureFromSubtyping(typeArgument: TypeArgumentMarker): TypeArgumentMarker {

compiler/testData/diagnostics/tests/inference/kt36819.fir.kt

@@ -0,0 +1,47 @@
+// !LANGUAGE: +NewInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE -UNUSED_ANONYMOUS_PARAMETER -CAST_NEVER_SUCCEEDS
+// ISSUE: KT-36819
+
+// Case 1
+
+fun <K> select(vararg x: K) = x[0]
+interface A
+class B: A
+class C: A
+fun <T> id1(x: T): T = x
+fun <R> id2(x: R): R = x
+
+class Out<out R>(x: R)
+
+fun main() {
+    val x1 = select(id1 { B() }, id2 { C() })
+    val x2 = select({ B() }, { C() })
+    val x3 = select(id1(Out(B())), id2(Out(C())))
+}
+
+// Case 2
+
+fun <R> fold(initial: R, operation: (R) -> Unit) {}
+
+fun foo() {
+    fold({ x: Int -> x }) { acc -> }
+}
+
+// Case 3
+
+class Foo
+
+typealias X = Foo.(Foo.() -> Unit) -> Unit
+
+fun bar() {
+    val y = when {
+        false -> { _ ->
+            Unit
+        }
+        true -> null as X
+        else -> { x ->
+            <!UNRESOLVED_REFERENCE!>x<!>()
+            Unit
+        }
+    }
+}

compiler/testData/diagnostics/tests/inference/kt36819.kt

@@ -1,8 +1,9 @@
-// FIR_IDENTICAL
 // !LANGUAGE: +NewInference
-// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE -UNUSED_ANONYMOUS_PARAMETER
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE -UNUSED_ANONYMOUS_PARAMETER -CAST_NEVER_SUCCEEDS
 // ISSUE: KT-36819
 
+// Case 1
+
 fun <K> select(vararg x: K) = x[0]
 interface A
 class B: A
@@ -18,8 +19,29 @@ fun main() {
     val x3 = select(id1(Out(B())), id2(Out(C())))
 }
 
+// Case 2
+
 fun <R> fold(initial: R, operation: (R) -> Unit) {}
 
 fun foo() {
     fold({ x: Int -> x }) { acc -> }
 }
+
+// Case 3
+
+class Foo
+
+typealias X = Foo.(Foo.() -> Unit) -> Unit
+
+fun bar() {
+    val y = when {
+        false -> { _ ->
+            Unit
+        }
+        true -> null as X
+        else -> { x ->
+            x()
+            Unit
+        }
+    }
+}

compiler/testData/diagnostics/tests/inference/kt36819.txt

@@ -1,5 +1,6 @@
 package
 
+public fun bar(): kotlin.Unit
 public fun </*0*/ R> fold(/*0*/ initial: R, /*1*/ operation: (R) -> kotlin.Unit): kotlin.Unit
 public fun foo(): kotlin.Unit
 public fun </*0*/ T> id1(/*0*/ x: T): T
@@ -27,9 +28,17 @@ public final class C : A {
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }
 
+public final class Foo {
+    public constructor Foo()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
 public final class Out</*0*/ out R> {
     public constructor Out</*0*/ out R>(/*0*/ x: R)
     public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
     public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }
+public typealias X = Foo.(Foo.() -> kotlin.Unit) -> kotlin.Unit

core/descriptors/src/org/jetbrains/kotlin/types/checker/ClassicTypeSystemContext.kt

@@ -11,11 +11,11 @@ import org.jetbrains.kotlin.builtins.PrimitiveType
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.descriptors.annotations.BuiltInAnnotationDescriptor
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.FqNameUnsafe
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.inference.CapturedType
-import org.jetbrains.kotlin.resolve.calls.inference.CapturedTypeConstructorImpl
 import org.jetbrains.kotlin.resolve.constants.IntegerLiteralTypeConstructor
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasExactAnnotation
@@ -434,11 +434,17 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
     override fun createSimpleType(
         constructor: TypeConstructorMarker,
         arguments: List<TypeArgumentMarker>,
-        nullable: Boolean
+        nullable: Boolean,
+        isExtensionFunction: Boolean
     ): SimpleTypeMarker {
         require(constructor is TypeConstructor, constructor::errorMessage)
+
+        val annotations = if (isExtensionFunction) {
+            Annotations.create(listOf(BuiltInAnnotationDescriptor(builtIns, FQ_NAMES.extensionFunctionType, emptyMap())))
+        } else Annotations.EMPTY
+
         @Suppress("UNCHECKED_CAST")
-        return KotlinTypeFactory.simpleType(Annotations.EMPTY, constructor, arguments as List<TypeProjection>, nullable)
+        return KotlinTypeFactory.simpleType(annotations, constructor, arguments as List<TypeProjection>, nullable)
     }
 
     override fun createTypeArgument(type: KotlinTypeMarker, variance: TypeVariance): TypeArgumentMarker {
@@ -458,6 +464,11 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
                 this is NewCapturedType
     }
 
+    override fun SimpleTypeMarker.isExtensionFunction(): Boolean {
+        require(this is SimpleType, this::errorMessage)
+        return this.hasAnnotation(FQ_NAMES.extensionFunctionType)
+    }
+
     override fun SimpleTypeMarker.replaceArguments(newArguments: List<TypeArgumentMarker>): SimpleTypeMarker {
         require(this is SimpleType, this::errorMessage)
         @Suppress("UNCHECKED_CAST")

core/type-system/src/org/jetbrains/kotlin/types/model/TypeSystemContext.kt

@@ -56,7 +56,13 @@ interface TypeSystemBuiltInsContext {
 
 interface TypeSystemTypeFactoryContext {
     fun createFlexibleType(lowerBound: SimpleTypeMarker, upperBound: SimpleTypeMarker): KotlinTypeMarker
-    fun createSimpleType(constructor: TypeConstructorMarker, arguments: List<TypeArgumentMarker>, nullable: Boolean): SimpleTypeMarker
+    fun createSimpleType(
+        constructor: TypeConstructorMarker,
+        arguments: List<TypeArgumentMarker>,
+        nullable: Boolean,
+        isExtensionFunction: Boolean = false
+    ): SimpleTypeMarker
+
     fun createTypeArgument(type: KotlinTypeMarker, variance: TypeVariance): TypeArgumentMarker
     fun createStarProjection(typeParameter: TypeParameterMarker): TypeArgumentMarker
 
@@ -88,6 +94,8 @@ interface TypeSystemCommonSuperTypesContext : TypeSystemContext, TypeSystemTypeF
 
     fun KotlinTypeMarker.canHaveUndefinedNullability(): Boolean
 
+    fun SimpleTypeMarker.isExtensionFunction(): Boolean
+
     fun SimpleTypeMarker.typeDepth(): Int
 
     fun KotlinTypeMarker.typeDepth(): Int = when (this) {