LL API: remove old resolve api

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

@@ -19,8 +19,6 @@ import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeMappingMode
 import org.jetbrains.kotlin.analysis.api.withValidityAssertion
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDeclaration
 import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
 import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.codegen.signature.BothSignatureWriter
@@ -31,9 +29,6 @@ import org.jetbrains.kotlin.descriptors.java.JavaVisibilities
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
 import org.jetbrains.kotlin.fir.backend.jvm.jvmTypeMapper
-import org.jetbrains.kotlin.fir.declarations.FirClass
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.declarations.utils.superConeTypes
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -61,7 +56,6 @@ internal class KtFirPsiTypeProvider(
     ): PsiType? = withValidityAssertion {
         type.coneType.asPsiType(
             rootModuleSession,
-            analysisSession.firResolveState,
             mode.toTypeMappingMode(type, isAnnotationMethod),
             useSitePosition
         )
@@ -85,10 +79,9 @@ internal class KtFirPsiTypeProvider(
 
 private fun ConeKotlinType.simplifyType(
     session: FirSession,
-    state: FirModuleResolveState,
     useSitePosition: PsiElement,
 ): ConeKotlinType {
-    val substitutor = AnonymousTypesSubstitutor(session, state)
+    val substitutor = AnonymousTypesSubstitutor(session)
     val visibilityForApproximation = useSitePosition.visibilityForApproximation
     // TODO: See if the given [useSitePosition] is an `inline` method
     val isInlineFunction = false
@@ -181,11 +174,10 @@ private fun ConeKotlinType.isLocalButAvailableAtPosition(
 
 internal fun ConeKotlinType.asPsiType(
     session: FirSession,
-    state: FirModuleResolveState,
     mode: TypeMappingMode,
     useSitePosition: PsiElement,
 ): PsiType? {
-    val correctedType = simplifyType(session, state, useSitePosition)
+    val correctedType = simplifyType(session, useSitePosition)
 
     if (correctedType is ConeClassErrorType || correctedType !is SimpleTypeMarker) return null
 
@@ -213,7 +205,6 @@ internal fun ConeKotlinType.asPsiType(
 
 private class AnonymousTypesSubstitutor(
     private val session: FirSession,
-    private val state: FirModuleResolveState,
 ) : AbstractConeSubstitutor(session.typeContext) {
     override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
         if (type !is ConeClassLikeType) return null
@@ -222,18 +213,14 @@ private class AnonymousTypesSubstitutor(
         if (!isAnonymous) return null
 
         fun ConeClassLikeType.isNotInterface(): Boolean {
-            val firClassNode = lookupTag.toSymbol(session)?.fir as? FirClass ?: return false
-            return firClassNode.withFirDeclaration(state) { firSuperClass ->
-                firSuperClass.classKind != ClassKind.INTERFACE
-            }
+            val firClassNode = lookupTag.toSymbol(session) as? FirClassSymbol<*> ?: return false
+            return firClassNode.classKind != ClassKind.INTERFACE
         }
 
-        val firClassNode = (type.lookupTag.toSymbol(session) as? FirClassSymbol)?.fir
+        val firClassNode = type.lookupTag.toSymbol(session) as? FirClassSymbol
         if (firClassNode != null) {
-            val superTypesCones = firClassNode.withFirDeclaration(state, FirResolvePhase.SUPER_TYPES) {
-                (it as? FirClass)?.superConeTypes
-            }
-            val superClass = superTypesCones?.firstOrNull { it.isNotInterface() }
+            val superTypesCones = firClassNode.resolvedSuperTypes
+            val superClass = superTypesCones.firstOrNull { (it as? ConeClassLikeType)?.isNotInterface() == true }
             if (superClass != null) return superClass
         }
 

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbolProvider.kt

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.analysis.api.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirSymbolOfType
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbolOfType
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
@@ -40,12 +40,14 @@ internal class KtFirSymbolProvider(
                 "Creating KtValueParameterSymbol for function type parameter is not possible. Please see the KDoc of getParameterSymbol"
             )
 
-            psi.isLoopParameter -> psi.withFirSymbolOfType<FirPropertySymbol, KtLocalVariableSymbol>(resolveState) {
-                firSymbolBuilder.variableLikeBuilder.buildLocalVariableSymbol(it)
+            psi.isLoopParameter -> {
+                firSymbolBuilder.variableLikeBuilder.buildLocalVariableSymbol(psi.resolveToFirSymbolOfType<FirPropertySymbol>(resolveState))
             }
 
-            else -> psi.withFirSymbolOfType<FirValueParameterSymbol, KtValueParameterSymbol>(resolveState) {
-                firSymbolBuilder.variableLikeBuilder.buildValueParameterSymbol(it)
+            else -> {
+                firSymbolBuilder.variableLikeBuilder.buildValueParameterSymbol(
+                    psi.resolveToFirSymbolOfType<FirValueParameterSymbol>(resolveState)
+                )
             }
         }
     }
@@ -55,95 +57,76 @@ internal class KtFirSymbolProvider(
     }
 
     override fun getFunctionLikeSymbol(psi: KtNamedFunction): KtFunctionLikeSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirFunctionSymbol<*>, KtFunctionLikeSymbol>(resolveState) { fir ->
-            when (fir) {
-                is FirNamedFunctionSymbol -> {
-                    if (fir.origin == FirDeclarationOrigin.SamConstructor) {
-                        firSymbolBuilder.functionLikeBuilder.buildSamConstructorSymbol(fir)
-                    } else {
-                        firSymbolBuilder.functionLikeBuilder.buildFunctionSymbol(fir)
-                    }
+        when (val firSymbol = psi.resolveToFirSymbolOfType<FirFunctionSymbol<*>>(resolveState)) {
+            is FirNamedFunctionSymbol -> {
+                if (firSymbol.origin == FirDeclarationOrigin.SamConstructor) {
+                    firSymbolBuilder.functionLikeBuilder.buildSamConstructorSymbol(firSymbol)
+                } else {
+                    firSymbolBuilder.functionLikeBuilder.buildFunctionSymbol(firSymbol)
                 }
-                is FirAnonymousFunctionSymbol -> firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(fir)
-                else -> error("Unexpected ${fir.fir.renderWithType()}")
             }
+            is FirAnonymousFunctionSymbol -> firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(firSymbol)
+            else -> error("Unexpected ${firSymbol.fir.renderWithType()}")
         }
     }
 
     override fun getConstructorSymbol(psi: KtConstructor<*>): KtConstructorSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirConstructorSymbol, KtConstructorSymbol>(resolveState) {
-            firSymbolBuilder.functionLikeBuilder.buildConstructorSymbol(it)
-        }
+
+        firSymbolBuilder.functionLikeBuilder.buildConstructorSymbol(psi.resolveToFirSymbolOfType<FirConstructorSymbol>(resolveState))
+
     }
 
     override fun getTypeParameterSymbol(psi: KtTypeParameter): KtTypeParameterSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirTypeParameterSymbol, KtTypeParameterSymbol>(resolveState) {
-            firSymbolBuilder.classifierBuilder.buildTypeParameterSymbol(it)
-        }
+        firSymbolBuilder.classifierBuilder.buildTypeParameterSymbol(psi.resolveToFirSymbolOfType<FirTypeParameterSymbol>(resolveState))
     }
 
     override fun getTypeAliasSymbol(psi: KtTypeAlias): KtTypeAliasSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirTypeAliasSymbol, KtTypeAliasSymbol>(resolveState) {
-            firSymbolBuilder.classifierBuilder.buildTypeAliasSymbol(it)
-        }
+        firSymbolBuilder.classifierBuilder.buildTypeAliasSymbol(psi.resolveToFirSymbolOfType<FirTypeAliasSymbol>(resolveState))
     }
 
     override fun getEnumEntrySymbol(psi: KtEnumEntry): KtEnumEntrySymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirEnumEntrySymbol, KtEnumEntrySymbol>(resolveState) {
-            firSymbolBuilder.buildEnumEntrySymbol(it)
-        }
+        firSymbolBuilder.buildEnumEntrySymbol(psi.resolveToFirSymbolOfType<FirEnumEntrySymbol>(resolveState))
     }
 
     override fun getAnonymousFunctionSymbol(psi: KtNamedFunction): KtAnonymousFunctionSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirNamedFunctionSymbol, KtFunctionSymbol>(resolveState) {
-            firSymbolBuilder.functionLikeBuilder.buildFunctionSymbol(it)
-        }
         firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(psi.getOrBuildFirOfType(resolveState))
     }
 
     override fun getAnonymousFunctionSymbol(psi: KtFunctionLiteral): KtAnonymousFunctionSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirAnonymousFunctionSymbol, KtAnonymousFunctionSymbol>(resolveState) {
-            firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(it)
-        }
+        firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(
+            psi.resolveToFirSymbolOfType<FirAnonymousFunctionSymbol>(resolveState)
+        )
     }
 
     override fun getVariableSymbol(psi: KtProperty): KtVariableSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirPropertySymbol, KtVariableSymbol>(resolveState) {
-            firSymbolBuilder.variableLikeBuilder.buildVariableSymbol(it)
-        }
+        firSymbolBuilder.variableLikeBuilder.buildVariableSymbol(psi.resolveToFirSymbolOfType<FirPropertySymbol>(resolveState))
     }
 
     override fun getAnonymousObjectSymbol(psi: KtObjectLiteralExpression): KtAnonymousObjectSymbol = withValidityAssertion {
-        psi.objectDeclaration.withFirSymbolOfType<FirAnonymousObjectSymbol, KtAnonymousObjectSymbol>(resolveState) {
-            firSymbolBuilder.classifierBuilder.buildAnonymousObjectSymbol(it)
-        }
+        firSymbolBuilder.classifierBuilder.buildAnonymousObjectSymbol(
+            psi.objectDeclaration.resolveToFirSymbolOfType<FirAnonymousObjectSymbol>(resolveState)
+        )
     }
 
     override fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirClassSymbol<*>, KtClassOrObjectSymbol>(resolveState) {
-            firSymbolBuilder.classifierBuilder.buildClassOrObjectSymbol(it)
-        }
+        firSymbolBuilder.classifierBuilder.buildClassOrObjectSymbol(psi.resolveToFirSymbolOfType<FirClassSymbol<*>>(resolveState))
     }
 
     override fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol? = withValidityAssertion {
         require(psi !is KtObjectDeclaration || psi.parent !is KtObjectLiteralExpression)
         // A KtClassOrObject may also map to an FirEnumEntry. Hence, we need to return null in this case.
         if (psi is KtEnumEntry) return null
-        psi.withFirSymbolOfType<FirRegularClassSymbol, KtNamedClassOrObjectSymbol>(resolveState) {
-            firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(it)
-        }
+        firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(psi.resolveToFirSymbolOfType<FirRegularClassSymbol>(resolveState))
     }
 
     override fun getPropertyAccessorSymbol(psi: KtPropertyAccessor): KtPropertyAccessorSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirPropertyAccessorSymbol, KtPropertyAccessorSymbol>(resolveState) {
-            firSymbolBuilder.callableBuilder.buildPropertyAccessorSymbol(it)
-        }
+        firSymbolBuilder.callableBuilder.buildPropertyAccessorSymbol(psi.resolveToFirSymbolOfType<FirPropertyAccessorSymbol>(resolveState))
     }
 
     override fun getClassInitializerSymbol(psi: KtClassInitializer): KtClassInitializerSymbol = withValidityAssertion {
-        psi.withFirSymbolOfType<FirAnonymousInitializerSymbol, KtClassInitializerSymbol>(resolveState) {
-            firSymbolBuilder.anonymousInitializerBuilder.buildClassInitializer(it)
-        }
+        firSymbolBuilder.anonymousInitializerBuilder.buildClassInitializer(
+            psi.resolveToFirSymbolOfType<FirAnonymousInitializerSymbol>(resolveState)
+        )
     }
 
     override fun getClassOrObjectSymbolByClassId(classId: ClassId): KtClassOrObjectSymbol? = withValidityAssertion {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/utils/firUtils.kt

@@ -11,8 +11,6 @@ import org.jetbrains.kotlin.analysis.api.KtNonConstantInitializerValue
 import org.jetbrains.kotlin.analysis.api.fir.evaluate.FirCompileTimeConstantEvaluator
 import org.jetbrains.kotlin.analysis.api.fir.getCandidateSymbols
 import org.jetbrains.kotlin.analysis.api.types.KtTypeNullability
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDeclaration
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.classKind

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirModuleResolveStateDepended.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.analysis.api.impl.barebone.annotations.InternalForInline
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
@@ -15,12 +14,13 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.KtToFirMap
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingKtFileIfAny
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalKtFile
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -78,15 +78,15 @@ internal class FirModuleResolveStateDepended(
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
         TODO("Diagnostics are not implemented for depended state")
 
-    @OptIn(InternalForInline::class)
-    override fun findSourceFirDeclaration(ktDeclaration: KtLambdaExpression): FirDeclaration =
-        originalState.findSourceFirDeclaration(ktDeclaration)
+    override fun findSourceFirSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
+        return originalState.findSourceFirSymbol(ktDeclaration)
+    }
 
-    @OptIn(InternalForInline::class)
-    override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration =
-        originalState.findSourceFirDeclaration(ktDeclaration)
+    override fun findSourceFirSymbol(ktDeclaration: KtLambdaExpression): FirBasedSymbol<*> {
+        return originalState.findSourceFirSymbol(ktDeclaration)
+    }
 
-    @OptIn(InternalForInline::class)
-    override fun findSourceFirCompiledDeclaration(ktDeclaration: KtDeclaration) =
-        originalState.findSourceFirCompiledDeclaration(ktDeclaration)
+    override fun findSourceFirCompiledSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
+        return originalState.findSourceFirCompiledSymbol(ktDeclaration)
+    }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirModuleResolveStateImpl.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.analysis.api.impl.barebone.annotations.InternalForInline
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.DiagnosticsCollector
@@ -36,6 +35,7 @@ import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousObjectExpression
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.psi.*
 
 internal class FirModuleResolveStateImpl(
@@ -82,14 +82,32 @@ internal class FirModuleResolveStateImpl(
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
         diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
 
-    @OptIn(InternalForInline::class)
-    override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
-        return findSourceFirDeclarationByExpression(ktDeclaration.originalDeclaration ?: ktDeclaration)
+    override fun findSourceFirSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
+        return findSourceFirDeclarationByExpression(ktDeclaration.originalDeclaration ?: ktDeclaration).symbol
     }
 
-    @OptIn(InternalForInline::class)
-    override fun findSourceFirDeclaration(ktDeclaration: KtLambdaExpression): FirDeclaration =
-        findSourceFirDeclarationByExpression(ktDeclaration)
+    override fun findSourceFirSymbol(ktDeclaration: KtLambdaExpression): FirBasedSymbol<*> {
+        return findSourceFirDeclarationByExpression(ktDeclaration).symbol
+    }
+
+    override fun findSourceFirCompiledSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
+        require(ktDeclaration.containingKtFile.isCompiled) {
+            "This method will only work on compiled declarations, but this declaration is not compiled: ${ktDeclaration.getElementTextInContext()}"
+        }
+
+        val searcher = FirDeclarationForCompiledElementSearcher(rootModuleSession.symbolProvider)
+
+        val firDeclaration = when (ktDeclaration) {
+            is KtEnumEntry -> searcher.findNonLocalEnumEntry(ktDeclaration)
+            is KtClassOrObject -> searcher.findNonLocalClass(ktDeclaration)
+            is KtConstructor<*> -> searcher.findConstructorOfNonLocalClass(ktDeclaration)
+            is KtNamedFunction -> searcher.findNonLocalFunction(ktDeclaration)
+            is KtProperty -> searcher.findNonLocalProperty(ktDeclaration)
+
+            else -> error("Unsupported compiled declaration of type ${ktDeclaration::class}: ${ktDeclaration.getElementTextInContext()}")
+        }
+        return firDeclaration.symbol
+    }
 
     /**
      * [ktDeclaration] should be either [KtDeclaration] or [KtLambdaExpression]
@@ -118,32 +136,13 @@ internal class FirModuleResolveStateImpl(
         }
     }
 
-    @OptIn(InternalForInline::class)
-    override fun findSourceFirCompiledDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
-        require(ktDeclaration.containingKtFile.isCompiled) {
-            "This method will only work on compiled declarations, but this declaration is not compiled: ${ktDeclaration.getElementTextInContext()}"
-        }
-
-        val searcher = FirDeclarationForCompiledElementSearcher(rootModuleSession.symbolProvider)
-
-        return when (ktDeclaration) {
-            is KtEnumEntry -> searcher.findNonLocalEnumEntry(ktDeclaration)
-            is KtClassOrObject -> searcher.findNonLocalClass(ktDeclaration)
-            is KtConstructor<*> -> searcher.findConstructorOfNonLocalClass(ktDeclaration)
-            is KtNamedFunction -> searcher.findNonLocalFunction(ktDeclaration)
-            is KtProperty -> searcher.findNonLocalProperty(ktDeclaration)
-
-            else -> error("Unsupported compiled declaration of type ${ktDeclaration::class}: ${ktDeclaration.getElementTextInContext()}")
-        }
-    }
-
     override fun resolveFirToPhase(declaration: FirDeclaration, toPhase: FirResolvePhase) {
         if (toPhase == FirResolvePhase.RAW_FIR) return
         val fileCache = when (val session = declaration.moduleData.session) {
             is FirIdeSourcesSession -> session.cache
             else -> return
         }
-        firLazyDeclarationResolver. lazyResolveDeclaration(
+        firLazyDeclarationResolver.lazyResolveDeclaration(
             firDeclarationToResolve = declaration,
             moduleFileCache = fileCache,
             scopeSession = ScopeSession(),

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/FirModuleResolveState.kt

@@ -6,15 +6,18 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.api
 
 import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.analysis.api.impl.barebone.annotations.InternalForInline
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
+import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -50,26 +53,37 @@ abstract class FirModuleResolveState {
 
     internal abstract fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic>
 
-    @InternalForInline
-    abstract fun findSourceFirDeclaration(
+    abstract fun findSourceFirSymbol(
         ktDeclaration: KtDeclaration,
-    ): FirDeclaration
+    ): FirBasedSymbol<*>
 
-    @InternalForInline
-    abstract fun findSourceFirDeclaration(
+    abstract fun findSourceFirSymbol(
         ktDeclaration: KtLambdaExpression,
-    ): FirDeclaration
+    ): FirBasedSymbol<*>
 
     /**
      * Looks for compiled non-local [ktDeclaration] declaration by querying its classId/callableId from the SymbolProvider.
      *
      * Works only if [ktDeclaration] is compiled (i.e. comes from .class file).
      */
-    @InternalForInline
-    abstract fun findSourceFirCompiledDeclaration(
+    abstract fun findSourceFirCompiledSymbol(
         ktDeclaration: KtDeclaration
-    ): FirDeclaration
+    ): FirBasedSymbol<*>
 
 
+    open fun resolveToFirSymbol(
+        ktDeclaration: KtDeclaration,
+        resolveState: FirModuleResolveState,
+        phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
+    ): FirBasedSymbol<*> {
+        return when (val module = ktDeclaration.getKtModule()) {
+            is KtSourceModule -> resolveState.findSourceFirSymbol(ktDeclaration).also { resolveState.resolveFirToPhase(it.fir, phase) }
+            is KtLibraryModule -> resolveState.findSourceFirCompiledSymbol(ktDeclaration)
+            else -> error("unsupported module $module")
+        }
+    }
+
     internal abstract fun resolveFirToPhase(declaration: FirDeclaration, toPhase: FirResolvePhase)
+
+
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacade.kt

@@ -6,19 +6,18 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.api
 
 import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.analysis.api.impl.barebone.annotations.InternalForInline
 import org.jetbrains.kotlin.analysis.low.level.api.fir.FirIdeResolveStateService
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtLambdaExpression
 
 /**
  * Returns [FirModuleResolveState] which corresponds to containing module
@@ -36,79 +35,47 @@ fun KtModule.getResolveState(project: Project): FirModuleResolveState =
 
 
 /**
- * Creates [FirDeclaration] by [KtDeclaration] and executes an [action] on it
- * [FirDeclaration] passed to [action] will be resolved at least to [phase] when executing [action] on it
+ * Creates [FirBasedSymbol] by [KtDeclaration] .
+ * returned [FirDeclaration]  will be resolved at least to [phase]
  *
- * [FirDeclaration] passed to [action] should not be leaked outside [action] lambda
- * Otherwise, some threading problems may arise,
  */
-@OptIn(InternalForInline::class)
-inline fun <R> KtDeclaration.withFirSymbol(
+fun KtDeclaration.resolveToFirSymbol(
     resolveState: FirModuleResolveState,
     phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
-    action: (FirBasedSymbol<*>) -> R
-): R {
-    val firDeclaration = if (getKtModule(project) !is KtSourceModule) {
-        resolveState.findSourceFirCompiledDeclaration(this)
-    } else {
-        resolveState.findSourceFirDeclaration(this)
+): FirBasedSymbol<*> {
+    return resolveState.resolveToFirSymbol(this, resolveState, phase)
+}
+
+/**
+ * Creates [FirBasedSymbol] by [KtDeclaration] .
+ * returned [FirDeclaration] will be resolved at least to [phase]
+ *
+ * If resulted [FirBasedSymbol] is not subtype of [S], throws [InvalidFirElementTypeException]
+ */
+@Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
+inline fun <reified S : FirBasedSymbol<*>> KtDeclaration.resolveToFirSymbolOfType(
+    resolveState: FirModuleResolveState,
+    phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
+): @kotlin.internal.NoInfer S {
+    val symbol = resolveToFirSymbol(resolveState, phase)
+    if (symbol !is S) {
+        throwUnexpectedFirElementError(symbol, this, S::class)
     }
-
-    if (firDeclaration.resolvePhase < phase) {
-        firDeclaration.resolvedFirToPhase(phase, resolveState)
-    }
-
-    return action(firDeclaration.symbol)
+    return symbol
 }
 
 /**
- * Creates [FirDeclaration] by [KtDeclaration] and executes an [action] on it
- * [FirDeclaration] passed to [action] will be resolved at least to [phase] when executing [action] on it
+ * Creates [FirBasedSymbol] by [KtDeclaration] .
+ * returned [FirDeclaration] will be resolved at least to [phase]
  *
- * If resulted [FirDeclaration] is not [F] throws [InvalidFirElementTypeException]
- *
- * [FirDeclaration] passed to [action] should not be leaked outside [action] lambda
- * Otherwise, some threading problems may arise,
+ * If resulted [FirBasedSymbol] is not subtype of [S], returns `null`
  */
-@OptIn(InternalForInline::class)
-inline fun <reified F : FirBasedSymbol<*>, R> KtDeclaration.withFirSymbolOfType(
+@Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
+inline fun <reified S : FirBasedSymbol<*>> KtDeclaration.resolveToFirSymbolOfTypeSafe(
     resolveState: FirModuleResolveState,
     phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
-    action: (F) -> R
-): R = withFirSymbol(resolveState, phase) { firDeclaration ->
-    if (firDeclaration !is F) throwUnexpectedFirElementError(firDeclaration, this, F::class)
-    action(firDeclaration)
-}
-
-/**
- * Creates [FirDeclaration] by [KtLambdaExpression] and executes an [action] on it
- *
- * If resulted [FirDeclaration] is not [F] throws [InvalidFirElementTypeException]
- *
- * [FirDeclaration] passed to [action] should not be leaked outside [action] lambda
- * Otherwise, some threading problems may arise,
- */
-@OptIn(InternalForInline::class)
-inline fun <reified F : FirDeclaration, R> KtLambdaExpression.withFirDeclarationOfType(
-    resolveState: FirModuleResolveState,
-    action: (F) -> R
-): R {
-    val firDeclaration = resolveState.findSourceFirDeclaration(this)
-    if (firDeclaration !is F) throwUnexpectedFirElementError(firDeclaration, this, F::class)
-    return action(firDeclaration)
-}
-
-/**
- * Executes [action] with given [FirDeclaration] under read action, so resolve **is not possible** inside [action]
- * [FirDeclaration] passed to [action] will be resolved at least to [phase] when executing [action] on it
- */
-fun <D : FirDeclaration, R> D.withFirDeclaration(
-    resolveState: FirModuleResolveState,
-    phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
-    action: (D) -> R,
-): R {
-    resolvedFirToPhase(phase, resolveState)
-    return action(this)
+): @kotlin.internal.NoInfer S? {
+    return resolveToFirSymbol(resolveState, phase) as? S
 }
 
 
@@ -129,19 +96,6 @@ fun KtFile.collectDiagnosticsForFile(
 ): Collection<KtPsiDiagnostic> =
     resolveState.collectDiagnosticsForFile(this, filter)
 
-/**
- * Resolves a given [FirDeclaration] to [phase] and returns resolved declaration
- *
- * Should not be called form [withFirSymbol], [withFirDeclarationOfType] functions, as it it may cause deadlock
- */
-fun FirDeclaration.resolvedFirToPhase(
-    phase: FirResolvePhase,
-    resolveState: FirModuleResolveState
-) {
-    resolveState.resolveFirToPhase(this, phase)
-}
-
-
 /**
  * Get a [FirElement] which was created by [KtElement]
  * Returned [FirElement] is guaranteed to be resolved to [FirResolvePhase.BODY_RESOLVE] phase

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyDeclarationResolveTest.kt

@@ -5,17 +5,18 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirRenderer
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtVisitorVoid
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.directives.ConfigurationDirectives
 import org.jetbrains.kotlin.test.services.TestModuleStructure
@@ -58,21 +59,19 @@ abstract class AbstractFirLazyDeclarationResolveTest : AbstractLowLevelApiSingle
                 .findResolveMe()
             for (currentPhase in FirResolvePhase.values()) {
                 if (currentPhase.pluginPhase || currentPhase == FirResolvePhase.SEALED_CLASS_INHERITORS) continue
-                declarationToResolve.withFirDeclaration(firModuleResolveState, currentPhase) {
-                    val firFile = firModuleResolveState.getOrBuildFirFile(ktFile)
-                    resultBuilder.append("\n${currentPhase.name}:\n")
-                    resultBuilder.append(firFile.render(rendererOption))
-                }
+                declarationToResolve.ensureResolved(currentPhase)
+                val firFile = firModuleResolveState.getOrBuildFirFile(ktFile)
+                resultBuilder.append("\n${currentPhase.name}:\n")
+                resultBuilder.append(firFile.render(rendererOption))
             }
         }
 
         resolveWithClearCaches(ktFile) { firModuleResolveState ->
             check(firModuleResolveState is FirModuleResolveStateImpl)
             val firFile = firModuleResolveState.getOrBuildFirFile(ktFile)
-            firFile.withFirDeclaration(firModuleResolveState, FirResolvePhase.BODY_RESOLVE) {
-                resultBuilder.append("\nFILE RAW TO BODY:\n")
-                resultBuilder.append(firFile.render(rendererOption))
-            }
+            firFile.ensureResolved(FirResolvePhase.BODY_RESOLVE)
+            resultBuilder.append("\nFILE RAW TO BODY:\n")
+            resultBuilder.append(firFile.render(rendererOption))
         }
 
         testServices.assertions.assertEqualsToTestDataFileSibling(resultBuilder.toString())

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/compiler/based/LowLevelFirAnalyzerFacade.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.compiler.based
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.collectDiagnosticsForFile
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolvedFirToPhase
 import org.jetbrains.kotlin.fir.analysis.AbstractFirAnalyzerFacade
 import org.jetbrains.kotlin.diagnostics.KtDiagnostic
 import org.jetbrains.kotlin.fir.backend.Fir2IrResult
@@ -17,6 +16,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSealedClassInheritorsProcessor
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi2ir.generators.GeneratorExtensions
 import org.jetbrains.kotlin.test.model.TestFile
@@ -40,8 +40,8 @@ class LowLevelFirAnalyzerFacade(
     }
 
     private fun findSealedInheritors() {
-        allFirFiles.values.forEach {
-            it.resolvedFirToPhase(FirResolvePhase.SUPER_TYPES, resolveState)
+        allFirFiles.values.forEach { firFile ->
+            firFile.ensureResolved(FirResolvePhase.SUPER_TYPES)
         }
         val sealedProcessor = FirSealedClassInheritorsProcessor(allFirFiles.values.first().moduleData.session, ScopeSession())
         sealedProcessor.process(allFirFiles.values)