FIR IDE: handle importing scopes in completion in HL API

Co-authored-by: Roman Golyshev <roman.golyshev@jetbrains.com>

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirAbstractImportingScope.kt

@@ -1,6 +1,6 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * that can be found in the license/LICENSE.txt file.
+ * 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.scopes.impl
@@ -26,7 +26,7 @@ abstract class FirAbstractImportingScope(
 ) : FirAbstractProviderBasedScope(session, lookupInFir) {
 
     // TODO: Rewrite somehow?
-    private fun getStaticsScope(classId: ClassId): FirScope? {
+    fun getStaticsScope(classId: ClassId): FirScope? {
         val symbol = provider.getClassLikeSymbolByFqName(classId) ?: return null
         if (symbol is FirTypeAliasSymbol) {
             val expansionSymbol = symbol.fir.expandedConeType?.lookupTag?.toSymbol(session)

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/scopes/KtImportingScope.kt

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2020 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.frontend.api.scopes
+
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+interface KtImportingScope : KtScope {
+    val imports: List<Import>
+    val isDefaultImportingScope: Boolean
+}
+
+interface KtStarImportingScope : KtImportingScope {
+    override val imports: List<StarImport>
+}
+
+interface KtNonStarImportingScope : KtImportingScope {
+    override val imports: List<NonStarImport>
+}
+
+sealed class Import {
+    abstract val packageFqName: FqName
+    abstract val relativeClassName: FqName?
+    abstract val resolvedClassId: ClassId?
+}
+
+class NonStarImport(
+    override val packageFqName: FqName,
+    override val relativeClassName: FqName?,
+    override val resolvedClassId: ClassId?,
+    val callableName: Name?,
+) : Import()
+
+class StarImport(
+    override val packageFqName: FqName,
+    override val relativeClassName: FqName?,
+    override val resolvedClassId: ClassId?,
+) : Import()

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/scopes/KtScope.kt

@@ -10,6 +10,8 @@ import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.name.Name
 
 interface KtScope : ValidityOwner {
+    // TODO check that names are accessible
+    // maybe return some kind of lazy set
     fun getAllNames(): Set<Name>
     fun getCallableNames(): Set<Name>
     fun getClassLikeSymbolNames(): Set<Name>

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt

@@ -200,7 +200,8 @@ private fun FirScope.collectCallableSymbols(packageIndexHelper: PackageIndexHelp
     return allCallableSymbols
 }
 
-private class PackageIndexHelper(private val project: Project) {
+class PackageIndexHelper(private val project: Project) {
+    //todo use more concrete scope
     private val searchScope = GlobalSearchScope.allScope(project)
 
     private val functionByPackageIndex = KotlinTopLevelFunctionByPackageIndex.getInstance()
@@ -210,9 +211,10 @@ private class PackageIndexHelper(private val project: Project) {
         return getTopLevelCallables(packageFqName).mapTo(mutableSetOf()) { it.nameAsSafeName }
     }
 
-    private fun getTopLevelCallables(packageFqName: FqName): Sequence<KtCallableDeclaration> = sequence {
+    @OptIn(ExperimentalStdlibApi::class)
-        yieldAll(functionByPackageIndex.get(packageFqName.asString(), project, searchScope))
+    private fun getTopLevelCallables(packageFqName: FqName): List<KtCallableDeclaration> = buildList {
-        yieldAll(propertyByPackageIndex.get(packageFqName.asString(), project, searchScope))
+        addAll(functionByPackageIndex.get(packageFqName.asString(), project, searchScope))
+        addAll(propertyByPackageIndex.get(packageFqName.asString(), project, searchScope))
     }
 }
 

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

@@ -44,6 +44,7 @@ constructor(
     resolveState: FirModuleResolveState
 ) : KtAnalysisSession(element.project) {
     internal val token get() = validityToken
+    private val project = element.project
 
     internal val firResolveState = resolveState
 
@@ -58,7 +59,7 @@ constructor(
     internal val firSymbolBuilder = KtSymbolByFirBuilder(
         firSession.firSymbolProvider,
         typeCheckerContext,
-        element.project,
+        project,
         validityToken
     )
 
@@ -70,7 +71,7 @@ constructor(
             firSymbolBuilder
         )
 
-    override val scopeProvider: KtScopeProvider = KtFirScopeProvider(token, firSymbolBuilder, firSession, firResolveState)
+    override val scopeProvider: KtScopeProvider = KtFirScopeProvider(token, firSymbolBuilder, project, firSession, firResolveState)
 
     init {
         assertIsValid()

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirDelegatingScope.kt

@@ -29,7 +29,7 @@ internal abstract class KtFirDelegatingScope(private val builder: KtSymbolByFirB
 
     override fun getAllNames(): Set<Name> = withValidityAssertion {
         if (allNamesCached == null) {
-            allNamesCached = firScope.getCallableNames() + firScope.getClassifierNames()
+            allNamesCached = getCallableNames() + getClassLikeSymbolNames()
         }
         allNamesCached!!
     }
@@ -51,7 +51,7 @@ internal abstract class KtFirDelegatingScope(private val builder: KtSymbolByFirB
 
     override fun getCallableSymbols(): Sequence<KtCallableSymbol> = withValidityAssertion {
         sequence {
-            firScope.getCallableNames().forEach { name ->
+            getCallableNames().forEach { name ->
                 val callables = mutableListOf<KtCallableSymbol>()
                 firScope.processFunctionsByName(name) { firSymbol ->
                     (firSymbol.fir as? FirSimpleFunction)?.let { fir ->
@@ -74,7 +74,7 @@ internal abstract class KtFirDelegatingScope(private val builder: KtSymbolByFirB
 
     override fun getClassClassLikeSymbols(): Sequence<KtClassLikeSymbol> = withValidityAssertion {
         sequence {
-            firScope.getClassifierNames().forEach { name ->
+            getClassLikeSymbolNames().forEach { name ->
                 val classLikeSymbols = mutableListOf<KtClassLikeSymbol>()
                 firScope.processClassifiersByName(name) { firSymbol ->
                     (firSymbol.fir as? FirClassLikeDeclaration<*>)?.let {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirNonStarImportingScope.kt

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2010-2020 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.frontend.api.fir.scopes
+
+import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractSimpleImportingScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirDefaultSimpleImportingScope
+import org.jetbrains.kotlin.idea.frontend.api.ValidityOwner
+import org.jetbrains.kotlin.idea.frontend.api.ValidityOwnerByValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtNonStarImportingScope
+import org.jetbrains.kotlin.idea.frontend.api.scopes.NonStarImport
+import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.name.Name
+
+internal class KtFirNonStarImportingScope(
+    firScope: FirAbstractSimpleImportingScope,
+    builder: KtSymbolByFirBuilder,
+    override val token: ValidityOwner
+) : KtFirDelegatingScope(builder), KtNonStarImportingScope, ValidityOwnerByValidityToken {
+    override val firScope: FirAbstractSimpleImportingScope = firScope
+
+    @OptIn(ExperimentalStdlibApi::class)
+    override val imports: List<NonStarImport> by cached {
+        buildList {
+            firScope.simpleImports.values.forEach { imports ->
+                imports.forEach { import ->
+                    NonStarImport(
+                        import.packageFqName,
+                        import.relativeClassName,
+                        import.resolvedClassId,
+                        import.importedName
+                    ).let(::add)
+                }
+            }
+        }
+    }
+
+
+    override fun getCallableNames(): Set<Name> = withValidityAssertion {
+        imports.mapNotNullTo(hashSetOf()) { it.callableName }
+    }
+
+    override fun getClassLikeSymbolNames(): Set<Name> = withValidityAssertion {
+        imports.mapNotNullTo((hashSetOf())) { it.relativeClassName?.shortName() }
+    }
+
+
+    override val isDefaultImportingScope: Boolean = withValidityAssertion { firScope is FirDefaultSimpleImportingScope }
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirScopeProvider.kt

@@ -5,11 +5,14 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.scopes
 
+import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.scope
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractSimpleImportingScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractStarImportingScope
 import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
@@ -35,6 +38,7 @@ import java.util.*
 internal class KtFirScopeProvider(
     override val token: ValidityOwner,
     private val builder: KtSymbolByFirBuilder,
+    private val project: Project,
     session: FirSession,
     val firResolveState: FirModuleResolveState
 ) : KtScopeProvider(), ValidityOwnerByValidityToken {
@@ -96,15 +100,27 @@ internal class KtFirScopeProvider(
         val nonLocalScopes = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.scope }
         val firLocalScopes = towerDataContext.localScopes
 
-        val allKtScopes = (implicitReceiverScopes + nonLocalScopes + firLocalScopes).map { convertToKtScope(it) }
+        @OptIn(ExperimentalStdlibApi::class)
+        val allKtScopes = buildList {
+            implicitReceiverScopes.mapTo(this, ::convertToKtScope)
+            nonLocalScopes.mapTo(this, ::convertToKtScope)
+            firLocalScopes.mapTo(this, ::convertToKtScope)
+        }
 
         KtScopeContext(getCompositeScope(allKtScopes), implicitReceiversTypes)
     }
 
-    private fun convertToKtScope(firScope: FirScope): KtScope = KtFirDelegatingScopeImpl(firScope, builder, token)
+    private fun convertToKtScope(firScope: FirScope): KtScope = when (firScope) {
+        is FirAbstractSimpleImportingScope -> KtFirNonStarImportingScope(firScope, builder, token)
+        is FirAbstractStarImportingScope -> KtFirStarImportingScope(firScope, builder, project, token)
+        else -> {
+            // todo create concrete KtScope here instead of a generic one
+            KtFirDelegatingScopeImpl(firScope, builder, token)
+        }
+    }
 }
 
 private class KtFirDelegatingScopeImpl(firScope: FirScope, builder: KtSymbolByFirBuilder, override val token: ValidityOwner) :
     KtFirDelegatingScope(builder), ValidityOwnerByValidityToken {
-    override val firScope: FirScope by weakRef(firScope)
+    override val firScope: FirScope = firScope
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirStarImportingScope.kt

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2010-2020 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.frontend.api.fir.scopes
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractStarImportingScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirDefaultStarImportingScope
+import org.jetbrains.kotlin.idea.completion.PackageIndexHelper
+import org.jetbrains.kotlin.idea.frontend.api.ValidityOwner
+import org.jetbrains.kotlin.idea.frontend.api.ValidityOwnerByValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.scopes.Import
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtStarImportingScope
+import org.jetbrains.kotlin.idea.frontend.api.scopes.StarImport
+import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.name.Name
+
+internal class KtFirStarImportingScope(
+    firScope: FirAbstractStarImportingScope,
+    builder: KtSymbolByFirBuilder,
+    project: Project,
+    override val token: ValidityOwner,
+) : KtFirDelegatingScope(builder), KtStarImportingScope, ValidityOwnerByValidityToken {
+    override val firScope: FirAbstractStarImportingScope = firScope
+    override val isDefaultImportingScope: Boolean = withValidityAssertion { firScope is FirDefaultStarImportingScope }
+    private val packageHelper = PackageIndexHelper(project)
+
+    override val imports: List<StarImport> by cached {
+        firScope.starImports.map { import ->
+            StarImport(
+                import.packageFqName,
+                import.relativeClassName,
+                import.resolvedClassId
+            )
+        }
+    }
+
+    // todo cache?
+    @OptIn(ExperimentalStdlibApi::class)
+    override fun getCallableNames(): Set<Name> = withValidityAssertion {
+        imports.flatMapTo(hashSetOf()) { import: Import ->
+            if (import.relativeClassName == null) { // top level callable
+                packageHelper.getPackageTopLevelNames(import.packageFqName)
+            } else { //member
+                val classId = import.resolvedClassId ?: error("Class id should not be null as relativeClassName is not null")
+                firScope.getStaticsScope(classId)?.getCallableNames().orEmpty()
+            }
+        }
+    }
+
+    override fun getClassLikeSymbolNames(): Set<Name> = withValidityAssertion {
+        //TODO
+        setOf()
+    }
+
+}