FIR: Support SAM constructors for type aliases

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

@@ -14,9 +14,12 @@ import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
+import org.jetbrains.kotlin.fir.scopes.impl.FirClassSubstitutionScope
+import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
 import org.jetbrains.kotlin.fir.symbols.impl.*
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
-import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -79,12 +82,48 @@ private fun processSyntheticConstructors(
     processor: (FirFunctionSymbol<*>) -> ProcessorAction,
     bodyResolveComponents: BodyResolveComponents
 ): ProcessorAction {
-    if (matchedSymbol == null) return ProcessorAction.NEXT
+    val samConstructor = matchedSymbol.findSAMConstructor(bodyResolveComponents)
-    if (matchedSymbol !is FirRegularClassSymbol) return ProcessorAction.NEXT
+    if (samConstructor != null) return processor(samConstructor.symbol)
 
-    val function = bodyResolveComponents.samResolver.getSamConstructor(matchedSymbol.fir) ?: return ProcessorAction.NEXT
+    return ProcessorAction.NEXT
+}
 
-    return processor(function.symbol)
+private fun FirClassLikeSymbol<*>?.findSAMConstructor(
+    bodyResolveComponents: BodyResolveComponents
+): FirSimpleFunction? {
+    return when (this) {
+        is FirRegularClassSymbol -> bodyResolveComponents.samResolver.getSamConstructor(fir)
+        is FirTypeAliasSymbol -> findSAMConstructorForTypeAlias(bodyResolveComponents)
+        is FirAnonymousObjectSymbol, null -> null
+    }
+}
+
+private fun FirTypeAliasSymbol.findSAMConstructorForTypeAlias(
+    bodyResolveComponents: BodyResolveComponents
+): FirSimpleFunction? {
+    val session = bodyResolveComponents.session
+    val type =
+        fir.expandedTypeRef.coneTypeUnsafe<ConeClassLikeType>().fullyExpandedType(session)
+
+    val expansionRegularClass = type.lookupTag.toSymbol(session)?.fir as? FirRegularClass ?: return null
+    val samConstructorForClass = bodyResolveComponents.samResolver.getSamConstructor(expansionRegularClass) ?: return null
+
+    if (type.typeArguments.isEmpty()) return samConstructorForClass
+
+    val namedSymbol = samConstructorForClass.symbol as? FirNamedFunctionSymbol ?: return null
+
+    val substitutor = prepareSubstitutorForTypeAliasConstructors<FirSimpleFunction>(
+        this,
+        type,
+        session
+    ) { newReturnType, newParameterTypes, newTypeParameters ->
+        FirClassSubstitutionScope.createFakeOverrideFunction(
+            session, this, namedSymbol, null,
+            newReturnType, newParameterTypes, newTypeParameters
+        ).fir
+    } ?: return null
+
+    return substitutor.substitute(samConstructorForClass)
 }
 
 private fun processConstructors(
@@ -147,7 +186,8 @@ private class TypeAliasConstructorsSubstitutingScope(
     }
 }
 
