FIR IDE: fix package reference resolve

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/FirCompletionContributorBase.kt

@@ -16,9 +16,11 @@ import org.jetbrains.kotlin.idea.completion.context.FirRawPositionCompletionCont
 import org.jetbrains.kotlin.idea.completion.weighers.Weighers
 import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeNameFilter
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
+import org.jetbrains.kotlin.idea.stubindex.KotlinSourceFilterScope
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.psi.KtFile
 
@@ -32,6 +34,10 @@ internal abstract class FirCompletionContributorBase<C : FirRawPositionCompletio
     protected val targetPlatform: TargetPlatform get() = basicContext.targetPlatform
     protected val indexHelper: IndexHelper get() = basicContext.indexHelper
     protected val lookupElementFactory: KotlinFirLookupElementFactory get() = basicContext.lookupElementFactory
+    protected val visibleScope = basicContext.visibleScope
+
+    protected val scopeNameFilter: KtScopeNameFilter =
+        { name -> !name.isSpecial && prefixMatcher.prefixMatches(name.identifier) }
 
     abstract fun KtAnalysisSession.complete(positionContext: C)
 

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/FirContextCompletionContributorBase.kt

@@ -14,8 +14,4 @@ import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeNameFilter
 
 internal abstract class FirContextCompletionContributorBase<C : FirNameReferenceRawPositionContext>(
     basicContext: FirBasicCompletionContext
-) : FirCompletionContributorBase<C>(basicContext) {
+) : FirCompletionContributorBase<C>(basicContext)
-
-    protected val scopeNameFilter: KtScopeNameFilter =
-        { name -> !name.isSpecial && prefixMatcher.prefixMatches(name.identifier) }
-}

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/FirPackageCompletionContributor.kt

@@ -0,0 +1,42 @@
+/*
+ * 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.idea.completion.contributors
+
+import com.intellij.psi.JavaPsiFacade
+import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
+import org.jetbrains.kotlin.idea.completion.context.FirNameReferenceRawPositionContext
+import org.jetbrains.kotlin.idea.completion.context.FirRawPositionCompletionContext
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
+import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.idea.stubindex.PackageIndexUtil
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.platform.jvm.isJvm
+import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
+import org.jetbrains.kotlin.psi.KtExpression
+
+internal class FirPackageCompletionContributor(
+    basicContext: FirBasicCompletionContext,
+) : FirCompletionContributorBase<FirRawPositionCompletionContext>(basicContext) {
+
+    override fun KtAnalysisSession.complete(positionContext: FirRawPositionCompletionContext) {
+        val rootSymbol = if (positionContext !is FirNameReferenceRawPositionContext || positionContext.explicitReceiver == null) {
+            ROOT_PACKAGE_SYMBOL
+        } else {
+            positionContext.explicitReceiver?.reference()?.resolveToSymbol() as? KtPackageSymbol
+        } ?: return
+        rootSymbol.getPackageScope()
+            .getPackageSymbols(scopeNameFilter)
+            .forEach { packageSymbol ->
+                result.addElement(lookupElementFactory.createPackageLookupElement(packageSymbol.fqName))
+            }
+    }
+
+    private fun KtExpression.reference() = when (this) {
+        is KtDotQualifiedExpression -> selectorExpression?.mainReference
+        else -> mainReference
+    }
+}

idea/idea-fir/tests/org/jetbrains/kotlin/idea/fir/resolve/FirReferenceResolveTestGenerated.java

@@ -866,4 +866,37 @@ public class FirReferenceResolveTestGenerated extends AbstractFirReferenceResolv
             runTest("idea/testData/resolve/references/nestedTypes/ResolveTypeInTheStartOfType.kt");
         }
     }
+
+    @TestMetadata("idea/testData/resolve/references/packageReference")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class PackageReference extends AbstractFirReferenceResolveTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInPackageReference() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/resolve/references/packageReference"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("javaPackageFirstQualifier.kt")
+        public void testJavaPackageFirstQualifier() throws Exception {
+            runTest("idea/testData/resolve/references/packageReference/javaPackageFirstQualifier.kt");
+        }
+
+        @TestMetadata("javaPackageSecondQualifier.kt")
+        public void testJavaPackageSecondQualifier() throws Exception {
+            runTest("idea/testData/resolve/references/packageReference/javaPackageSecondQualifier.kt");
+        }
+
+        @TestMetadata("kotlinPackageFirstQualifier.kt")
+        public void testKotlinPackageFirstQualifier() throws Exception {
+            runTest("idea/testData/resolve/references/packageReference/kotlinPackageFirstQualifier.kt");
+        }
+
+        @TestMetadata("kotlinPackageSecondQualifier.kt")
+        public void testKotlinPackageSecondQualifier() throws Exception {
+            runTest("idea/testData/resolve/references/packageReference/kotlinPackageSecondQualifier.kt");
+        }
+    }
 }

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

@@ -26,7 +26,8 @@ class KtFirPackageSymbol(
     override val token: ValidityToken
 ) : KtPackageSymbol(), ValidityTokenOwner {
     override val psi: PsiElement? by cached {
-        KtPackage(PsiManager.getInstance(project), fqName, GlobalSearchScope.allScope(project)/*TODO*/)
+        JavaPsiFacade.getInstance(project).findPackage(fqName.asString())
+            ?: KtPackage(PsiManager.getInstance(project), fqName, GlobalSearchScope.allScope(project)/*TODO*/)
     }
 
     override val origin: KtSymbolOrigin

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

