FIR: move typealias SAM constructor computation to FirSamResolver

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirTypeInfoProvider.kt

@@ -27,7 +27,7 @@ internal class KtFirTypeInfoProvider(
             firSession,
             analysisSession.getScopeSessionFor(firSession),
         )
-        return samResolver.getSamInfoForPossibleSamType(coneType) != null
+        return samResolver.isSamType(coneType)
     }
 
     override fun getFunctionClassKind(type: KtType): FunctionClassKind? {

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/AdapterGenerator.kt

@@ -485,7 +485,7 @@ internal class AdapterGenerator(
     }
 
     internal fun getFunctionTypeForPossibleSamType(parameterType: ConeKotlinType): ConeKotlinType? {
-        return samResolver.getSamInfoForPossibleSamType(parameterType)?.type
+        return samResolver.getFunctionTypeForPossibleSamType(parameterType)
     }
 
     /**

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SamResolution.kt

@@ -26,7 +26,9 @@ import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticFunctionSymbol
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
+import org.jetbrains.kotlin.fir.scopes.impl.FirFakeOverrideGenerator
 import org.jetbrains.kotlin.fir.scopes.impl.hasTypeOf
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
 import org.jetbrains.kotlin.fir.symbols.impl.*
@@ -40,12 +42,12 @@ import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.types.Variance
 
 abstract class FirSamResolver {
-    abstract fun getSamInfoForPossibleSamType(type: ConeKotlinType): SAMInfo<ConeKotlinType>?
     abstract fun shouldRunSamConversionForFunction(firFunction: FirFunction): Boolean
-    abstract fun getSamConstructor(firRegularClass: FirRegularClass): FirSimpleFunction?
+    abstract fun getSamConstructor(firClassOrTypeAlias: FirClassLikeDeclaration): FirSimpleFunction?
+    abstract fun getFunctionTypeForPossibleSamType(type: ConeKotlinType): ConeKotlinType?
 
-    fun getFunctionTypeForPossibleSamType(type: ConeKotlinType): ConeKotlinType? =
-        getSamInfoForPossibleSamType(type)?.type
+    fun isSamType(type: ConeKotlinType): Boolean =
+        getFunctionTypeForPossibleSamType(type) != null
 }
 
 private val SAM_PARAMETER_NAME = Name.identifier("function")
@@ -61,18 +63,15 @@ class FirSamResolverImpl(
     private val samConstructorsCache = session.samConstructorStorage.samConstructors
     private val samConversionTransformers = session.extensionService.samConversionTransformers
 
-    override fun getSamInfoForPossibleSamType(type: ConeKotlinType): SAMInfo<ConeKotlinType>? {
+    override fun getFunctionTypeForPossibleSamType(type: ConeKotlinType): ConeKotlinType? {
         return when (type) {
             is ConeClassLikeType -> getFunctionTypeForPossibleSamType(type.fullyExpandedType(session))
             is ConeFlexibleType -> {
-                val (lowerSymbol, lowerType) = getSamInfoForPossibleSamType(type.lowerBound) ?: return null
-                val (_, upperType) = getSamInfoForPossibleSamType(type.upperBound) ?: return null
-                SAMInfo(
-                    lowerSymbol,
-                    ConeFlexibleType(lowerType.lowerBoundIfFlexible(), upperType.upperBoundIfFlexible())
-                )
+                val lowerType = getFunctionTypeForPossibleSamType(type.lowerBound) ?: return null
+                val upperType = getFunctionTypeForPossibleSamType(type.upperBound) ?: return null
+                ConeFlexibleType(lowerType.lowerBoundIfFlexible(), upperType.upperBoundIfFlexible())
             }
-            is ConeErrorType, is ConeStubType -> null
+            is ConeStubType -> null
             // TODO: support those types as well
             is ConeTypeParameterType, is ConeTypeVariableType,
             is ConeCapturedType, is ConeDefinitelyNotNullType, is ConeIntersectionType,
@@ -83,55 +82,32 @@ class FirSamResolverImpl(
         }
     }
 
-    private fun getFunctionTypeForPossibleSamType(type: ConeClassLikeType): SAMInfo<ConeLookupTagBasedType>? {
+    private fun getFunctionTypeForPossibleSamType(type: ConeClassLikeType): ConeLookupTagBasedType? {
         @OptIn(LookupTagInternals::class)
         val firRegularClass = type.lookupTag.toFirRegularClass(session) ?: return null
 
-        val (functionSymbol, unsubstitutedFunctionType) = resolveFunctionTypeIfSamInterface(firRegularClass) ?: return null
+        val (_, unsubstitutedFunctionType) = resolveFunctionTypeIfSamInterface(firRegularClass) ?: return null
 
-        if (firRegularClass.typeParameters.isEmpty()) {
-            return SAMInfo(
-                functionSymbol,
-                unsubstitutedFunctionType.withNullability(ConeNullability.create(type.isMarkedNullable), session.typeContext)
-            )
+        val functionType = firRegularClass.buildSubstitutorForSamTypeAlias(session, type)?.substituteOrNull(unsubstitutedFunctionType)
+            ?: unsubstitutedFunctionType
+
+        require(functionType is ConeLookupTagBasedType) {
+            "Function type should always be ConeLookupTagBasedType, but ${functionType::class} was found"
         }
-
-        val substitutor =
-            substitutorByMap(
-                firRegularClass.typeParameters
-                    .map { it.symbol }
-                    .zip(
-                        type.typeArguments,
-                    ).map { (parameterSymbol, projection) ->
-                        val typeArgument =
-                            (projection as? ConeKotlinTypeProjection)?.type
-                            // TODO: Consider using `parameterSymbol.fir.bounds.first().coneType` once sure that it won't fail with exception
-                                ?: parameterSymbol.fir.bounds.firstOrNull()?.coneTypeSafe()
-                                ?: session.builtinTypes.nullableAnyType.type
-
-                        Pair(parameterSymbol, typeArgument)
-                    }
-                    .toMap(),
-                session
-            )
-
-        val result =
-            substitutor
-                .substituteOrSelf(unsubstitutedFunctionType)
-                .withNullability(ConeNullability.create(type.isMarkedNullable), session.typeContext)
-
-        require(result is ConeLookupTagBasedType) {
-            "Function type should always be ConeLookupTagBasedType, but ${result::class} was found"
-        }
-
-        return SAMInfo(functionSymbol, result)
+        return functionType.withNullability(ConeNullability.create(type.isMarkedNullable), session.typeContext)
     }
 
-    override fun getSamConstructor(firRegularClass: FirRegularClass): FirSimpleFunction? {
-        return samConstructorsCache.getValue(firRegularClass.symbol, this)?.fir
+    override fun getSamConstructor(firClassOrTypeAlias: FirClassLikeDeclaration): FirSimpleFunction? {
+        if (firClassOrTypeAlias is FirTypeAlias) {
+            // Precompute the constructor for the base type to avoid deadlocks in the IDE.
+            firClassOrTypeAlias.expandedTypeRef.coneTypeSafe<ConeClassLikeType>()
+                ?.fullyExpandedType(session)?.lookupTag?.toSymbol(session)
+                ?.let { samConstructorsCache.getValue(it, this) }
+        }
+        return samConstructorsCache.getValue(firClassOrTypeAlias.symbol, this)?.fir
     }
 
-    fun buildSamConstructor(classSymbol: FirRegularClassSymbol): FirNamedFunctionSymbol? {
+    fun buildSamConstructorForRegularClass(classSymbol: FirRegularClassSymbol): FirNamedFunctionSymbol? {
         val firRegularClass = classSymbol.fir
         val (functionSymbol, functionType) = resolveFunctionTypeIfSamInterface(firRegularClass) ?: return null
 
@@ -240,6 +216,39 @@ class FirSamResolverImpl(
         }.symbol
     }
 
+    fun buildSamConstructorForTypeAlias(typeAliasSymbol: FirTypeAliasSymbol): FirNamedFunctionSymbol? {
+        val type =
+            typeAliasSymbol.fir.expandedTypeRef.coneTypeUnsafe<ConeClassLikeType>().fullyExpandedType(session)
+
+        val expansionRegularClass = type.lookupTag.toSymbol(session)?.fir as? FirRegularClass ?: return null
+        val samConstructorForClass = getSamConstructor(expansionRegularClass) ?: return null
+
+        val substitutor = expansionRegularClass.buildSubstitutorForSamTypeAlias(session, type)
+            ?: return samConstructorForClass.symbol
+        val newReturnType = substitutor.substituteOrNull(samConstructorForClass.returnTypeRef.coneType)
+        val newParameterTypes = samConstructorForClass.valueParameters.map {
+            substitutor.substituteOrNull(it.returnTypeRef.coneType)
+        }
+        val newContextReceiverTypes = samConstructorForClass.contextReceivers.map {
+            substitutor.substituteOrNull(it.typeRef.coneType)
+        }
+        if (newReturnType == null && newParameterTypes.all { it == null } && newContextReceiverTypes.all { it == null }) {
+            return samConstructorForClass.symbol
+        }
+
+        val namedSymbol = samConstructorForClass.symbol
+        val symbolForOverride = FirFakeOverrideGenerator.createSymbolForSubstitutionOverride(namedSymbol, expansionRegularClass.classId)
+
+        return FirFakeOverrideGenerator.createSubstitutionOverrideFunction(
+            session, symbolForOverride, samConstructorForClass,
+            derivedClassLookupTag = null,
+            newDispatchReceiverType = null,
+            newReceiverType = null,
+            newContextReceiverTypes,
+            newReturnType, newParameterTypes, typeAliasSymbol.fir.typeParameters,
+        )
+    }
+
     private fun resolveFunctionTypeIfSamInterface(firRegularClass: FirRegularClass): SAMInfo<ConeLookupTagBasedType>? {
         return resolvedFunctionType.getOrPut(firRegularClass) {
             if (!firRegularClass.status.isFun) return@getOrPut null
@@ -260,6 +269,19 @@ class FirSamResolverImpl(
     }
 }
 
+private fun FirRegularClass.buildSubstitutorForSamTypeAlias(session: FirSession, type: ConeClassLikeType): ConeSubstitutor? {
+    if (typeParameters.isEmpty()) return null
+    val mapping = typeParameters.zip(type.typeArguments).associate { (parameter, projection) ->
+        val typeArgument =
+            (projection as? ConeKotlinTypeProjection)?.type
+            // TODO: Consider using `parameterSymbol.fir.bounds.first().coneType` once sure that it won't fail with exception
+                ?: parameter.symbol.fir.bounds.firstOrNull()?.coneTypeSafe()
+                ?: session.builtinTypes.nullableAnyType.type
+        Pair(parameter.symbol, typeArgument)
+    }
+    return substitutorByMap(mapping, session)
+}
+
 private fun FirRegularClass.getSingleAbstractMethodOrNull(
     session: FirSession,
     scopeSession: ScopeSession,
@@ -400,8 +422,14 @@ private fun FirSimpleFunction.getFunctionTypeForAbstractMethod(): ConeLookupTagB
 }
 
 class FirSamConstructorStorage(session: FirSession) : FirSessionComponent {
-    val samConstructors: FirCache<FirRegularClassSymbol, FirNamedFunctionSymbol?, FirSamResolverImpl> =
-        session.firCachesFactory.createCache { classSymbol, samResolver -> samResolver.buildSamConstructor(classSymbol) }
+    val samConstructors: FirCache<FirClassLikeSymbol<*>, FirNamedFunctionSymbol?, FirSamResolverImpl> =
+        session.firCachesFactory.createCache { classSymbol, samResolver ->
+            when (classSymbol) {
+                is FirRegularClassSymbol -> samResolver.buildSamConstructorForRegularClass(classSymbol)
+                is FirTypeAliasSymbol -> samResolver.buildSamConstructorForTypeAlias(classSymbol)
+                is FirAnonymousObjectSymbol -> null
+            }
+        }
 }
 
 private val FirSession.samConstructorStorage: FirSamConstructorStorage by FirSession.sessionComponentAccessor()

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

@@ -190,7 +190,7 @@ abstract class AbstractConeCallConflictResolver(
                     if (!call.usesSAM) {
                         TypeWithConversion(argumentType)
                     } else {
-                        val functionType = samResolver.getSamInfoForPossibleSamType(argumentType)?.type
+                        val functionType = samResolver.getFunctionTypeForPossibleSamType(argumentType)
                         if (functionType == null) TypeWithConversion(argumentType)
                         else TypeWithConversion(functionType, argumentType)
                     }

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

@@ -507,7 +507,7 @@ private fun Candidate.getExpectedTypeWithSAMConversion(
 
     // TODO: resolvedCall.registerArgumentWithSamConversion(argument, SamConversionDescription(convertedTypeByOriginal, convertedTypeByCandidate!!))
 
-    val (_, expectedFunctionType) = context.bodyResolveComponents.samResolver.getSamInfoForPossibleSamType(candidateExpectedType)
+    val expectedFunctionType = context.bodyResolveComponents.samResolver.getFunctionTypeForPossibleSamType(candidateExpectedType)
         ?: return null
     return runIf(argument.isFunctional(session, scopeSession, expectedFunctionType, context.returnTypeCalculator)) {
         usesSAM = true

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

@@ -11,15 +11,12 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildConstructedClassTypeParameterRef
 import org.jetbrains.kotlin.fir.declarations.builder.buildConstructorCopy
 import org.jetbrains.kotlin.fir.declarations.builder.buildReceiverParameter
-import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.declarations.utils.isInner
 import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
-import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.scopes.impl.FirFakeOverrideGenerator
 import org.jetbrains.kotlin.fir.scopes.processClassifiersByName
 import org.jetbrains.kotlin.fir.scopes.scopeForClass
 import org.jetbrains.kotlin.fir.symbols.impl.*
@@ -155,82 +152,12 @@ private fun processSyntheticConstructors(
     processor: (FirFunctionSymbol<*>) -> Unit,
     bodyResolveComponents: BodyResolveComponents
 ) {
-    val samConstructor = matchedSymbol.findSAMConstructor(bodyResolveComponents)
+    val samConstructor = bodyResolveComponents.samResolver.getSamConstructor(matchedSymbol.fir)
     if (samConstructor != null) {
         processor(samConstructor.symbol)
     }
 }
 
-private fun FirClassLikeSymbol<*>.findSAMConstructor(
-    bodyResolveComponents: BodyResolveComponents
-): FirSimpleFunction? {
-    return when (this) {
-        is FirRegularClassSymbol -> bodyResolveComponents.samResolver.getSamConstructor(fir)
-        is FirTypeAliasSymbol -> findSAMConstructorForTypeAlias(bodyResolveComponents)
-        is FirAnonymousObjectSymbol -> 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
-
-    val substitutor = prepareSubstitutorForTypeAliasConstructors(
-        type,
-        session
-    ) ?: return null
-
-    val typeParameters = this@findSAMConstructorForTypeAlias.fir.typeParameters
-    val newReturnType = samConstructorForClass.returnTypeRef.coneType.let(substitutor::substituteOrNull)
-
-    val newParameterTypes = samConstructorForClass.valueParameters.map { valueParameter ->
-        valueParameter.returnTypeRef.coneType.let(substitutor::substituteOrNull)
-    }
-
-    val newContextReceiverTypes = samConstructorForClass.contextReceivers.map { contextReceiver ->
-        contextReceiver.typeRef.coneType.let(substitutor::substituteOrNull)
-    }
-
-    if (newReturnType == null && newParameterTypes.all { it == null }) return samConstructorForClass
-
-    val symbolForOverride = FirFakeOverrideGenerator.createSymbolForSubstitutionOverride(namedSymbol, expansionRegularClass.classId)
-
-    return FirFakeOverrideGenerator.createSubstitutionOverrideFunction(
-        session, symbolForOverride, samConstructorForClass,
-        derivedClassLookupTag = null,
-        newDispatchReceiverType = null,
-        newReceiverType = null,
-        newContextReceiverTypes,
-        newReturnType, newParameterTypes, typeParameters,
-    ).fir
-}
-
-private fun prepareSubstitutorForTypeAliasConstructors(
-    expandedType: ConeClassLikeType,
-    session: FirSession
-): ConeSubstitutor? {
-    val expandedClass = expandedType.lookupTag.toSymbol(session)?.fir as? FirRegularClass ?: return null
-
-    val resultingTypeArguments = expandedType.typeArguments.map {
-        // We don't know how to handle cases like yet
-        // typealias A = ArrayList<*>()
-        it as? ConeKotlinType ?: return null
-    }
-    return substitutorByMap(
-        expandedClass.typeParameters.map { it.symbol }.zip(resultingTypeArguments).toMap(), session
-    )
-}
-
 private fun processConstructors(
     matchedSymbol: FirClassLikeSymbol<*>,
     substitutor: ConeSubstitutor,

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

@@ -588,7 +588,7 @@ class FirCallCompletionResultsWriterTransformer(
                         session.symbolProvider.getClassLikeSymbolByClassId(expectedArgumentType.lookupTag.classId)?.fir as? FirRegularClass
 
                     firRegularClass?.let answer@{
-                        val (_, functionType) = samResolver.getSamInfoForPossibleSamType(firRegularClass.defaultType())
+                        val functionType = samResolver.getFunctionTypeForPossibleSamType(firRegularClass.defaultType())
                             ?: return@answer null
                         createFunctionalType(
                             functionType.typeArguments.dropLast(1).map { it as ConeKotlinType },