-private typealias ConstructorCopyFactory<F> = F.(FirTypeRef, List<FirValueParameter>, List<FirTypeParameter>) -> F
+private typealias ConstructorCopyFactory<F> =
+        F.(newReturnType: ConeKotlinType?, newValueParameterTypes: List<ConeKotlinType?>, newTypeParameters: List<FirTypeParameter>) -> F
 
 private class TypeAliasConstructorsSubstitutor<F : FirMemberFunction<F>>(
     private val typeAliasSymbol: FirTypeAliasSymbol,
@@ -156,33 +196,17 @@ private class TypeAliasConstructorsSubstitutor<F : FirMemberFunction<F>>(
 ) {
     fun substitute(baseFunction: F): F {
         val typeParameters = typeAliasSymbol.fir.typeParameters
-        val newReturnTypeRef = baseFunction.returnTypeRef.substitute(substitutor)
+        val newReturnType = baseFunction.returnTypeRef.coneTypeUnsafe<ConeKotlinType>().let(substitutor::substituteOrNull)
 
-        val newParameterTypeRefs = baseFunction.valueParameters.map { valueParameter ->
+        val newParameterTypes = baseFunction.valueParameters.map { valueParameter ->
-            valueParameter.returnTypeRef.substitute(substitutor)
+            valueParameter.returnTypeRef.coneTypeUnsafe<ConeKotlinType>().let(substitutor::substituteOrNull)
         }
 
-        if (newReturnTypeRef == null && newParameterTypeRefs.all { it == null }) return baseFunction
+        if (newReturnType == null && newParameterTypes.all { it == null }) return baseFunction
 
         return baseFunction.copyFactory(
-            newReturnTypeRef ?: baseFunction.returnTypeRef,
+            newReturnType,
-            baseFunction.valueParameters.zip(
+            newParameterTypes,
-                newParameterTypeRefs
-            ) { valueParameter, newTypeRef ->
-                with(valueParameter) {
-                    FirValueParameterImpl(
-                        source,
-                        session,
-                        newTypeRef ?: returnTypeRef,
-                        name,
-                        FirVariableSymbol(valueParameter.symbol.callableId),
-                        defaultValue,
-                        isCrossinline,
-                        isNoinline,
-                        isVararg
-                    )
-                }
-            },
             typeParameters
         )
     }
@@ -199,13 +223,32 @@ private fun prepareSubstitutingScopeForTypeAliasConstructors(
             typeAliasSymbol,
             expandedType,
             session
-        ) factory@{ newReturnTypeRef, newValueParameters, newTypeParameters ->
+        ) factory@{ newReturnType, newParameterTypes, newTypeParameters ->
             FirConstructorImpl(
-                source, session, newReturnTypeRef, receiverTypeRef, status,
+                source, session,
+                returnTypeRef.withReplacedConeType(newReturnType),
+                receiverTypeRef, status,
                 FirConstructorSymbol(symbol.callableId, overriddenSymbol = symbol)
             ).apply {
                 resolvePhase = this@factory.resolvePhase
-                valueParameters += newValueParameters
+                valueParameters +=
+                    this@factory.valueParameters.zip(
+                        newParameterTypes
+                    ) { valueParameter, newParameterType ->
+                        with(valueParameter) {
+                            FirValueParameterImpl(
+                                source,
+                                session,
+                                returnTypeRef.withReplacedConeType(newParameterType),
+                                name,
+                                FirVariableSymbol(valueParameter.symbol.callableId),
+                                defaultValue,
+                                isCrossinline,
+                                isNoinline,
+                                isVararg
+                            )
+                        }
+                    }
                 this.typeParameters += newTypeParameters
             }
         } ?: return null
@@ -236,8 +279,3 @@ private fun <F : FirMemberFunction<F>> prepareSubstitutorForTypeAliasConstructor
 
     return TypeAliasConstructorsSubstitutor(typeAliasSymbol, substitutor, copyFactory)
 }
-
-private fun FirTypeRef.substitute(substitutor: ConeSubstitutor): FirResolvedTypeRef? =
-    coneTypeUnsafe<ConeKotlinType>()
-        .let(substitutor::substituteOrNull)
-        ?.let { FirResolvedTypeRefImpl(source, it) }

compiler/fir/resolve/testData/resolve/stdlib/kotlinComparatorAlias.kt

@@ -1,5 +1,5 @@
 fun test_1() {
-    val comp = <!UNRESOLVED_REFERENCE!>Comparator<!><Int> { x, y -> 1 }
+    val comp = Comparator<Int> { x, y -> 1 }
 }
 
 fun test_3(comparator: java.util.Comparator<Int>) {

compiler/fir/resolve/testData/resolve/stdlib/kotlinComparatorAlias.txt

@@ -1,6 +1,6 @@
-FILE: unresolvedComparator.kt
+FILE: kotlinComparatorAlias.kt
     public final fun test_1(): R|kotlin/Unit| {
-        lval comp: <ERROR TYPE REF: Unresolved name: Comparator> = <Unresolved name: Comparator>#<R|kotlin/Int|>(<L> = Comparator@fun <anonymous>(x: R|class error: No type for parameter|, y: R|class error: No type for parameter|): R|kotlin/Int| {
+        lval comp: R|java/util/Comparator<kotlin/Int>| = R|java/util/Comparator|<R|kotlin/Int|>(<L> = Comparator@fun <anonymous>(x: R|kotlin/Int|, y: R|kotlin/Int|): R|kotlin/Int| {
             Int(1)
         }
         )

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithStdlibTestGenerated.java

@@ -123,6 +123,11 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
         runTest("compiler/fir/resolve/testData/resolve/stdlib/javaLangComparator.kt");
     }
 
+    @TestMetadata("kotlinComparatorAlias.kt")
+    public void testKotlinComparatorAlias() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/stdlib/kotlinComparatorAlias.kt");
+    }
+
     @TestMetadata("listPlusAssign.kt")
     public void testListPlusAssign() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/stdlib/listPlusAssign.kt");
@@ -571,10 +576,5 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
         public void testCloneArray() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/stdlib/problems/cloneArray.kt");
         }
-
-        @TestMetadata("unresolvedComparator.kt")
-        public void testUnresolvedComparator() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolve/stdlib/problems/unresolvedComparator.kt");
-        }
     }
 }

