FIR resolution API: fix handling of expressions inside local functions

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt

@@ -15,8 +15,8 @@ import org.jetbrains.kotlin.fir.resolve.FirProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.scopes.impl.FirTopLevelDeclaredMemberScope
-import org.jetbrains.kotlin.fir.symbols.ConeCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.CallableId
-import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.name.ClassId
@@ -37,15 +37,14 @@ private fun KtClassOrObject.relativeFqName(): FqName {
 private fun FirFile.findCallableMember(
     provider: FirProvider, callableMember: KtCallableDeclaration,
     packageFqName: FqName, klassFqName: FqName?, declName: Name
-): FirCallableMemberDeclaration<*> {
+): FirCallableDeclaration<*> {
     val memberScope =
         if (klassFqName == null) FirTopLevelDeclaredMemberScope(this, session)
         else provider.getClassDeclaredMemberScope(ClassId(packageFqName, klassFqName, false))!!
-    var result: FirCallableMemberDeclaration<*>? = null
+    var result: FirCallableDeclaration<*>? = null
-    val processor = { symbol: ConeCallableSymbol ->
+    val processor = { symbol: FirCallableSymbol<*> ->
-        val firSymbol = symbol as? FirBasedSymbol<*>
+        val fir = symbol.fir
-        val fir = firSymbol?.fir as? FirCallableMemberDeclaration<*>
+        if (fir.psi == callableMember) {
-        if (fir?.psi == callableMember) {
             result = fir
             ProcessorAction.STOP
         } else {
@@ -58,13 +57,15 @@ private fun FirFile.findCallableMember(
         memberScope.processPropertiesByName(declName, processor)
     }
 
-    return result!!
+    return result
+        ?: error("Cannot find FIR callable declaration ${CallableId(packageFqName, klassFqName, declName)}")
 }
 
+// NB: at this moment it crashes with ISE when called on local declaration
 fun KtCallableDeclaration.getOrBuildFir(
     state: FirResolveState,
     phase: FirResolvePhase = FirResolvePhase.DECLARATIONS
-): FirCallableMemberDeclaration<*> {
+): FirCallableDeclaration<*> {
     val session = state.getSession(this)
 
     val file = this.containingKtFile
@@ -131,12 +132,20 @@ private fun FirDeclaration.runResolve(
     }
 }
 
+private fun KtElement.containingNonLocalDeclaration(): KtDeclaration {
+    var container = this.containingDeclarationForPseudocode
+    while (container != null && KtPsiUtil.isLocal(container)) {
+        container = container.containingDeclarationForPseudocode
+    }
+    return container ?: error("No containing non-local declaration: $text")
+}
+
 fun KtElement.getOrBuildFir(
     state: FirResolveState,
     phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
 ): FirElement {
     val containerFir: FirDeclaration =
-        when (val container = this.containingDeclarationForPseudocode ?: error("No containing declaration: $text")) {
+        when (val container = this.containingNonLocalDeclaration()) {
             is KtCallableDeclaration -> container.getOrBuildFir(state, phase)
             is KtClassOrObject -> container.getOrBuildFir(state, phase)
             else -> error("Unsupported: ${container.text}")

idea/testData/fir/lazyResolve/inLocal/inLocal.test

@@ -0,0 +1,3 @@
+{
+  "mainFile": "main.kt"
+}

idea/testData/fir/lazyResolve/inLocal/main.kt

@@ -0,0 +1,14 @@
+package main
+
+fun main() {
+    fun foo(): Int {
+        return <caret>bar()
+    }
+
+    foo()
+}
+
+fun bar(): Int {
+    val x = 4
+    return 9 * x
+}

idea/testData/fir/lazyResolve/inLocal/main.txt

@@ -0,0 +1,12 @@
+FILE: main.kt
+    public final fun main(): R|kotlin/Unit| {
+        local final fun foo(): R|kotlin/Int| {
+            ^foo R|main/bar|()
+        }
+
+        R|<local>/foo|()
+    }
+    public final fun bar(): R|kotlin/Int| {
+        lval x: <implicit> = Int(4)
+        ^bar Int(9).times#(x#)
+    }

idea/testData/fir/lazyResolve/inLocal/results.txt

@@ -0,0 +1 @@
+R|main/bar|()

idea/tests/org/jetbrains/kotlin/idea/fir/FirLazyResolveTestGenerated.java

@@ -29,6 +29,11 @@ public class FirLazyResolveTestGenerated extends AbstractFirLazyResolveTest {
         KotlinTestUtils.assertAllTestsPresentInSingleGeneratedClass(this.getClass(), new File("idea/testData/fir/lazyResolve"), Pattern.compile("^(.+)\\.test$"), TargetBackend.ANY);
     }
 
+    @TestMetadata("inLocal/inLocal.test")
+    public void testInLocal_InLocal() throws Exception {
+        runTest("idea/testData/fir/lazyResolve/inLocal/inLocal.test");
+    }
+
     @TestMetadata("simple/simple.test")
     public void testSimple_Simple() throws Exception {
         runTest("idea/testData/fir/lazyResolve/simple/simple.test");