Completion with FIR (prototype, work-in-progress)

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

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtNamedFunction
 
 internal class FirIdeProvider(
     val project: Project,
@@ -80,6 +81,10 @@ internal class FirIdeProvider(
         }
     }
 
+    fun buildFunctionWithBody(ktNamedFunction: KtNamedFunction): FirFunction<*> {
+        return RawFirBuilder(session, cacheProvider.kotlinScopeProvider, stubMode = false).buildFunctionWithBody(ktNamedFunction)
+    }
+
     fun getFile(ktFile: KtFile): FirFile? {
         val (firFile, stamp) = files[ktFile] ?: return null
         if (stamp == ktFile.modificationStamp) return firFile

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

@@ -158,7 +158,7 @@ internal fun KtFile.getOrBuildFirWithDiagnostics(state: FirModuleResolveState):
     return firFile
 }
 
-private fun FirDeclaration.runResolve(
+internal fun FirDeclaration.runResolve(
     file: FirFile,
     firProvider: FirIdeProvider,
     toPhase: FirResolvePhase,
@@ -255,47 +255,59 @@ internal fun KtElement.getOrBuildFir(
         else -> this
     }
     return state.getCachedMapping(this) ?: run {
-        containerFir.accept(object : FirVisitorVoid() {
+        state.recordElementsFrom(containerFir)
-            override fun visitElement(element: FirElement) {
-                (element.realPsi as? KtElement)?.let {
-                    state.record(it, element)
-                }
-                element.acceptChildren(this)
-            }
 
-            override fun visitReference(reference: FirReference) {}
+        val (current, mappedFir) = psi.getFirOfClosestParent(state) ?: error("FirElement is not found for: $text")
-
+        if (current != this) {
-            override fun visitControlFlowGraphReference(controlFlowGraphReference: FirControlFlowGraphReference) {}
+            state.record(current, mappedFir)
-
-            override fun visitNamedReference(namedReference: FirNamedReference) {}
-
-            override fun visitResolvedNamedReference(resolvedNamedReference: FirResolvedNamedReference) {}
-
-            override fun visitDelegateFieldReference(delegateFieldReference: FirDelegateFieldReference) {}
-
-            override fun visitBackingFieldReference(backingFieldReference: FirBackingFieldReference) {}
-
-            override fun visitSuperReference(superReference: FirSuperReference) {}
-
-            override fun visitThisReference(thisReference: FirThisReference) {}
-
-            override fun visitErrorTypeRef(errorTypeRef: FirErrorTypeRef) {}
-
-            override fun visitUserTypeRef(userTypeRef: FirUserTypeRef) {
-                userTypeRef.acceptChildren(this)
-            }
-        })
-        var current: PsiElement? = psi
-        while (current is KtElement) {
-            val mappedFir = state.getCachedMapping(current)
-            if (mappedFir != null) {
-                if (current != this) {
-                    state.record(current, mappedFir)
-                }
-                return mappedFir
-            }
-            current = current.parent
         }
-        error("FirElement is not found for: $text")
+
+        mappedFir
     }
 }
+
+internal fun KtElement.getFirOfClosestParent(state: FirModuleResolveState): Pair<KtElement, FirElement>? {
+    var current: PsiElement? = this
+    while (current is KtElement) {
+        val mappedFir = state.getCachedMapping(current)
+        if (mappedFir != null) {
+            return current to mappedFir
+        }
+        current = current.parent
+    }
+
+    return null
+}
+
+internal fun FirModuleResolveState.recordElementsFrom(containerFir: FirDeclaration) {
+    containerFir.accept(object : FirVisitorVoid() {
+        override fun visitElement(element: FirElement) {
+            (element.realPsi as? KtElement)?.let {
+                record(it, element)
+            }
+            element.acceptChildren(this)
+        }
+
+        override fun visitReference(reference: FirReference) {}
+
+        override fun visitControlFlowGraphReference(controlFlowGraphReference: FirControlFlowGraphReference) {}
+
+        override fun visitNamedReference(namedReference: FirNamedReference) {}
+
+        override fun visitResolvedNamedReference(resolvedNamedReference: FirResolvedNamedReference) {}
+
+        override fun visitDelegateFieldReference(delegateFieldReference: FirDelegateFieldReference) {}
+
+        override fun visitBackingFieldReference(backingFieldReference: FirBackingFieldReference) {}
+
+        override fun visitSuperReference(superReference: FirSuperReference) {}
+
+        override fun visitThisReference(thisReference: FirThisReference) {}
+
+        override fun visitErrorTypeRef(errorTypeRef: FirErrorTypeRef) {}
+
+        override fun visitUserTypeRef(userTypeRef: FirUserTypeRef) {
+            userTypeRef.acceptChildren(this)
+        }
+    })
+}

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

@@ -7,9 +7,12 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtNamedFunction
 
 object LowLevelFirApiFacade {
     fun getResolveStateFor(element: KtElement): FirModuleResolveState =
@@ -25,6 +28,26 @@ object LowLevelFirApiFacade {
     fun getOrBuildFirFor(element: KtElement, resolveState: FirModuleResolveState, phase: FirResolvePhase): FirElement =
         element.getOrBuildFir(resolveState, phase)
 
+
+    fun getFirOfClosestParent(element: KtElement): FirElement? = element.getFirOfClosestParent(element.firResolveState())?.second
+
+    fun buildFunctionWithResolvedBody(
+        firFile: FirFile,
+        element: KtNamedFunction,
+        phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
+    ): FirFunction<*> {
+        val state = element.firResolveState()
+        val firIdeProvider = firFile.session.firIdeProvider
+        val builtFunction = firIdeProvider.buildFunctionWithBody(element)
+
+        return builtFunction.apply {
+            runResolve(firFile, firIdeProvider, phase, state)
+
+            // TODO this PSI caching should be somewhere else
+            state.recordElementsFrom(this)
+        }
+    }
+
     fun getDiagnosticsFor(element: KtElement, resolveState: FirModuleResolveState): Collection<Diagnostic> {
         val file = element.containingKtFile
         file.getOrBuildFirWithDiagnostics(resolveState)

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

@@ -0,0 +1,120 @@
+/*
+ * 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.completion
+
+import com.intellij.codeInsight.completion.*
+import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.patterns.PlatformPatterns
+import com.intellij.util.ProcessingContext
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByTypeRef
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
+import org.jetbrains.kotlin.idea.fir.getFirOfClosestParent
+import org.jetbrains.kotlin.idea.fir.getOrBuildFir
+import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
+import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+
+class KotlinFirCompletionContributor : CompletionContributor() {
+    init {
+        extend(CompletionType.BASIC, PlatformPatterns.psiElement(), KotlinFirCompletionProvider)
+    }
+}
+
+/**
+ * 1. Collect all receivers for the expression (implicit and explicit)
+ */
+private object KotlinFirCompletionProvider : CompletionProvider<CompletionParameters>() {
+    override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
+        val originalFile = parameters.originalFile as? KtFile ?: return
+
+        val originalFileFir = originalFile.getOrBuildFirSafe<FirFile>(LowLevelFirApiFacade.getResolveStateFor(originalFile)) ?: return
+
+        val reference = (parameters.position.parent as? KtSimpleNameExpression)?.mainReference ?: return
+        val nameExpression = reference.expression.takeIf { it !is KtLabelReferenceExpression } ?: return
+
+        val parentFunction = nameExpression.getNonStrictParentOfType<KtNamedFunction>() ?: return
+
+        val firFunction = LowLevelFirApiFacade.buildFunctionWithResolvedBody(originalFileFir, parentFunction)
+
+        val element = nameExpression.getFirOfClosestParent() as? FirQualifiedAccessExpression ?: return
+
+        for (scope in getScopes(element, firFunction.session)) {
+            for (name in scope.availableNames()) {
+                result.addElement(LookupElementBuilder.create(name.asString()))
+            }
+        }
+
+//        val receiver = element.explicitReceiver
+//        if (receiver != null) {
+//
+//            val firScope =
+//                receiver.typeRef.coneTypeUnsafe<ConeKotlinType>().scope(firFunction.session, ScopeSession()) ?: return
+//                // (symbolProvider.getSymbolByTypeRef<FirClassSymbol<*>>(receiver.typeRef))?.buildUseSiteMemberScope(firFunction.session, ScopeSession()) ?: return
+//
+//            for (name in firScope.getCallableNames()) {
+//                firScope.processFunctionsByName(name) {
+//                    it.fir
+//                    result.addElement(LookupElementBuilder.create(name.asString()))
+//                }
+//
+//                firScope.processPropertiesByName(name) {
+//                    result.addElement(LookupElementBuilder.create(name.asString()))
+//                }
+//            }
+//        }
+
+//        val symbolProvider = originalFileFir.session.firSymbolProvider
+//
+//        val allCallableNamesInPackage =
+////            symbolProvider.getTopLevelCallableSymbols(originalFile.packageFqName, Name.identifier("bar"))
+//            symbolProvider.getAllCallableNamesInPackage(originalFileFir.packageFqName)
+//
+//        for (it in allCallableNamesInPackage) {
+//            result.addElement(LookupElementBuilder.create(it.asString()))
+//        }
+//
+//        val scope = ((nameExpression.getOrBuildFir() as? FirQualifiedAccessExpression)?.explicitReceiver as? FirFunctionCall)?.typeRef?.let {
+//            (symbolProvider.getSymbolByTypeRef<AbstractFirBasedSymbol<*>>(it) as? FirClassSymbol<*>)?.classId?.let {
+//                symbolProvider.getAllCallableNamesInClass(it)
+//            }
+//        }
+//
+////        val type = (((nameExpression.getOrBuildFir() as FirQualifiedAccessExpression).explicitReceiver as FirFunctionCall).typeRef as FirResolvedTypeRef).type as ConeClassLikeType
+////        originalFileFir.session.declaredMemberScopeProvider.declaredMemberScope(type.lookupTag.toSymbol(originalFileFir.session).fir as FirRegularClass)
+//
+        return
+    }
+
+    interface ScopeWrapper {
+        fun availableNames(): Set<Name>
+    }
+
+    fun getScopes(element: FirQualifiedAccessExpression, session: FirSession): Sequence<ScopeWrapper> = sequence {
+        val receiver = element.explicitReceiver
+        if (receiver != null) {
+
+            val firScope = receiver.typeRef.coneTypeUnsafe<ConeKotlinType>().scope(session, ScopeSession())
+            // (symbolProvider.getSymbolByTypeRef<FirClassSymbol<*>>(receiver.typeRef))?.buildUseSiteMemberScope(firFunction.session, ScopeSession()) ?: return
+
+            if (firScope != null) {
+                yield(object : ScopeWrapper {
+                    override fun availableNames(): Set<Name> {
+                        return firScope.getCallableNames()
+                    }
+                })
+            }
+        }
+    }
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt

@@ -24,6 +24,8 @@ fun KtElement.getOrBuildFir(
     phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
 ) = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState, phase)
 
+fun KtElement.getFirOfClosestParent() = LowLevelFirApiFacade.getFirOfClosestParent(this)
+
 inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
     resolveState: FirModuleResolveState,
     phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE

idea/resources-fir/META-INF/plugin.xml

@@ -91,6 +91,11 @@ The Kotlin FIR plugin provides language support in IntelliJ IDEA and Android Stu
     <projectService serviceInterface="org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService"
                     serviceImplementation="org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateServiceImpl"/>
     <projectService serviceImplementation="org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade"/>
+
+    <completion.contributor language="kotlin"
+                            id="KotlinCompletionContributor"
+                            order="first"
+                            implementationClass="org.jetbrains.kotlin.idea.completion.KotlinFirCompletionContributor"/>
   </extensions>
 
   <!--  scripts  -->