FIR IDE: Add implementation of classifiers completion from indices

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

@@ -10,7 +10,9 @@ import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.patterns.PsiJavaPatterns
 import com.intellij.util.ProcessingContext
+import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.completion.weighers.Weighers
+import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
@@ -39,7 +41,12 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
         if (shouldSuppressCompletion(parameters, result.prefixMatcher)) return
 
         val resultSet = createResultSet(parameters, result)
-        KotlinAvailableScopesCompletionProvider(resultSet.prefixMatcher).addCompletions(parameters, resultSet)
+        val indexHelper = IndexHelper(
+            parameters.position.project,
+            parameters.position.getModuleInfo().contentScope()
+        )
+
+        KotlinCommonCompletionProvider(resultSet.prefixMatcher, indexHelper).addCompletions(parameters, resultSet)
     }
 
     private fun createResultSet(parameters: CompletionParameters, result: CompletionResultSet): CompletionResultSet =
@@ -72,7 +79,15 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
     }
 }
 
-private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatcher) {
+/**
+ * Currently, this class is responsible for collecting all possible completion variants.
+ *
+ * TODO refactor it, try to split into several classes, or decompose it into several classes.
+ */
+private class KotlinCommonCompletionProvider(
+    prefixMatcher: PrefixMatcher,
+    private val indexHelper: IndexHelper
+) {
     private val lookupElementFactory = KotlinFirLookupElementFactory()
 
     private val scopeNameFilter: KtScopeNameFilter =
@@ -138,8 +153,11 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
         implicitScopes: KtScope,
         expectedType: KtType?,
     ) {
-        val availableClasses = implicitScopes.getClassifierSymbols(scopeNameFilter)
+        val classesFromScopes = implicitScopes.getClassifierSymbols(scopeNameFilter)
-        availableClasses.forEach { addSymbolToCompletion(result, expectedType, it) }
+        classesFromScopes.forEach { addSymbolToCompletion(result, expectedType, it) }
+
+        val kotlinClassesFromIndices = indexHelper.getKotlinClasses(scopeNameFilter, psiFilter = { it !is KtEnumEntry })
+        kotlinClassesFromIndices.forEach { addSymbolToCompletion(result, expectedType, it.getSymbol()) }
     }
 
     private fun KtAnalysisSession.collectDotCompletion(

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirLookupElementFactory.kt

@@ -16,13 +16,19 @@ import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.core.withRootPrefixIfNeeded
+import org.jetbrains.kotlin.idea.frontend.api.HackToForceAllowRunningAnalyzeOnEDT
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.analyze
+import org.jetbrains.kotlin.idea.frontend.api.hackyAllowRunningOnEdt
 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.KtFunctionalType
-import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtTypeArgumentList
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.renderer.render
@@ -53,9 +59,19 @@ internal class KotlinFirLookupElementFactory {
  */
 private class UniqueLookupObject
 
+private data class ClassifierLookupObject(val shortName: Name, val classId: ClassId?)
+
 private class ClassLookupElementFactory {
     fun createLookup(symbol: KtClassLikeSymbol): LookupElementBuilder {
-        return LookupElementBuilder.create(UniqueLookupObject(), symbol.name.asString())
+        return LookupElementBuilder.create(ClassifierLookupObject(symbol.name, symbol.classIdIfNonLocal), symbol.name.asString())
+            .withInsertHandler(createInsertHandler(symbol))
+    }
+
+    private fun createInsertHandler(symbol: KtClassLikeSymbol): InsertHandler<LookupElement> {
+        val classFqName = symbol.classIdIfNonLocal?.asSingleFqName()
+            ?: return QuotedNamesAwareInsertionHandler(symbol.name)
+
+        return ClassifierInsertionHandler(classFqName)
     }
 }
 
@@ -115,11 +131,25 @@ private class FunctionLookupElementFactory {
     }
 
     private fun KtAnalysisSession.createInsertHandler(symbol: KtFunctionSymbol): InsertHandler<LookupElement> {
-        return FunctionInsertionHandler(
+        val functionFqName = symbol.callableIdIfNonLocal
-            symbol.name,
+
-            inputValueArguments = symbol.valueParameters.isNotEmpty(),
+        return if (functionFqName != null && canBeCalledByFqName(symbol)) {
-            insertEmptyLambda = insertLambdaBraces(symbol)
+            ShorteningFunctionInsertionHandler(
-        )
+                functionFqName,
+                inputValueArguments = symbol.valueParameters.isNotEmpty(),
+                insertEmptyLambda = insertLambdaBraces(symbol)
+            )
+        } else {
+            SimpleFunctionInsertionHandler(
+                symbol.name,
+                inputValueArguments = symbol.valueParameters.isNotEmpty(),
+                insertEmptyLambda = insertLambdaBraces(symbol)
+            )
+        }
+    }
+
+    private fun canBeCalledByFqName(symbol: KtFunctionSymbol): Boolean {
+        return !symbol.isExtension && symbol.dispatchType == null
     }
 
     companion object {
@@ -128,23 +158,25 @@ private class FunctionLookupElementFactory {
 }
 
 /**
- * A partial copy from `KotlinFunctionInsertHandler.Normal`.
+ * The simplest implementation of the insertion handler for a classifiers.
  */
-private class FunctionInsertionHandler(
+private class ClassifierInsertionHandler(private val fqName: FqName) : InsertHandler<LookupElement> {
+    override fun handleInsert(context: InsertionContext, item: LookupElement) {
+        val targetFile = context.file as? KtFile ?: return
+
+        context.document.replaceString(context.startOffset, context.tailOffset, fqName.render())
+        context.commitDocument()
+
+        shortenReferences(targetFile, TextRange(context.startOffset, context.tailOffset))
+    }
+}
+
+private abstract class AbstractFunctionInsertionHandler(
     name: Name,
     private val inputValueArguments: Boolean,
     private val insertEmptyLambda: Boolean
 ) : QuotedNamesAwareInsertionHandler(name) {
-    override fun handleInsert(context: InsertionContext, item: LookupElement) {
+    protected fun addArguments(context: InsertionContext, offsetElement: PsiElement) {
-        super.handleInsert(context, item)
-
-        val startOffset = context.startOffset
-        val element = context.file.findElementAt(startOffset) ?: return
-
-        addArguments(context, element)
-    }
-
-    private fun addArguments(context: InsertionContext, offsetElement: PsiElement) {
         val completionChar = context.completionChar
         if (completionChar == '(') { //TODO: more correct behavior related to braces type
             context.setAddCompletionChar(false)
@@ -215,6 +247,47 @@ private class FunctionInsertionHandler(
     }
 }
 
+private class SimpleFunctionInsertionHandler(
+    name: Name,
+    inputValueArguments: Boolean,
+    insertEmptyLambda: Boolean
+) : AbstractFunctionInsertionHandler(name, inputValueArguments, insertEmptyLambda) {
+    override fun handleInsert(context: InsertionContext, item: LookupElement) {
+        super.handleInsert(context, item)
+
+        val startOffset = context.startOffset
+        val element = context.file.findElementAt(startOffset) ?: return
+
+        addArguments(context, element)
+    }
+}
+
+private class ShorteningFunctionInsertionHandler(
+    private val name: FqName,
+    inputValueArguments: Boolean,
+    insertEmptyLambda: Boolean,
+) : AbstractFunctionInsertionHandler(name.shortName(), inputValueArguments, insertEmptyLambda) {
+    override fun handleInsert(context: InsertionContext, item: LookupElement) {
+        val targetFile = context.file as? KtFile ?: return
+        super.handleInsert(context, item)
+
+        val startOffset = context.startOffset
+        val element = context.file.findElementAt(startOffset) ?: return
+
+        context.document.replaceString(
+            context.startOffset,
+            context.tailOffset,
+            name.withRootPrefixIfNeeded().render()
+        )
+        context.commitDocument()
+
+        addArguments(context, element)
+        context.commitDocument()
+
+        shortenReferences(targetFile, TextRange(context.startOffset, context.tailOffset))
+    }
+}
+
 private open class QuotedNamesAwareInsertionHandler(private val name: Name) : InsertHandler<LookupElement> {
     override fun handleInsert(context: InsertionContext, item: LookupElement) {
         val startOffset = context.startOffset
@@ -250,3 +323,19 @@ private fun CharSequence.indexOfSkippingSpace(c: Char, startIndex: Int): Int? {
     }
     return null
 }
+
+private fun shortenReferences(targetFile: KtFile, textRange: TextRange) {
+    val shortenings = withAllowedResolve {
+        analyze(targetFile) {
+            collectPossibleReferenceShortenings(targetFile, textRange)
+        }
+    }
+
+    shortenings.invokeShortening()
+}
+
+// FIXME: This is a hack, we should think how we can get rid of it
+private inline fun <T> withAllowedResolve(action: () -> T): T {
+    @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class)
+    return hackyAllowRunningOnEdt(action)
+}

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

@@ -101,7 +101,7 @@ internal class FirModuleResolveStateForCompletion(
 
     @OptIn(InternalForInline::class)
     override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
-        error("Should not be used in completion")
+        return originalState.findSourceFirDeclaration(ktDeclaration)
     }
 
     @OptIn(InternalForInline::class)

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

@@ -5,17 +5,18 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.stubs.StubIndex
 import com.intellij.psi.stubs.StubIndexKey
 import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper.Companion.asStringForIndexes
 import org.jetbrains.kotlin.idea.stubindex.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtProperty
@@ -72,6 +73,18 @@ public class IndexHelper(val project: Project, private val scope: GlobalSearchSc
             .get(packageFqName.asStringForIndexes(), project, scope)
             .mapNotNullTo(hashSetOf()) { it.nameAsName }
 
+    fun getKotlinClasses(
+        nameFilter: (Name) -> Boolean,
+        psiFilter: (element: KtClassOrObject) -> Boolean = { true }
+    ): Collection<KtClassOrObject> {
+        val index = KotlinFullClassNameIndex.getInstance()
+        return index.getAllKeys(project).asSequence()
+            .onEach { ProgressManager.checkCanceled() }
+            .filter { fqName -> nameFilter(getShortName(fqName)) }
+            .flatMap { fqName -> index[fqName, project, scope] }
+            .filter(psiFilter)
+            .toList()
+    }
 
     companion object {
         private fun CallableId.asStringForIndexes(): String =
@@ -82,5 +95,7 @@ public class IndexHelper(val project: Project, private val scope: GlobalSearchSc
 
         private fun ClassId.asStringForIndexes(): String =
             asSingleFqName().asStringForIndexes()
+
+        private fun getShortName(fqName: String) = Name.identifier(fqName.substringAfterLast('.'))
     }
 }