@@ -366,7 +366,7 @@ internal object FirReferenceResolveHelper {
         analysisSession: KtFirAnalysisSession
     ): Collection<KtSymbol> {
         // TODO refactor that block
-        val classId = fir.classId ?: return emptyList()
+        val classId = fir.classId ?: return listOfNotNull(getPackageSymbolFor(expression, symbolBuilder, forQualifiedType = false))
 
         var parent = expression.parent as? KtDotQualifiedExpression
         // Distinguish A.foo() from A(.Companion).foo()

idea/testData/resolve/references/packageReference/javaPackageFirstQualifier.kt

@@ -0,0 +1,4 @@
+fun x() {
+    val z = jav<caret>a.util.ArrayList()
+}
+// REF: java

idea/testData/resolve/references/packageReference/javaPackageSecondQualifier.kt

@@ -0,0 +1,4 @@
+fun x() {
+    val z = java.uti<caret>l.ArrayList()
+}
+// REF: util

idea/testData/resolve/references/packageReference/kotlinPackageFirstQualifier.kt

@@ -0,0 +1,4 @@
+fun x() {
+    val z = kotli<caret>n.text.toString()
+}
+// REF: kotlin

idea/testData/resolve/references/packageReference/kotlinPackageSecondQualifier.kt

@@ -0,0 +1,4 @@
+fun x() {
+    val z = kotlin.tex<caret>t.toString()
+}
+// REF: text

idea/tests/org/jetbrains/kotlin/idea/resolve/ReferenceResolveTestGenerated.java

@@ -866,4 +866,37 @@ public class ReferenceResolveTestGenerated extends AbstractReferenceResolveTest
             runTest("idea/testData/resolve/references/nestedTypes/ResolveTypeInTheStartOfType.kt");
         }
     }
+
+    @TestMetadata("idea/testData/resolve/references/packageReference")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class PackageReference extends AbstractReferenceResolveTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInPackageReference() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/resolve/references/packageReference"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("javaPackageFirstQualifier.kt")
+        public void testJavaPackageFirstQualifier() throws Exception {
+            runTest("idea/testData/resolve/references/packageReference/javaPackageFirstQualifier.kt");
+        }
+
+        @TestMetadata("javaPackageSecondQualifier.kt")
+        public void testJavaPackageSecondQualifier() throws Exception {
+            runTest("idea/testData/resolve/references/packageReference/javaPackageSecondQualifier.kt");
+        }
+
+        @TestMetadata("kotlinPackageFirstQualifier.kt")
+        public void testKotlinPackageFirstQualifier() throws Exception {
+            runTest("idea/testData/resolve/references/packageReference/kotlinPackageFirstQualifier.kt");
+        }
+
+        @TestMetadata("kotlinPackageSecondQualifier.kt")
+        public void testKotlinPackageSecondQualifier() throws Exception {
+            runTest("idea/testData/resolve/references/packageReference/kotlinPackageSecondQualifier.kt");
+        }
+    }
 }