Expose local scopes to use in completion

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/PersistentMultimap.kt

@@ -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()
     }
-}
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt

@@ -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> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -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) {

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

@@ -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
 }

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

@@ -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
+    }
 }
 
 

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

@@ -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 {
+        val function = builtFunction.apply {
-            runResolve(firFile, firIdeProvider, phase, state)
+            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)
     }
-}
+}

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

@@ -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.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByTypeRef
+import org.jetbrains.kotlin.fir.resolve.scope
-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.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) {
 //