FIR: introduce package directive

The package directive is needed for KtElement -> FirElement mapping in IDE.
There should be some FirElement which will be result of mapping of elements
inside KtPackageDirective.

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirAnnotationChecker.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.isInner
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.argumentMapping
+import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.resolve.fqName
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.StandardClassIds

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.declarations.utils.isSynthetic
 import org.jetbrains.kotlin.fir.declarations.utils.primaryConstructor
 import org.jetbrains.kotlin.fir.descriptors.FirModuleDescriptor
+import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.signaturer.FirMangler

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaVisibilityChecker.kt

@@ -7,12 +7,9 @@ package org.jetbrains.kotlin.fir.java
 
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.descriptors.java.JavaVisibilities
-import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.FirVisibilityChecker
-import org.jetbrains.kotlin.fir.NoMutableState
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.getOwnerId
 import org.jetbrains.kotlin.fir.resolve.calls.ReceiverValue
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt

@@ -90,10 +90,13 @@ class DeclarationsConverter(
         val importList = mutableListOf<FirImport>()
         val firDeclarationList = mutableListOf<FirDeclaration>()
         context.packageFqName = FqName.ROOT
+        var packageDirective: FirPackageDirective? = null
         file.forEachChildren {
             when (it.tokenType) {
                 FILE_ANNOTATION_LIST -> fileAnnotationList += convertFileAnnotationList(it)
-                PACKAGE_DIRECTIVE -> context.packageFqName = convertPackageName(it)
+                PACKAGE_DIRECTIVE -> {
+                    packageDirective = convertPackageDirective(it).also { context.packageFqName = it.packageFqName }
+                }
                 IMPORT_LIST -> importList += convertImportDirectives(it)
                 CLASS -> firDeclarationList += convertClass(it)
                 FUN -> firDeclarationList += convertFunctionDeclaration(it) as FirDeclaration
@@ -108,7 +111,7 @@ class DeclarationsConverter(
             origin = FirDeclarationOrigin.Source
             moduleData = baseModuleData
             name = fileName
-            packageFqName = context.packageFqName
+            this.packageDirective = packageDirective ?: buildPackageDirective { packageFqName = context.packageFqName }
             annotations += fileAnnotationList
             imports += importList
             declarations += firDeclarationList
@@ -151,7 +154,7 @@ class DeclarationsConverter(
     /**
      * @see org.jetbrains.kotlin.parsing.KotlinParsing.parsePackageName
      */
-    private fun convertPackageName(packageNode: LighterASTNode): FqName {
+    private fun convertPackageDirective(packageNode: LighterASTNode): FirPackageDirective {
         var packageName: FqName = FqName.ROOT
         packageNode.forEachChildren {
             when (it.tokenType) {
@@ -159,7 +162,10 @@ class DeclarationsConverter(
                 DOT_QUALIFIED_EXPRESSION, REFERENCE_EXPRESSION -> packageName = FqName(it.getAsStringWithoutBacktick())
             }
         }
-        return packageName
+        return buildPackageDirective {
+            packageFqName = packageName
+            source = packageNode.toFirSourceElement()
+        }
     }
 
     private fun convertImportAlias(importAlias: LighterASTNode): String? {

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -740,7 +740,10 @@ open class RawFirBuilder(
                 moduleData = baseModuleData
                 origin = FirDeclarationOrigin.Source
                 name = file.name
-                packageFqName = context.packageFqName
+                packageDirective = buildPackageDirective {
+                    packageFqName = context.packageFqName
+                    source = file.packageDirective?.toFirPsiSourceElement()
+                }
                 for (annotationEntry in file.annotationEntries) {
                     annotations += annotationEntry.convert<FirAnnotationCall>()
                 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/Scopes.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.scopes
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.scopes.impl.*
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/FirPackageDirective.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2010-2021 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.
+ */
+
+package org.jetbrains.kotlin.fir
+
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.fir.visitors.*
+
+/*
+ * This file was generated automatically
+ * DO NOT MODIFY IT MANUALLY
+ */
+
+abstract class FirPackageDirective : FirPureAbstractElement(), FirElement {
+    abstract override val source: FirSourceElement?
+    abstract val packageFqName: FqName
+
+    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitPackageDirective(this, data)
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <E: FirElement, D> transform(transformer: FirTransformer<D>, data: D): E = 
+        transformer.transformPackageDirective(this, data) as E
+}

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/builder/FirPackageDirectiveBuilder.kt

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2021 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.
+ */
+
+package org.jetbrains.kotlin.fir.builder
+
+import kotlin.contracts.*
+import org.jetbrains.kotlin.fir.FirPackageDirective
+import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
+import org.jetbrains.kotlin.fir.impl.FirPackageDirectiveImpl
+import org.jetbrains.kotlin.fir.visitors.*
+import org.jetbrains.kotlin.name.FqName
+
+/*
+ * This file was generated automatically
+ * DO NOT MODIFY IT MANUALLY
+ */
+
+@FirBuilderDsl
+class FirPackageDirectiveBuilder {
+    var source: FirSourceElement? = null
+    lateinit var packageFqName: FqName
+
+    fun build(): FirPackageDirective {
+        return FirPackageDirectiveImpl(
+            source,
+            packageFqName,
+        )
+    }
+
+}
+
+@OptIn(ExperimentalContracts::class)
+inline fun buildPackageDirective(init: FirPackageDirectiveBuilder.() -> Unit): FirPackageDirective {
+    contract {
+        callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
+    }
+    return FirPackageDirectiveBuilder().apply(init).build()
+}

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirFile.kt

@@ -7,10 +7,10 @@ package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirModuleData
+import org.jetbrains.kotlin.fir.FirPackageDirective
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.symbols.impl.FirFileSymbol
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -25,10 +25,10 @@ abstract class FirFile : FirAnnotatedDeclaration() {
     abstract override val origin: FirDeclarationOrigin
     abstract override val attributes: FirDeclarationAttributes
     abstract override val annotations: List<FirAnnotationCall>
+    abstract val packageDirective: FirPackageDirective
     abstract val imports: List<FirImport>
     abstract val declarations: List<FirDeclaration>
     abstract val name: String
-    abstract val packageFqName: FqName
     abstract override val symbol: FirFileSymbol
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitFile(this, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirFileBuilder.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.declarations.builder
 
 import kotlin.contracts.*
 import org.jetbrains.kotlin.fir.FirModuleData
+import org.jetbrains.kotlin.fir.FirPackageDirective
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
@@ -20,7 +21,6 @@ import org.jetbrains.kotlin.fir.declarations.impl.FirFileImpl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.symbols.impl.FirFileSymbol
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.name.FqName
 
 /*
  * This file was generated automatically
@@ -35,10 +35,10 @@ class FirFileBuilder : FirAnnotationContainerBuilder {
     lateinit var origin: FirDeclarationOrigin
     var attributes: FirDeclarationAttributes = FirDeclarationAttributes()
     override val annotations: MutableList<FirAnnotationCall> = mutableListOf()
+    lateinit var packageDirective: FirPackageDirective
     val imports: MutableList<FirImport> = mutableListOf()
     val declarations: MutableList<FirDeclaration> = mutableListOf()
     lateinit var name: String
-    lateinit var packageFqName: FqName
 
     override fun build(): FirFile {
         return FirFileImpl(
@@ -48,10 +48,10 @@ class FirFileBuilder : FirAnnotationContainerBuilder {
             origin,
             attributes,
             annotations,
+            packageDirective,
             imports,
             declarations,
             name,
-            packageFqName,
         )
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirFileImpl.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.declarations.impl
 
 import org.jetbrains.kotlin.fir.FirModuleData
+import org.jetbrains.kotlin.fir.FirPackageDirective
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationAttributes
@@ -15,7 +16,6 @@ import org.jetbrains.kotlin.fir.declarations.FirImport
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.symbols.impl.FirFileSymbol
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -30,10 +30,10 @@ internal class FirFileImpl(
     override val origin: FirDeclarationOrigin,
     override val attributes: FirDeclarationAttributes,
     override val annotations: MutableList<FirAnnotationCall>,
+    override var packageDirective: FirPackageDirective,
     override val imports: MutableList<FirImport>,
     override val declarations: MutableList<FirDeclaration>,
     override val name: String,
-    override val packageFqName: FqName,
 ) : FirFile() {
     override val symbol: FirFileSymbol = FirFileSymbol()
 
@@ -43,12 +43,14 @@ internal class FirFileImpl(
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         annotations.forEach { it.accept(visitor, data) }
+        packageDirective.accept(visitor, data)
         imports.forEach { it.accept(visitor, data) }
         declarations.forEach { it.accept(visitor, data) }
     }
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirFileImpl {
         transformAnnotations(transformer, data)
+        packageDirective = packageDirective.transform(transformer, data)
         transformImports(transformer, data)
         transformDeclarations(transformer, data)
         return this

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/impl/FirPackageDirectiveImpl.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2021 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.
+ */
+
+package org.jetbrains.kotlin.fir.impl
+
+import org.jetbrains.kotlin.fir.FirPackageDirective
+import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.fir.visitors.*
+
+/*
+ * This file was generated automatically
+ * DO NOT MODIFY IT MANUALLY
+ */
+
+internal class FirPackageDirectiveImpl(
+    override val source: FirSourceElement?,
+    override val packageFqName: FqName,
+) : FirPackageDirective() {
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {}
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirPackageDirectiveImpl {
+        return this
+    }
+}

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirTransformer.kt

@@ -43,6 +43,7 @@ import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.FirPackageDirective
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
@@ -291,6 +292,10 @@ abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {
         return transformElement(file, data)
     }
 
+    open fun transformPackageDirective(packageDirective: FirPackageDirective, data: D): FirPackageDirective {
+        return transformElement(packageDirective, data)
+    }
+
     open fun transformAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: D): FirStatement {
         return transformElement(anonymousFunction, data)
     }
@@ -803,6 +808,10 @@ abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {
         return transformFile(file, data)
     }
 
+    final override fun visitPackageDirective(packageDirective: FirPackageDirective, data: D): FirPackageDirective {
+        return transformPackageDirective(packageDirective, data)
+    }
+
     final override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: D): FirStatement {
         return transformAnonymousFunction(anonymousFunction, data)
     }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirVisitor.kt

@@ -43,6 +43,7 @@ import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.FirPackageDirective
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
@@ -216,6 +217,8 @@ abstract class FirVisitor<out R, in D> {
 
     open fun visitFile(file: FirFile, data: D): R  = visitElement(file, data)
 
+    open fun visitPackageDirective(packageDirective: FirPackageDirective, data: D): R  = visitElement(packageDirective, data)
+
     open fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: D): R  = visitElement(anonymousFunction, data)
 
     open fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression, data: D): R  = visitElement(anonymousFunctionExpression, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirVisitorVoid.kt

@@ -43,6 +43,7 @@ import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.FirPackageDirective
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
@@ -290,6 +291,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitElement(file)
     }
 
+    open fun visitPackageDirective(packageDirective: FirPackageDirective) {
+        visitElement(packageDirective)
+    }
+
     open fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
         visitElement(anonymousFunction)
     }
@@ -802,6 +807,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitFile(file)
     }
 
+    final override fun visitPackageDirective(packageDirective: FirPackageDirective, data: Nothing?) {
+        visitPackageDirective(packageDirective)
+    }
+
     final override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: Nothing?) {
         visitAnonymousFunction(anonymousFunction)
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/Utils.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.types.FirDynamicTypeRef
@@ -14,6 +15,7 @@ import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.builder.*
 import org.jetbrains.kotlin.fir.types.impl.*
+import org.jetbrains.kotlin.name.FqName
 
 fun ModuleInfo.dependenciesWithoutSelf(): Sequence<ModuleInfo> = dependencies().asSequence().filter { it != this }
 
@@ -87,3 +89,6 @@ fun FirSourceElement.getWholeQualifierSourceIfPossible(stepsToWholeQualifier: In
         }
     }
 }
+
+val FirFile.packageFqName: FqName
+    get() = packageDirective.packageFqName

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/FirTreeBuilder.kt

@@ -56,6 +56,7 @@ object FirTreeBuilder : AbstractFirTreeBuilder() {
     val propertyAccessor = element("PropertyAccessor", Declaration, function, contractDescriptionOwner, typeParametersOwner)
     val constructor = element("Constructor", Declaration, function, typeParameterRefsOwner)
     val file = element("File", Declaration, annotatedDeclaration)
+    val packageDirective = element("PackageDirective", Other)
 
     val anonymousFunction = element("AnonymousFunction", Declaration, function, typeParametersOwner)
     val anonymousFunctionExpression = element("AnonymousFunctionExpression", Expression, expression)

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -389,13 +389,17 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
 
         file.configure {
+            +field("packageDirective", packageDirective)
             +fieldList(import).withTransform()
             +declarations.withTransform()
             +stringField("name")
-            +field("packageFqName", fqNameType)
             +symbol("FirFileSymbol")
         }
 
+        packageDirective.configure {
+            +field("packageFqName", fqNameType)
+        }
+
         import.configure {
             +field("importedFqName", fqNameType, nullable = true)
             +booleanField("isAllUnder")

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

@@ -112,6 +112,7 @@ internal object FileDiagnosticRetriever : FileStructureElementDiagnosticRetrieve
             withSuppressedDiagnostics(file) {
                 visitWithDeclaration(file) {
                     file.annotations.forEach { it.accept(this, data) }
+                    file.packageDirective.accept(this, data)
                     file.imports.forEach { it.accept(this, data) }
                     // do not visit declarations here
                 }

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

@@ -6,9 +6,7 @@
 package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.tree.TokenSet
-import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildImport
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
@@ -17,7 +15,6 @@ import org.jetbrains.kotlin.fir.declarations.utils.isCompanion
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.declarations.utils.isStatic
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticPropertySymbol
@@ -46,7 +43,6 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
 import org.jetbrains.kotlin.psi.psiUtil.unwrapNullability
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -184,7 +180,8 @@ internal object FirReferenceResolveHelper {
                 getSymbolsForResolvedQualifier(fir, expression, session, symbolBuilder)
             is FirAnnotationCall -> getSymbolsForAnnotationCall(fir, session, symbolBuilder)
             is FirResolvedImport -> getSymbolsByResolvedImport(expression, symbolBuilder, fir, session)
-            is FirFile -> getSymbolsByFirFile(expression, symbolBuilder, fir)
+            is FirPackageDirective -> getSymbolsForPackageDirective(expression, symbolBuilder)
+            is FirFile -> getSymbolsByFirFile(symbolBuilder, fir)
             is FirArrayOfCall -> {
                 // We can't yet find PsiElement for arrayOf, intArrayOf, etc.
                 emptyList()
@@ -199,6 +196,14 @@ internal object FirReferenceResolveHelper {
         }
     }
 
+    private fun getSymbolsForPackageDirective(
+        expression: KtSimpleNameExpression,
+        symbolBuilder: KtSymbolByFirBuilder
+    ): List<KtFirPackageSymbol> {
+        return listOfNotNull(getPackageSymbolFor(expression, symbolBuilder, forQualifiedType = false))
+    }
+
+
     private fun getSymbolByResolvedNameReference(
         fir: FirResolvedNamedReference,
         session: FirSession,
@@ -319,14 +324,9 @@ internal object FirReferenceResolveHelper {
     }
 
     private fun getSymbolsByFirFile(
-        expression: KtSimpleNameExpression,
         symbolBuilder: KtSymbolByFirBuilder,
         fir: FirFile
     ): List<KtSymbol> {
-        if (expression.getNonStrictParentOfType<KtPackageDirective>() != null) {
-            // Special: package reference in the middle of package directive
-            return listOfNotNull(getPackageSymbolFor(expression, symbolBuilder, forQualifiedType = false))
-        }
         return listOf(symbolBuilder.buildSymbol(fir))
     }
 

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenTopLevelDeclarationsGenerator.kt

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.has
 import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.plugin.fqn
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol