FIR IDE: introduce scoped fir functions

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

@@ -9,12 +9,14 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 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.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
+import kotlin.reflect.KClass
 
 object LowLevelFirApiFacade {
     fun getResolveStateFor(element: KtElement): FirModuleResolveState =
@@ -26,12 +28,22 @@ object LowLevelFirApiFacade {
     fun getSessionFor(element: KtElement): FirSession =
         getResolveStateFor(element).getSessionFor(element.getModuleInfo())
 
+    @Deprecated("Consider using withFirElement")
     fun getOrBuildFirFor(element: KtElement, resolveState: FirModuleResolveState): FirElement =
         resolveState.getOrBuildFirFor(element)
 
+    @Suppress("DEPRECATION")
+    inline fun <R> withFirElement(element: KtElement, resolveState: FirModuleResolveState, action: (FirElement) -> R): R =
+        action(getOrBuildFirFor(element, resolveState))
+
+    @Deprecated("Consider using withFirFile")
     fun getFirFile(ktFile: KtFile, resolveState: FirModuleResolveState) =
         resolveState.getFirFile(ktFile)
 
+    @Suppress("DEPRECATION")
+    inline fun <R> withFirFile(ktFile: KtFile, resolveState: FirModuleResolveState, action: (FirFile) -> R): R =
+        action(getFirFile(ktFile, resolveState))
+
     fun getDiagnosticsFor(element: KtElement, resolveState: FirModuleResolveState): Collection<Diagnostic> =
         resolveState.getDiagnostics(element)
 
@@ -44,5 +56,48 @@ object LowLevelFirApiFacade {
         resolveState: FirModuleResolveState
     ): D =
         resolveState.resolvedFirToPhase(firDeclaration, phase)
-
 }
+
+@Deprecated("Consider using withFir")
+fun KtElement.getOrBuildFir(
+    resolveState: FirModuleResolveState,
+) = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState)
+
+@Deprecated("Consider using withFirSafe")
+inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
+    resolveState: FirModuleResolveState,
+) = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState) as? E
+
+@Deprecated("Consider using withFirOfType")
+inline fun <reified E : FirElement> KtElement.getOrBuildFirOfType(
+    resolveState: FirModuleResolveState,
+): E {
+    val fir = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState)
+    if (fir is E) return fir
+    throw InvalidFirElementTypeException(this, E::class, fir::class)
+}
+
+inline fun <R> KtElement.withFir(
+    resolveState: FirModuleResolveState,
+    action: (FirElement) -> R,
+) = LowLevelFirApiFacade.withFirElement(this, resolveState, action)
+
+
+inline fun <R : Any, reified E : FirElement> KtElement.withFirSafe(
+    resolveState: FirModuleResolveState,
+    action: (E) -> R?
+) = LowLevelFirApiFacade.withFirElement(this, resolveState) { element -> (element as? E)?.let(action) }
+
+
+@Suppress("DEPRECATION")
+inline fun <R, reified E : FirElement> KtElement.withFirOfType(
+    resolveState: FirModuleResolveState,
+    action: (E) -> R
+): R = action(getOrBuildFirOfType(resolveState))
+
+
+class InvalidFirElementTypeException(
+    ktElement: KtElement,
+    expectedFirClass: KClass<out FirElement>,
+    actualFirClass: KClass<out FirElement>
+) : IllegalStateException("For $ktElement with text `${ktElement.text}` the $expectedFirClass expected, but $actualFirClass found")

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt

@@ -24,30 +24,8 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import kotlin.reflect.KClass
 
-fun KtElement.getOrBuildFir(
-    resolveState: FirModuleResolveState,
-) = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState)
 
 
-inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
-    resolveState: FirModuleResolveState,
-) = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState) as? E
-
-inline fun <reified E : FirElement> KtElement.getOrBuildFirOfType(
-    resolveState: FirModuleResolveState,
-): E {
-    val fir = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState)
-    if (fir is E) return fir
-    throw InvalidFirElementTypeException(this, E::class, fir::class)
-}
-
-
-class InvalidFirElementTypeException(
-    ktElement: KtElement,
-    expectedFirClass: KClass<out FirElement>,
-    actualFirClass: KClass<out FirElement>
-) : IllegalStateException("For $ktElement with text `${ktElement.text}` the $expectedFirClass expected, but $actualFirClass found")
-
 
 fun FirFunctionCall.isImplicitFunctionCall(): Boolean {
     if (dispatchReceiver !is FirQualifiedAccessExpression) return false

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCallResolver.kt

@@ -14,9 +14,9 @@ import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.idea.fir.getOrBuildFir
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.fir.isImplicitFunctionCall
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.*
 import org.jetbrains.kotlin.idea.frontend.api.components.KtCallResolver
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt

@@ -9,9 +9,9 @@ import com.jetbrains.rd.util.getOrCreate
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForCompletion
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.ResolutionParameters
 import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolutionMode
 import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolver

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirSmartcastProvider.kt

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartCast
 import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartcastKind
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirTypeProvider.kt

@@ -10,8 +10,7 @@ import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.types.coneType
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartCast
 import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartcastKind
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken

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

@@ -7,8 +7,8 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirArrayAccessReference.kt

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.idea.references
 
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.idea.fir.getCalleeSymbol
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirConstructorDelegationReference.kt

@@ -6,8 +6,8 @@
 package org.jetbrains.kotlin.idea.references
 
 import org.jetbrains.kotlin.fir.expressions.FirDelegatedConstructorCall
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.fir.getResolvedKtSymbolOfNameReference
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirDestructuringDeclarationReference.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.idea.references
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.expressions.FirComponentCall
 import org.jetbrains.kotlin.idea.fir.getCalleeSymbol
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirForLoopInReference.kt

@@ -9,8 +9,8 @@ import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhileLoop
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.fir.getResolvedSymbolOfNameReference
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirPropertyDelegationMethodsReference.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol