[LL FIR] FileStructureElementDiagnosticRetriever: deduplicate code

2024-02-17 00:45:54 +01:00
parent 3f31bdb877
commit 2e17d8d524
3 changed files with 100 additions and 99 deletions
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -17,39 +17,49 @@ import org.jetbrains.kotlin.fir.correspondingProperty
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirScript
 import org.jetbrains.kotlin.fir.declarations.impl.FirPrimaryConstructor
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.util.withSourceCodeAnalysisExceptionUnwrapping

-internal abstract class FileStructureElementDiagnosticRetriever {
-    abstract fun retrieve(
-        firFile: FirFile,
-        collector: FileStructureElementDiagnosticsCollector,
-        moduleComponents: LLFirModuleResolveComponents,
-    ): FileStructureElementDiagnosticList
-}
-
-internal class ClassDiagnosticRetriever(
-    private val structureElementDeclaration: FirRegularClass
-) : FileStructureElementDiagnosticRetriever() {
-    override fun retrieve(
-        firFile: FirFile,
-        collector: FileStructureElementDiagnosticsCollector,
-        moduleComponents: LLFirModuleResolveComponents,
-    ): FileStructureElementDiagnosticList {
+internal sealed class FileStructureElementDiagnosticRetriever(
+    val declaration: FirDeclaration,
+    private val file: FirFile,
+    private val moduleComponents: LLFirModuleResolveComponents,
+) {
+    fun retrieve(collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
        val sessionHolder = SessionHolderImpl(moduleComponents.session, moduleComponents.scopeSessionProvider.getScopeSession())
-        val context = PersistenceContextCollector.collectContext(sessionHolder, firFile, structureElementDeclaration)
+        val context = if (declaration is FirFile) {
+            PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(sessionHolder)
+        } else {
+            PersistenceContextCollector.collectContext(sessionHolder, file, declaration)
+        }
+
        return withSourceCodeAnalysisExceptionUnwrapping {
-            collector.collectForStructureElement(structureElementDeclaration) { components ->
-                Visitor(structureElementDeclaration, context, components)
+            collector.collectForStructureElement(declaration) { components ->
+                createVisitor(context, components)
            }
        }
    }

+    abstract fun createVisitor(context: CheckerContextForProvider, components: DiagnosticCollectorComponents): LLFirDiagnosticVisitor
+}
+
+internal class ClassDiagnosticRetriever(
+    declaration: FirDeclaration,
+    file: FirFile,
+    moduleComponents: LLFirModuleResolveComponents,
+) : FileStructureElementDiagnosticRetriever(
+    declaration,
+    file,
+    moduleComponents,
+) {
+    override fun createVisitor(context: CheckerContextForProvider, components: DiagnosticCollectorComponents): LLFirDiagnosticVisitor {
+        return Visitor(declaration, context, components)
+    }
+
    private class Visitor(
-        private val structureElementDeclaration: FirRegularClass,
+        private val structureElementDeclaration: FirDeclaration,
        context: CheckerContextForProvider,
        components: DiagnosticCollectorComponents,
    ) : LLFirDiagnosticVisitor(context, components) {
@@ -89,25 +99,21 @@ internal val FirElement.isImplicitConstructor: Boolean
    get() = this is FirConstructor && source?.kind == KtFakeSourceElementKind.ImplicitConstructor

 internal class SingleNonLocalDeclarationDiagnosticRetriever(
-    private val structureElementDeclaration: FirDeclaration
-) : FileStructureElementDiagnosticRetriever() {
-    override fun retrieve(
-        firFile: FirFile,
-        collector: FileStructureElementDiagnosticsCollector,
-        moduleComponents: LLFirModuleResolveComponents,
-    ): FileStructureElementDiagnosticList {
-        val sessionHolder = SessionHolderImpl(moduleComponents.session, moduleComponents.scopeSessionProvider.getScopeSession())
-        val context = PersistenceContextCollector.collectContext(sessionHolder, firFile, structureElementDeclaration)
-        return withSourceCodeAnalysisExceptionUnwrapping {
-            collector.collectForStructureElement(structureElementDeclaration) { components ->
-                Visitor(context, components)
-            }
-        }
+    declaration: FirDeclaration,
+    file: FirFile,
+    moduleComponents: LLFirModuleResolveComponents,
+) : FileStructureElementDiagnosticRetriever(
+    declaration,
+    file,
+    moduleComponents,
+) {
+    override fun createVisitor(context: CheckerContextForProvider, components: DiagnosticCollectorComponents): LLFirDiagnosticVisitor {
+        return Visitor(context, components)
    }

    private class Visitor(
        context: CheckerContextForProvider,
-        components: DiagnosticCollectorComponents
+        components: DiagnosticCollectorComponents,
    ) : LLFirDiagnosticVisitor(context, components) {
        override fun visitConstructor(constructor: FirConstructor, data: Nothing?) {
            super.visitConstructor(constructor, data)
@@ -123,27 +129,23 @@ internal class SingleNonLocalDeclarationDiagnosticRetriever(
    }
 }

-internal object FileDiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
-    override fun retrieve(
-        firFile: FirFile,
-        collector: FileStructureElementDiagnosticsCollector,
-        moduleComponents: LLFirModuleResolveComponents,
-    ): FileStructureElementDiagnosticList =
-        withSourceCodeAnalysisExceptionUnwrapping {
-            collector.collectForStructureElement(firFile) { components ->
-                Visitor(components, moduleComponents)
-            }
-        }
+internal class FileDiagnosticRetriever(
+    declaration: FirDeclaration,
+    file: FirFile,
+    moduleComponents: LLFirModuleResolveComponents,
+) : FileStructureElementDiagnosticRetriever(
+    declaration,
+    file,
+    moduleComponents,
+) {
+    override fun createVisitor(context: CheckerContextForProvider, components: DiagnosticCollectorComponents): LLFirDiagnosticVisitor {
+        return Visitor(context, components)
+    }

    private class Visitor(
+        context: CheckerContextForProvider,
        components: DiagnosticCollectorComponents,
-        moduleComponents: LLFirModuleResolveComponents,
-    ) : LLFirDiagnosticVisitor(
-        PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(
-            SessionHolderImpl(moduleComponents.session, moduleComponents.scopeSessionProvider.getScopeSession())
-        ),
-        components,
-    ) {
+    ) : LLFirDiagnosticVisitor(context, components) {
        override fun visitFile(file: FirFile, data: Nothing?) {
            withAnnotationContainer(file) {
                visitWithFile(file) {
@@ -157,19 +159,17 @@ internal object FileDiagnosticRetriever : FileStructureElementDiagnosticRetrieve
    }
 }

-internal class ScriptDiagnosticRetriever(private val script: FirScript) : FileStructureElementDiagnosticRetriever() {
-    override fun retrieve(
-        firFile: FirFile,
-        collector: FileStructureElementDiagnosticsCollector,
-        moduleComponents: LLFirModuleResolveComponents,
-    ): FileStructureElementDiagnosticList {
-        val sessionHolder = SessionHolderImpl(moduleComponents.session, moduleComponents.scopeSessionProvider.getScopeSession())
-        val context = PersistenceContextCollector.collectContext(sessionHolder, firFile, script)
-        return withSourceCodeAnalysisExceptionUnwrapping {
-            collector.collectForStructureElement(script) { components ->
-                Visitor(context, components)
-            }
-        }
+internal class ScriptDiagnosticRetriever(
+    declaration: FirDeclaration,
+    file: FirFile,
+    moduleComponents: LLFirModuleResolveComponents,
+) : FileStructureElementDiagnosticRetriever(
+    declaration,
+    file,
+    moduleComponents,
+) {
+    override fun createVisitor(context: CheckerContextForProvider, components: DiagnosticCollectorComponents): LLFirDiagnosticVisitor {
+        return Visitor(context, components)
    }

    private class Visitor(
@@ -7,22 +7,16 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics

 import com.intellij.psi.PsiElement
 import com.intellij.util.SmartList
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
-import org.jetbrains.kotlin.fir.declarations.FirFile

-internal class FileStructureElementDiagnostics(
-    private val firFile: FirFile,
-    private val retriever: FileStructureElementDiagnosticRetriever,
-    private val moduleComponents: LLFirModuleResolveComponents,
-) {
+internal class FileStructureElementDiagnostics(private val retriever: FileStructureElementDiagnosticRetriever) {
    private val diagnosticByCommonCheckers: FileStructureElementDiagnosticList by lazy {
-        retriever.retrieve(firFile, FileStructureElementDiagnosticsCollector.USUAL_COLLECTOR, moduleComponents)
+        retriever.retrieve(FileStructureElementDiagnosticsCollector.USUAL_COLLECTOR)
    }

    private val diagnosticByExtendedCheckers: FileStructureElementDiagnosticList by lazy {
-        retriever.retrieve(firFile, FileStructureElementDiagnosticsCollector.EXTENDED_COLLECTOR, moduleComponents)
+        retriever.retrieve(FileStructureElementDiagnosticsCollector.EXTENDED_COLLECTOR)
    }

    fun diagnosticsFor(filter: DiagnosticCheckerFilter, element: PsiElement): List<KtPsiDiagnostic> =
@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.ClassDiagnosticRetriever
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.FileDiagnosticRetriever
-import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.FileStructureElementDiagnosticRetriever
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.FileStructureElementDiagnostics
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.ScriptDiagnosticRetriever
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
@@ -27,9 +26,7 @@ import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment

 internal sealed class FileStructureElement(
    val declaration: FirDeclaration,
-    firFile: FirFile,
-    retriever: FileStructureElementDiagnosticRetriever,
-    moduleComponents: LLFirModuleResolveComponents,
+    val diagnostics: FileStructureElementDiagnostics,
 ) {
    init {
        val actualResolvePhase = declaration.resolvePhase
@@ -48,12 +45,6 @@ internal sealed class FileStructureElement(

    val mappings: KtToFirMapping = KtToFirMapping(declaration)

-    val diagnostics: FileStructureElementDiagnostics = FileStructureElementDiagnostics(
-        firFile,
-        retriever,
-        moduleComponents,
-    )
-
    companion object {
        fun recorderFor(fir: FirDeclaration): FirElementsRecorder = when (fir) {
            is FirFile -> RootStructureElement.Recorder
@@ -146,9 +137,13 @@ internal class RootScriptStructureElement(
    moduleComponents: LLFirModuleResolveComponents,
 ) : FileStructureElement(
    declaration = script,
-    firFile = file,
-    retriever = ScriptDiagnosticRetriever(script),
-    moduleComponents = moduleComponents,
+    diagnostics = FileStructureElementDiagnostics(
+        ScriptDiagnosticRetriever(
+            declaration = script,
+            file = file,
+            moduleComponents = moduleComponents,
+        )
+    ),
 ) {
    object Recorder : FirElementsRecorder() {
        override fun visitScript(script: FirScript, data: MutableMap<KtElement, FirElement>) {
@@ -169,9 +164,13 @@ internal class ClassDeclarationStructureElement(
    moduleComponents: LLFirModuleResolveComponents,
 ) : FileStructureElement(
    declaration = clazz,
-    firFile = file,
-    retriever = ClassDiagnosticRetriever(clazz),
-    moduleComponents = moduleComponents
+    diagnostics = FileStructureElementDiagnostics(
+        ClassDiagnosticRetriever(
+            declaration = clazz,
+            file = file,
+            moduleComponents = moduleComponents,
+        )
+    ),
 ) {
    class Recorder(private val firClass: FirRegularClass) : FirElementsRecorder() {
        override fun visitProperty(property: FirProperty, data: MutableMap<KtElement, FirElement>) {
@@ -217,9 +216,13 @@ internal class DeclarationStructureElement(
    moduleComponents: LLFirModuleResolveComponents,
 ) : FileStructureElement(
    declaration = declaration,
-    firFile = file,
-    retriever = SingleNonLocalDeclarationDiagnosticRetriever(declaration),
-    moduleComponents = moduleComponents,
+    diagnostics = FileStructureElementDiagnostics(
+        SingleNonLocalDeclarationDiagnosticRetriever(
+            declaration = declaration,
+            file = file,
+            moduleComponents = moduleComponents,
+        )
+    ),
 ) {
    object Recorder : FirElementsRecorder() {
        override fun visitConstructor(constructor: FirConstructor, data: MutableMap<KtElement, FirElement>) {
@@ -241,9 +244,13 @@ internal class RootStructureElement(
    moduleComponents: LLFirModuleResolveComponents,
 ) : FileStructureElement(
    declaration = file,
-    firFile = file,
-    retriever = FileDiagnosticRetriever,
-    moduleComponents = moduleComponents
+    diagnostics = FileStructureElementDiagnostics(
+        FileDiagnosticRetriever(
+            declaration = file,
+            file = file,
+            moduleComponents = moduleComponents,
+        )
+    ),
 ) {
    object Recorder : FirElementsRecorder() {
        override fun visitElement(element: FirElement, data: MutableMap<KtElement, FirElement>) {