[LL FIR] create separate FileStructureElement for script

It is required to avoid duplication visiting on declarations

^KT-60728

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

@@ -9,12 +9,14 @@ 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.fir.PersistenceContextCollector
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.fir.PersistentCheckerContextFactory
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.visitScriptDependentElements
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContextForProvider
 import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorComponents
 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.FirValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
@@ -132,3 +134,33 @@ 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)
+            }
+        }
+    }
+
+    private class Visitor(
+        context: CheckerContextForProvider,
+        components: DiagnosticCollectorComponents,
+    ) : LLFirDiagnosticVisitor(context, components) {
+        override fun visitScript(script: FirScript, data: Nothing?) {
+            withAnnotationContainer(script) {
+                checkElement(script)
+                withDeclaration(script) {
+                    visitScriptDependentElements(script, this, data)
+                }
+            }
+        }
+    }
+}

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

@@ -54,6 +54,8 @@ internal object FileElementFactory {
             )
         }
 
+        ktDeclaration is KtScript -> RootScriptStructureElement(firFile, firDeclaration as FirScript, ktDeclaration, moduleComponents)
+
         else -> {
             NonReanalyzableNonClassDeclarationStructureElement(
                 firFile,

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

@@ -11,7 +11,9 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveCompone
 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.FileStructureElementDiagnostics
+import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.ScriptDiagnosticRetriever
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptStatement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.correspondingProperty
 import org.jetbrains.kotlin.fir.declarations.*
@@ -21,6 +23,7 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.psi.*
 
 internal sealed class FileStructureElement(val firFile: FirFile, protected val moduleComponents: LLFirModuleResolveComponents) {
@@ -123,6 +126,38 @@ internal sealed class NonReanalyzableDeclarationStructureElement(
     moduleComponents: LLFirModuleResolveComponents,
 ) : FileStructureElement(firFile, moduleComponents)
 
+internal class RootScriptStructureElement(
+    firFile: FirFile,
+    val script: FirScript,
+    override val psi: KtScript,
+    moduleComponents: LLFirModuleResolveComponents,
+) : NonReanalyzableDeclarationStructureElement(firFile, moduleComponents) {
+    override val mappings: KtToFirMapping = KtToFirMapping(script, Recorder)
+
+    override val diagnostics: FileStructureElementDiagnostics = FileStructureElementDiagnostics(
+        firFile,
+        ScriptDiagnosticRetriever(script),
+        moduleComponents,
+    )
+
+    private object Recorder : FirElementsRecorder() {
+        override fun visitScript(script: FirScript, data: MutableMap<KtElement, FirElement>) {
+            cacheElement(script, data)
+            visitScriptDependentElements(script, this, data)
+        }
+    }
+}
+
+internal fun <T, R> visitScriptDependentElements(script: FirScript, visitor: FirVisitor<T, R>, data: R) {
+    script.annotations.forEach { it.accept(visitor, data) }
+    script.contextReceivers.forEach { it.accept(visitor, data) }
+    script.statements.forEach {
+        if (it.isScriptStatement) {
+            it.accept(visitor, data)
+        }
+    }
+}
+
 internal class NonReanalyzableClassDeclarationStructureElement(
     firFile: FirFile,
     val fir: FirRegularClass,
@@ -217,17 +252,14 @@ internal class RootStructureElement(
     override val psi: KtFile,
     moduleComponents: LLFirModuleResolveComponents,
 ) : FileStructureElement(firFile, moduleComponents) {
-    override val mappings = KtToFirMapping(firFile, recorder)
+    override val mappings = KtToFirMapping(firFile, Recorder)
 
-    override val diagnostics =
-        FileStructureElementDiagnostics(firFile, FileDiagnosticRetriever, moduleComponents)
+    override val diagnostics = FileStructureElementDiagnostics(firFile, FileDiagnosticRetriever, moduleComponents)
 
-    companion object {
-        private val recorder = object : FirElementsRecorder() {
-            override fun visitElement(element: FirElement, data: MutableMap<KtElement, FirElement>) {
-                if (element !is FirDeclaration || element is FirFile) {
-                    super.visitElement(element, data)
-                }
+    private object Recorder : FirElementsRecorder() {
+        override fun visitElement(element: FirElement, data: MutableMap<KtElement, FirElement>) {
+            if (element !is FirDeclaration || element is FirFile) {
+                super.visitElement(element, data)
             }
         }
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FirElementsRecorder.kt

@@ -119,7 +119,7 @@ internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement,
         userTypeRef.acceptChildren(this, data)
     }
 
-    private fun cacheElement(element: FirElement, cache: MutableMap<KtElement, FirElement>) {
+    protected fun cacheElement(element: FirElement, cache: MutableMap<KtElement, FirElement>) {
         val psi = element.source
             ?.takeIf {
                 it is KtRealPsiSourceElement ||

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

@@ -53,7 +53,12 @@ internal fun KtDeclaration.findSourceNonLocalFirDeclaration(firFile: FirFile, pr
                     } else {
                         if (declaration.containingKtFile.isScript()) {
                             // .kts will have a single [FirScript] as a declaration. We need to unwrap statements in it.
-                            (firFile.declarations.singleOrNull() as? FirScript)?.statements?.filterIsInstance<FirDeclaration>()
+                            val firScript = firFile.declarations.singleOrNull() as? FirScript
+                            if (declaration is KtScript) {
+                                return@findSourceNonLocalFirDeclarationByProvider firScript?.takeIf { it.psi == declaration }
+                            }
+
+                            firScript?.statements?.filterIsInstance<FirDeclaration>()
                         } else {
                             firFile.declarations
                         }

analysis/low-level-api-fir/testdata/fileStructure/classMemberPropertyScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class X {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */class X {/* NonReanalyzableClassDeclarationStructureElement */
     var x: Int/* ReanalyzablePropertyStructureElement */
         get() = field
         set(value) {

analysis/low-level-api-fir/testdata/fileStructure/classScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
     fun x() {/* ReanalyzableFunctionStructureElement */
 
     }

analysis/low-level-api-fir/testdata/fileStructure/constructorsScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class A
+/* RootScriptStructureElement */class A
 
     (val a: Int)/* NonReanalyzableNonClassDeclarationStructureElement */
 

analysis/low-level-api-fir/testdata/fileStructure/danglingAnnotationClassLevelScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class Foo {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */class Foo {/* NonReanalyzableClassDeclarationStructureElement */
     @Suppress("") @MustBeDocumented
 }
 class Bar {/* NonReanalyzableClassDeclarationStructureElement */

analysis/low-level-api-fir/testdata/fileStructure/declarationsInPropertyInitScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
     val a = run {
         class X()
 

analysis/low-level-api-fir/testdata/fileStructure/enumClassScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */enum class A {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */enum class A {/* NonReanalyzableClassDeclarationStructureElement */
     X,/* NonReanalyzableNonClassDeclarationStructureElement */
     Y,/* NonReanalyzableNonClassDeclarationStructureElement */
     Z

analysis/low-level-api-fir/testdata/fileStructure/enumClassWithBodyScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */enum class A {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */enum class A {/* NonReanalyzableClassDeclarationStructureElement */
     X {/* NonReanalyzableNonClassDeclarationStructureElement */
         fun localInX() = 1
     },

analysis/low-level-api-fir/testdata/fileStructure/initBlockScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
     init {/* NonReanalyzableNonClassDeclarationStructureElement */
         val x = 10
         class B

analysis/low-level-api-fir/testdata/fileStructure/localClassScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */fun a() {/* ReanalyzableFunctionStructureElement */
+/* RootScriptStructureElement */fun a() {/* ReanalyzableFunctionStructureElement */
     class X
 }
 

analysis/low-level-api-fir/testdata/fileStructure/localFunScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */fun x() {/* ReanalyzableFunctionStructureElement */
+/* RootScriptStructureElement */fun x() {/* ReanalyzableFunctionStructureElement */
     fun y() {
 
     }

analysis/low-level-api-fir/testdata/fileStructure/localPropertyScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */fun foo() {/* ReanalyzableFunctionStructureElement */
+/* RootScriptStructureElement */fun foo() {/* ReanalyzableFunctionStructureElement */
     var x: Int
 }
 class A {/* NonReanalyzableClassDeclarationStructureElement */

analysis/low-level-api-fir/testdata/fileStructure/memberTypeAliasScript.kts

@@ -1,3 +1,3 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
     typealias X = Int/* NonReanalyzableNonClassDeclarationStructureElement */
 }

analysis/low-level-api-fir/testdata/fileStructure/nestedClassesScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */class A {/* NonReanalyzableClassDeclarationStructureElement */
     class B {/* NonReanalyzableClassDeclarationStructureElement */
         fun x() {/* ReanalyzableFunctionStructureElement */
         }

analysis/low-level-api-fir/testdata/fileStructure/propertyAccessorsScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */var x: Int = 10/* ReanalyzablePropertyStructureElement */
+/* RootScriptStructureElement */var x: Int = 10/* ReanalyzablePropertyStructureElement */
     get() = field
     set(value) {
         field = value

analysis/low-level-api-fir/testdata/fileStructure/script.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */class Builder {/* NonReanalyzableClassDeclarationStructureElement */
+/* RootScriptStructureElement */class Builder {/* NonReanalyzableClassDeclarationStructureElement */
     var version: String = ""/* NonReanalyzableNonClassDeclarationStructureElement */
 
     fun execute() {/* ReanalyzableFunctionStructureElement */

analysis/low-level-api-fir/testdata/fileStructure/superClassCallScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */open class A
+/* RootScriptStructureElement */open class A
     (init: A.() -> Unit)/* NonReanalyzableNonClassDeclarationStructureElement */
 {/* NonReanalyzableClassDeclarationStructureElement */
     val prop: String = ""/* NonReanalyzableNonClassDeclarationStructureElement */

analysis/low-level-api-fir/testdata/fileStructure/topLevelExpressionBodyFunWithTypeScript.kts

@@ -1 +1 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */fun foo(): Int = 42/* ReanalyzableFunctionStructureElement */
+/* RootScriptStructureElement */fun foo(): Int = 42/* ReanalyzableFunctionStructureElement */

analysis/low-level-api-fir/testdata/fileStructure/topLevelExpressionBodyFunWithoutTypeScript.kts

@@ -1 +1 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */fun foo() = 42/* NonReanalyzableNonClassDeclarationStructureElement */
+/* RootScriptStructureElement */fun foo() = 42/* NonReanalyzableNonClassDeclarationStructureElement */

analysis/low-level-api-fir/testdata/fileStructure/topLevelFunWithTypeScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */fun foo(): Int {/* ReanalyzableFunctionStructureElement */
+/* RootScriptStructureElement */fun foo(): Int {/* ReanalyzableFunctionStructureElement */
     println("")
     return 10
 }

analysis/low-level-api-fir/testdata/fileStructure/topLevelPropertyScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */var x: Int/* ReanalyzablePropertyStructureElement */
+/* RootScriptStructureElement */var x: Int/* ReanalyzablePropertyStructureElement */
     get() = field
     set(value) {
         field = value

analysis/low-level-api-fir/testdata/fileStructure/topLevelUnitFunScript.kts

@@ -1,3 +1,3 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */fun foo() {/* ReanalyzableFunctionStructureElement */
+/* RootScriptStructureElement */fun foo() {/* ReanalyzableFunctionStructureElement */
     println("")
 }

analysis/low-level-api-fir/testdata/fileStructure/withoutNameScript.kts

@@ -1,4 +1,4 @@
-/* NonReanalyzableNonClassDeclarationStructureElement */fun (a: Int = 1): String = "str"/* ReanalyzableFunctionStructureElement */
+/* RootScriptStructureElement */fun (a: Int = 1): String = "str"/* ReanalyzableFunctionStructureElement */
 
 fun () {/* ReanalyzableFunctionStructureElement */
 

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

@@ -68,6 +68,7 @@ abstract class AbstractFirContextCollectionTest : AbstractLowLevelApiSingleFileT
         is DanglingTopLevelModifierListStructureElement -> fir
         is NonReanalyzableClassDeclarationStructureElement -> fir
         is NonReanalyzableNonClassDeclarationStructureElement -> fir
+        is RootScriptStructureElement -> script
     }
 
     private class BeforeElementLLFirSessionConfigurator(private val testServices: TestServices) : LLFirSessionConfigurator {