compiler/testData/ir/irText/expressions/sam/samConstructors.fir.txt

@@ -20,14 +20,16 @@ FILE fqName:<root> fileName:/samConstructors.kt
       RETURN type=kotlin.Nothing from='public final fun test3 (): java.lang.Runnable declared in <root>'
         CALL 'public final fun Runnable (block: kotlin.Function0<kotlin.Unit>): java.lang.Runnable declared in java.lang' type=java.lang.Runnable origin=null
           block: FUNCTION_REFERENCE 'public final fun foo (): kotlin.Unit declared in <root>' type=kotlin.Function0<kotlin.Unit> origin=null
-  FUN name:test4 visibility:public modality:FINAL <> () returnType:IrErrorType
+  FUN name:test4 visibility:public modality:FINAL <> () returnType:java.util.Comparator<kotlin.Int>
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun test4 (): IrErrorType declared in <root>'
+      RETURN type=kotlin.Nothing from='public final fun test4 (): java.util.Comparator<kotlin.Int> declared in <root>'
-        ERROR_CALL 'Unresolved reference: <Unresolved name: Comparator>#' type=IrErrorType
+        CALL 'public final fun Comparator <T> (block: kotlin.Function2<T of java.util.Comparator?, T of java.util.Comparator?, kotlin.Int>): java.util.Comparator<T of java.util.Comparator> declared in java.util' type=java.util.Comparator<kotlin.Int> origin=null
-          FUN_EXPR type=kotlin.Function2<IrErrorType, IrErrorType, IrErrorType> origin=LAMBDA
+          <T>: kotlin.Int
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (a:IrErrorType, b:IrErrorType) returnType:IrErrorType
+          block: FUN_EXPR type=kotlin.Function2<kotlin.Int, kotlin.Int, kotlin.Int> origin=LAMBDA
-              VALUE_PARAMETER name:a index:0 type:IrErrorType
+            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (a:kotlin.Int, b:kotlin.Int) returnType:kotlin.Int
-              VALUE_PARAMETER name:b index:1 type:IrErrorType
+              VALUE_PARAMETER name:a index:0 type:kotlin.Int
+              VALUE_PARAMETER name:b index:1 type:kotlin.Int
               BLOCK_BODY
-                ERROR_CALL 'Unresolved reference: <Unresolved name: minus>#' type=IrErrorType
+                CALL 'public final fun minus (other: kotlin.Int): kotlin.Int [operator] declared in kotlin.Int' type=kotlin.Int origin=null
-                  GET_VAR 'b: IrErrorType declared in <root>.test4.<anonymous>' type=IrErrorType origin=null
+                  $this: GET_VAR 'a: kotlin.Int declared in <root>.test4.<anonymous>' type=kotlin.Int origin=null
+                  other: GET_VAR 'b: kotlin.Int declared in <root>.test4.<anonymous>' type=kotlin.Int origin=null