[FIR] Add ExtensionFunctionType attribute

#KT-39034 Fixed

compiler/fir/analysis-tests/testData/resolve/inference/nestedExtensionFunctionType.kt

@@ -0,0 +1,19 @@
+// ISSUE: KT-39034
+
+interface A
+
+fun test_0(a: A, f: A.() -> Unit) {
+    a.f()
+}
+
+fun test_1(a: A, ys: List<A.() -> Unit>) {
+    for (y in ys) {
+        a.y()
+    }
+}
+
+fun test_2(a: A, vararg zs: A.() -> Unit) {
+    for (z in zs) {
+        a.z()
+    }
+}

compiler/fir/analysis-tests/testData/resolve/inference/nestedExtensionFunctionType.txt

@@ -0,0 +1,22 @@
+FILE: nestedExtensionFunctionType.kt
+    public abstract interface A : R|kotlin/Any| {
+    }
+    public final fun test_0(a: R|A|, f: R|A.() -> kotlin/Unit|): R|kotlin/Unit| {
+        R|<local>/f|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/a|)
+    }
+    public final fun test_1(a: R|A|, ys: R|kotlin/collections/List<kotlin/Function1<A, kotlin/Unit>>|): R|kotlin/Unit| {
+        lval <iterator>: R|kotlin/collections/Iterator<kotlin/Function1<A, kotlin/Unit>>| = R|<local>/ys|.R|FakeOverride<kotlin/collections/List.iterator: R|kotlin/collections/Iterator<kotlin/Function1<A, kotlin/Unit>>|>|()
+        while(R|<local>/<iterator>|.R|kotlin/collections/Iterator.hasNext|()) {
+            lval y: R|(A) -> kotlin/Unit| = R|<local>/<iterator>|.R|FakeOverride<kotlin/collections/Iterator.next: R|(A) -> kotlin/Unit|>|()
+            R|<local>/y|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/a|)
+        }
+
+    }
+    public final fun test_2(a: R|A|, vararg zs: @R|kotlin/ExtensionFunctionType|() R|kotlin/Array<out kotlin/Function1<A, kotlin/Unit>>|): R|kotlin/Unit| {
+        lval <iterator>: R|kotlin/collections/Iterator<kotlin/Function1<A, kotlin/Unit>>| = R|<local>/zs|.R|FakeOverride<kotlin/Array.iterator: R|kotlin/collections/Iterator<kotlin/Function1<A, kotlin/Unit>>|>|()
+        while(R|<local>/<iterator>|.R|kotlin/collections/Iterator.hasNext|()) {
+            lval z: R|(A) -> kotlin/Unit| = R|<local>/<iterator>|.R|FakeOverride<kotlin/collections/Iterator.next: R|(A) -> kotlin/Unit|>|()
+            R|<local>/z|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/a|)
+        }
+
+    }

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

@@ -1479,6 +1479,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/lambdaAsReturnStatementOfLambda.kt");
         }
 
+        @TestMetadata("nestedExtensionFunctionType.kt")
+        public void testNestedExtensionFunctionType() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/inference/nestedExtensionFunctionType.kt");
+        }
+
         @TestMetadata("nestedLambdas.kt")
         public void testNestedLambdas() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/nestedLambdas.kt");

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

@@ -1479,6 +1479,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/lambdaAsReturnStatementOfLambda.kt");
         }
 
+        @TestMetadata("nestedExtensionFunctionType.kt")
+        public void testNestedExtensionFunctionType() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/inference/nestedExtensionFunctionType.kt");
+        }
+
         @TestMetadata("nestedLambdas.kt")
         public void testNestedLambdas() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/nestedLambdas.kt");

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/CompilerConeAttributes.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.types
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import kotlin.reflect.KClass
 
 object CompilerConeAttributes {
@@ -30,7 +31,18 @@ object CompilerConeAttributes {
 
         override val key: KClass<out NoInfer> = NoInfer::class
     }
+
+    object ExtensionFunctionType : ConeAttribute<ExtensionFunctionType>() {
+        val ANNOTATION_CLASS_ID = ClassId(FqName("kotlin"), Name.identifier("ExtensionFunctionType"))
+
+        override fun union(other: ExtensionFunctionType?): ExtensionFunctionType? = other
+        override fun intersect(other: ExtensionFunctionType?): ExtensionFunctionType? = this
+        override fun isSubtypeOf(other: ExtensionFunctionType?): Boolean = true
+
+        override val key: KClass<out ExtensionFunctionType> = ExtensionFunctionType::class
+    }
 }
 
 val ConeAttributes.exact: CompilerConeAttributes.Exact? by ConeAttributes.attributeAccessor<CompilerConeAttributes.Exact>()
 val ConeAttributes.noInfer: CompilerConeAttributes.NoInfer? by ConeAttributes.attributeAccessor<CompilerConeAttributes.NoInfer>()
