FIR IDE: collect diagnostics context under read lock

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

@@ -12,17 +12,28 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.fir.PersistenceContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.fir.PersistentCheckerContextFactory
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.LockProvider
 
 internal abstract class FileStructureElementDiagnosticRetriever {
-    abstract fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList
+    abstract fun retrieve(
+        firFile: FirFile,
+        collector: FileStructureElementDiagnosticsCollector,
+        lockProvider: LockProvider<FirFile>
+    ): FileStructureElementDiagnosticList
 }
 
-internal class SingleNonLocalDeclarationDiagnosticRetriever(
+internal class SingleNonLocalDeclarationDiagnosticRetriever (
     private val declaration: FirDeclaration
 ) : FileStructureElementDiagnosticRetriever() {
-    override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
+    override fun retrieve(
+        firFile: FirFile,
+        collector: FileStructureElementDiagnosticsCollector,
+        lockProvider: LockProvider<FirFile>
+    ): FileStructureElementDiagnosticList {
         val sessionHolder = SessionHolderImpl.createWithEmptyScopeSession(firFile.session)
-        val context = PersistenceContextCollector.collectContext(sessionHolder, firFile, declaration)
+        val context = lockProvider.withReadLock(firFile) {
+            PersistenceContextCollector.collectContext(sessionHolder, firFile, declaration)
+        }
         return collector.collectForStructureElement(declaration) { components ->
             FirIdeDiagnosticVisitor(context, components)
         }
@@ -30,7 +41,11 @@ internal class SingleNonLocalDeclarationDiagnosticRetriever(
 }
 
 internal object FileDiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
-    override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList =
+    override fun retrieve(
+        firFile: FirFile,
+        collector: FileStructureElementDiagnosticsCollector,
+        lockProvider: LockProvider<FirFile>
+    ): FileStructureElementDiagnosticList =
         collector.collectForStructureElement(firFile) { components ->
             Visitor(firFile, components)
         }

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

@@ -10,18 +10,19 @@ import com.intellij.util.SmartList
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
-import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureElement
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.LockProvider
 
 internal class FileStructureElementDiagnostics(
     private val firFile: FirFile,
+    private val lockProvider: LockProvider<FirFile>,
     private val retriever: FileStructureElementDiagnosticRetriever
 ) {
     private val diagnosticByCommonCheckers: FileStructureElementDiagnosticList by lazy {
-        retriever.retrieve(firFile, FileStructureElementDiagnosticsCollector.USUAL_COLLECTOR)
+        retriever.retrieve(firFile, FileStructureElementDiagnosticsCollector.USUAL_COLLECTOR, lockProvider)
     }
 
     private val diagnosticByExtendedCheckers: FileStructureElementDiagnosticList by lazy {
-        retriever.retrieve(firFile, FileStructureElementDiagnosticsCollector.EXTENDED_COLLECTOR)
+        retriever.retrieve(firFile, FileStructureElementDiagnosticsCollector.EXTENDED_COLLECTOR, lockProvider)
     }
 
     fun diagnosticsFor(filter: DiagnosticCheckerFilter, element: PsiElement): List<FirPsiDiagnostic<*>> =

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileElementFactory.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.LockProvider
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtProperty
@@ -21,25 +22,29 @@ internal object FileElementFactory {
         firDeclaration: FirDeclaration,
         ktDeclaration: KtDeclaration,
         firFile: FirFile,
+        firFileLockProvider: LockProvider<FirFile>,
     ): FileStructureElement = when {
         ktDeclaration is KtNamedFunction && ktDeclaration.isReanalyzableContainer() -> ReanalyzableFunctionStructureElement(
             firFile,
             ktDeclaration,
             (firDeclaration as FirSimpleFunction).symbol,
-            ktDeclaration.modificationStamp
+            ktDeclaration.modificationStamp,
+            firFileLockProvider,
         )
 
         ktDeclaration is KtProperty && ktDeclaration.isReanalyzableContainer() -> ReanalyzablePropertyStructureElement(
             firFile,
             ktDeclaration,
             (firDeclaration as FirProperty).symbol,
-            ktDeclaration.modificationStamp
+            ktDeclaration.modificationStamp,
+            firFileLockProvider,
         )
 
         else -> NonReanalyzableDeclarationStructureElement(
             firFile,
             firDeclaration,
             ktDeclaration,
+            firFileLockProvider,
         )
     }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructure.kt

@@ -115,7 +115,7 @@ internal class FileStructure(
             towerDataContextCollector = collector
         )
         return moduleFileCache.firFileLockProvider.withReadLock(firFile) {
-            FileElementFactory.createFileStructureElement(firDeclaration, declaration, firFile)
+            FileElementFactory.createFileStructureElement(firDeclaration, declaration, firFile, moduleFileCache.firFileLockProvider)
         }
     }
 
@@ -133,6 +133,7 @@ internal class FileStructure(
             RootStructureElement(
                 firFile,
                 container,
+                moduleFileCache.firFileLockProvider,
             )
         }
         is KtDeclaration -> createDeclarationStructure(container)

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt

@@ -15,22 +15,24 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileDiagnosticRet
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnostics
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.LockProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.hasFqName
 import org.jetbrains.kotlin.psi.*
 
-internal sealed class FileStructureElement(val firFile: FirFile) {
+internal sealed class FileStructureElement(val firFile: FirFile, protected val lockProvider: LockProvider<FirFile>) {
     abstract val psi: KtAnnotated
     abstract val mappings: Map<KtElement, FirElement>
     abstract val diagnostics: FileStructureElementDiagnostics
 }
 
-internal sealed class ReanalyzableStructureElement<KT : KtDeclaration, S: AbstractFirBasedSymbol<*>>(
+internal sealed class ReanalyzableStructureElement<KT : KtDeclaration, S : AbstractFirBasedSymbol<*>>(
     firFile: FirFile,
     protected val firSymbol: S,
-) : FileStructureElement(firFile) {
+    lockProvider: LockProvider<FirFile>,
+) : FileStructureElement(firFile, lockProvider) {
     abstract override val psi: KtDeclaration
     abstract val timestamp: Long
 
@@ -48,7 +50,7 @@ internal sealed class ReanalyzableStructureElement<KT : KtDeclaration, S: Abstra
 
     fun isUpToDate(): Boolean = psi.getModificationStamp() == timestamp
 
-    override val diagnostics = FileStructureElementDiagnostics(firFile, SingleNonLocalDeclarationDiagnosticRetriever(firSymbol.fir as FirDeclaration))
+    override val diagnostics = FileStructureElementDiagnostics(firFile, lockProvider, SingleNonLocalDeclarationDiagnosticRetriever(firSymbol.fir as FirDeclaration))
 
     companion object {
         val recorder = FirElementsRecorder()
@@ -59,8 +61,9 @@ internal class ReanalyzableFunctionStructureElement(
     firFile: FirFile,
     override val psi: KtNamedFunction,
     firSymbol: FirFunctionSymbol<*>,
-    override val timestamp: Long
+    override val timestamp: Long,
-) : ReanalyzableStructureElement<KtNamedFunction, FirFunctionSymbol<*>>(firFile, firSymbol) {
+    lockProvider: LockProvider<FirFile>,
+) : ReanalyzableStructureElement<KtNamedFunction, FirFunctionSymbol<*>>(firFile, firSymbol, lockProvider) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
@@ -89,6 +92,7 @@ internal class ReanalyzableFunctionStructureElement(
                     newKtDeclaration,
                     newFunction.symbol,
                     newKtDeclaration.modificationStamp,
+                    lockProvider,
                 )
             }
         }
@@ -99,8 +103,9 @@ internal class ReanalyzablePropertyStructureElement(
     firFile: FirFile,
     override val psi: KtProperty,
     firSymbol: FirPropertySymbol,
-    override val timestamp: Long
+    override val timestamp: Long,
-) : ReanalyzableStructureElement<KtProperty, FirPropertySymbol>(firFile, firSymbol) {
+    lockProvider: LockProvider<FirFile>,
+) : ReanalyzableStructureElement<KtProperty, FirPropertySymbol>(firFile, firSymbol, lockProvider) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
@@ -129,6 +134,7 @@ internal class ReanalyzablePropertyStructureElement(
                     newKtDeclaration,
                     newProperty.symbol,
                     newKtDeclaration.modificationStamp,
+                    lockProvider,
                 )
             }
         }
@@ -139,11 +145,12 @@ internal class NonReanalyzableDeclarationStructureElement(
     firFile: FirFile,
     fir: FirDeclaration,
     override val psi: KtDeclaration,
-) : FileStructureElement(firFile) {
+    lockProvider: LockProvider<FirFile>,
+) : FileStructureElement(firFile, lockProvider) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(fir, recorder)
 
-    override val diagnostics = FileStructureElementDiagnostics(firFile, SingleNonLocalDeclarationDiagnosticRetriever(fir))
+    override val diagnostics = FileStructureElementDiagnostics(firFile, lockProvider, SingleNonLocalDeclarationDiagnosticRetriever(fir))
 
 
     companion object {
@@ -169,11 +176,12 @@ internal class NonReanalyzableDeclarationStructureElement(
 internal class RootStructureElement(
     firFile: FirFile,
     override val psi: KtFile,
-) : FileStructureElement(firFile) {
+    lockProvider: LockProvider<FirFile>,
+) : FileStructureElement(firFile, lockProvider) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firFile, recorder)
 
-    override val diagnostics = FileStructureElementDiagnostics(firFile, FileDiagnosticRetriever)
+    override val diagnostics = FileStructureElementDiagnostics(firFile, lockProvider, FileDiagnosticRetriever)
 
     companion object {
         private val recorder = object : FirElementsRecorder() {