Expose local scopes to use in completion

2020-06-19 15:16:16 +03:00
parent 6129f4bcef
commit 97d4918ed3
7 changed files with 79 additions and 17 deletions
@@ -34,6 +34,8 @@ class PersistentMultimap<K, V> private constructor(private val map: PersistentMa
    operator fun get(key: K): List<V> {
        return map[key] ?: emptyList()
    }
+
+    val keys: ImmutableSet<K> get() = map.keys
 }

 class PersistentSetMultimap<K, V> private constructor(private val map: PersistentMap<K, PersistentSet<V>>) {
@@ -63,4 +65,4 @@ class PersistentSetMultimap<K, V> private constructor(private val map: Persisten
    operator fun get(key: K): Set<V> {
        return map[key] ?: emptySet()
    }
-}
+}
@@ -83,6 +83,8 @@ open class FirBodyResolveTransformer(
        return data.expectedTypeRef.compose()
    }

+    open fun onBeforeStatementResolution(statement: FirStatement) {}
+
    // ------------------------------------- Expressions -------------------------------------

    override fun transformExpression(expression: FirExpression, data: ResolutionMode): CompositeTransformResult<FirStatement> {
@@ -291,6 +291,9 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform

        block.transformStatementsIndexed(transformer) { index ->
            val value = if (index == numberOfStatements - 1) data else ResolutionMode.ContextIndependent
+
+            transformer.onBeforeStatementResolution(block.statements[index])
+
            TransformData.Data(value)
        }
        if (data == ResolutionMode.ContextIndependent) {
@@ -70,4 +70,6 @@ class FirLocalScope private constructor(
    }

    override fun mayContainName(name: Name) = properties.containsKey(name) || functions[name].isNotEmpty() || classes.containsKey(name)
+
+    override fun getCallableNames(): Set<Name> = properties.keys + functions.keys + classes.keys
 }
@@ -9,8 +9,12 @@ import com.intellij.openapi.progress.ProgressIndicatorProvider
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.*
-import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredCallableSymbols
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
@@ -162,7 +166,8 @@ internal fun FirDeclaration.runResolve(
    file: FirFile,
    firProvider: FirIdeProvider,
    toPhase: FirResolvePhase,
-    state: FirModuleResolveState
+    state: FirModuleResolveState,
+    towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
 ) {
    val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
    file.runResolve(toPhase = nonLazyPhase, fromPhase = this.resolvePhase)
@@ -193,7 +198,8 @@ internal fun FirDeclaration.runResolve(
    val transformer = FirDesignatedBodyResolveTransformerForIDE(
        designation.iterator(), state.getSession(psi as KtElement),
        scopeSession,
-        implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE
+        implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+        towerDataContextForStatement
    )
    file.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
 }
@@ -202,7 +208,8 @@ private class FirDesignatedBodyResolveTransformerForIDE(
    private val designation: Iterator<FirElement>,
    session: FirSession,
    scopeSession: ScopeSession,
-    implicitTypeOnly: Boolean
+    implicitTypeOnly: Boolean,
+    private val towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null
 ) : FirBodyResolveTransformer(
    session,
    phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
@@ -210,6 +217,7 @@ private class FirDesignatedBodyResolveTransformerForIDE(
    scopeSession = scopeSession,
    returnTypeCalculator = createReturnTypeCalculatorForIDE(session, scopeSession)
 ) {
+
    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
        if (designation.hasNext()) {
            designation.next().visitNoTransform(this, data)
@@ -218,6 +226,11 @@ private class FirDesignatedBodyResolveTransformerForIDE(

        return super.transformDeclarationContent(declaration, data)
    }
+
+    override fun onBeforeStatementResolution(statement: FirStatement) {
+        if (towerDataContextForStatement == null) return
+        towerDataContextForStatement[statement] = context.towerDataContext
+    }
 }


@@ -5,12 +5,15 @@

 package org.jetbrains.kotlin.idea.fir.low.level.api

+import com.intellij.psi.PsiElement
 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.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtNamedFunction

@@ -31,21 +34,48 @@ object LowLevelFirApiFacade {

    fun getFirOfClosestParent(element: KtElement): FirElement? = element.getFirOfClosestParent(element.firResolveState())?.second

-    fun buildFunctionWithResolvedBody(
+    class FirCompletionContext internal constructor(
+        val session: FirSession,
+        private val towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>,
+        private val state: FirModuleResolveState,
+    ) {
+        fun getTowerDataContext(element: KtElement): FirTowerDataContext {
+            var current: PsiElement? = element
+            while (current is KtElement) {
+                val mappedFir = state.getCachedMapping(current)
+
+                if (mappedFir is FirStatement) {
+                    towerDataContextForStatement[mappedFir]?.let { return it }
+                }
+                current = current.parent
+            }
+
+            error("No context for $element")
+        }
+    }
+
+    fun buildCompletionContextForFunction(
        firFile: FirFile,
        element: KtNamedFunction,
        phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
-    ): FirFunction<*> {
+    ): FirCompletionContext {
        val state = element.firResolveState()
        val firIdeProvider = firFile.session.firIdeProvider
        val builtFunction = firIdeProvider.buildFunctionWithBody(element)
+        val towerDataContextForStatement = mutableMapOf<FirStatement, FirTowerDataContext>()

-        return builtFunction.apply {
-            runResolve(firFile, firIdeProvider, phase, state)
+        val function = builtFunction.apply {
+            runResolve(firFile, firIdeProvider, phase, state, towerDataContextForStatement)

            // TODO this PSI caching should be somewhere else
            state.recordElementsFrom(this)
        }
+
+        return FirCompletionContext(
+            function.session,
+            towerDataContextForStatement,
+            state
+        )
    }

    fun getDiagnosticsFor(element: KtElement, resolveState: FirModuleResolveState): Collection<Diagnostic> {
@@ -53,4 +83,4 @@ object LowLevelFirApiFacade {
        file.getOrBuildFirWithDiagnostics(resolveState)
        return resolveState.getDiagnostics(element)
    }
-}
+}
@@ -12,18 +12,19 @@ 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.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.scope
 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.KtFile
+import org.jetbrains.kotlin.psi.KtLabelReferenceExpression
+import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType

 class KotlinFirCompletionContributor : CompletionContributor() {
@@ -46,16 +47,25 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame

        val parentFunction = nameExpression.getNonStrictParentOfType<KtNamedFunction>() ?: return

-        val firFunction = LowLevelFirApiFacade.buildFunctionWithResolvedBody(originalFileFir, parentFunction)
+        val completionContext = LowLevelFirApiFacade.buildCompletionContextForFunction(originalFileFir, parentFunction)

        val element = nameExpression.getFirOfClosestParent() as? FirQualifiedAccessExpression ?: return

-        for (scope in getScopes(element, firFunction.session)) {
+        for (scope in getScopes(element, completionContext.session)) {
            for (name in scope.availableNames()) {
                result.addElement(LookupElementBuilder.create(name.asString()))
            }
        }

+        val towerDataContext = completionContext.getTowerDataContext(nameExpression)
+
+        for (localScope in towerDataContext.localScopes) {
+            for (name in localScope.getCallableNames()) {
+                result.addElement(LookupElementBuilder.create(name.asString()))
+            }
+        }
+
+
 //        val receiver = element.explicitReceiver
 //        if (receiver != null) {
 //