+val ConeAttributes.extensionFunctionType: CompilerConeAttributes.ExtensionFunctionType? by ConeAttributes.attributeAccessor<CompilerConeAttributes.ExtensionFunctionType>()

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

@@ -359,19 +359,3 @@ internal fun FirExpression.getExpectedType(
 fun ConeKotlinType.varargElementType(session: FirSession): ConeKotlinType {
     return this.arrayElementType() ?: error("Failed to extract! ${this.render()}!")
 }
-
-fun FirTypeRef.isExtensionFunctionType(session: FirSession): Boolean {
-    if (annotations.any(FirAnnotationCall::isExtensionFunctionAnnotationCall)) return true
-
-    if (this !is FirResolvedTypeRef) return false
-    val type = type as? ConeClassLikeType ?: return false
-    if (type.fullyExpandedType(session) === type) return false
-
-    val typeAlias = type.lookupTag
-        .toSymbol(session)
-        ?.safeAs<FirTypeAliasSymbol>()?.fir ?: return false
-
-    if (typeAlias.expandedTypeRef.annotations.any(FirAnnotationCall::isExtensionFunctionAnnotationCall)) return true
-
-    return false
-}

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

@@ -97,7 +97,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
             val receiverType = callable.receiverTypeRef?.coneTypeUnsafe<ConeKotlinType>()
             if (receiverType != null) return receiverType
             val returnTypeRef = callable.returnTypeRef as? FirResolvedTypeRef ?: return null
-            if (!returnTypeRef.isExtensionFunctionType(bodyResolveComponents.session)) return null
+            if (!returnTypeRef.type.isExtensionFunctionType(bodyResolveComponents.session)) return null
             return (returnTypeRef.type.typeArguments.firstOrNull() as? ConeKotlinTypeProjection)?.type
         }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirTypeResolverImpl.kt

@@ -101,7 +101,8 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver {
         return ConeClassLikeTypeImpl(
             resolveBuiltInQualified(classId, session).toLookupTag(),
             parameters.toTypedArray(),
-            typeRef.isMarkedNullable
+            typeRef.isMarkedNullable,
+            attributes
         )
     }
 

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

@@ -5,11 +5,15 @@
 
 package org.jetbrains.kotlin.fir.types
 
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.name.ClassId
@@ -18,6 +22,7 @@ import org.jetbrains.kotlin.types.AbstractNullabilityChecker
 import org.jetbrains.kotlin.types.AbstractStrictEqualityTypeChecker
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 object ConeNullabilityChecker {
     fun isSubtypeOfAny(context: ConeTypeContext, type: ConeKotlinType): Boolean {
@@ -217,3 +222,12 @@ fun ConeClassifierLookupTag.constructType(typeArguments: Array<out ConeTypeProje
         else -> error("! ${this::class}")
     }
 }
+
+fun FirTypeRef.isExtensionFunctionType(session: FirSession): Boolean {
+    return coneTypeSafe<ConeKotlinType>()?.isExtensionFunctionType(session) == true
+}
+
+fun ConeKotlinType.isExtensionFunctionType(session: FirSession): Boolean {
+    val type = this.lowerBoundIfFlexible().fullyExpandedType(session)
+    return type.attributes.extensionFunctionType != null
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirTypeUtils.kt

@@ -86,6 +86,7 @@ fun List<FirAnnotationCall>.computeTypeAttributes(): ConeAttributes {
         when (type.lookupTag.classId) {
             CompilerConeAttributes.Exact.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.Exact
             CompilerConeAttributes.NoInfer.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.NoInfer
+            CompilerConeAttributes.ExtensionFunctionType.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.ExtensionFunctionType
         }
     }
     return ConeAttributes.create(attributes)

compiler/testData/diagnostics/tests/inline/labeled.fir.kt

@@ -2,11 +2,11 @@
 
 inline fun foo(bar1: (String.() -> Int) -> Int, bar2: (()->Int) -> Int) {
     bar1 label@ {
-        this@label.<!UNRESOLVED_REFERENCE!>length<!>
+        this@label.length
     }
 
     bar1 {
-        this.<!UNRESOLVED_REFERENCE!>length<!>
+        this.length
     }
     //unmute after KT-4247 fix
     //bar1  {

compiler/testData/diagnostics/tests/labels/kt4247.fir.kt

@@ -1,11 +0,0 @@
-//KT-4247 LABEL_NAME_CLASH
-
-fun foo(bar1: (String.() -> Int) -> Int) {
-    bar1 {
-        this.<!UNRESOLVED_REFERENCE!>length<!>
-    }
-
-    bar1 {
-        this@bar1.<!UNRESOLVED_REFERENCE!>length<!>
-    }
-}

compiler/testData/diagnostics/tests/labels/kt4247.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 //KT-4247 LABEL_NAME_CLASH
 
 fun foo(bar1: (String.() -> Int) -> Int) {