FIR IDE: fix local class containing file search

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

@@ -30,10 +30,10 @@ abstract class FirProvider : FirSymbolProvider() {
 
     abstract fun getFirClassifierContainerFileIfAny(fqName: ClassId): FirFile?
 
-    fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile =
+    open fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile =
         getFirClassifierContainerFile(symbol.classId)
 
-    fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? =
+    open fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? =
         getFirClassifierContainerFileIfAny(symbol.classId)
 
     abstract fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile?

compiler/fir/resolve/testData/resolve/expresssions/localObjects.kt

@@ -0,0 +1,21 @@
+object A {
+    val x : Int = 0
+}
+
+open class Foo {
+    fun foo() : Int = 1
+}
+
+fun test() {
+    A.x
+    val b = object : Foo() {
+    }
+    b.foo()
+
+    object B {
+        fun foo() {}
+    }
+    B.foo()
+}
+
+val bb = <!UNRESOLVED_REFERENCE!>B<!>.<!UNRESOLVED_REFERENCE!>foo<!>()

compiler/fir/resolve/testData/resolve/expresssions/localObjects.txt

@@ -0,0 +1,44 @@
+FILE: localObjects.kt
+    public final object A : R|kotlin/Any| {
+        private constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final val x: R|kotlin/Int| = Int(0)
+            public get(): R|kotlin/Int|
+
+    }
+    public open class Foo : R|kotlin/Any| {
+        public constructor(): R|Foo| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final fun foo(): R|kotlin/Int| {
+            ^foo Int(1)
+        }
+
+    }
+    public final fun test(): R|kotlin/Unit| {
+        Q|A|.R|/A.x|
+        lval b: R|anonymous| = object : R|Foo| {
+            private constructor(): R|Foo| {
+                super<R|Foo|>()
+            }
+
+        }
+
+        R|<local>/b|.R|/Foo.foo|()
+        local final object B : R|kotlin/Any| {
+            private constructor(): R|B| {
+                super<R|kotlin/Any|>()
+            }
+
+            public final fun foo(): R|kotlin/Unit| {
+            }
+
+        }
+
+        Q|B|.R|/B.foo|()
+    }
+    public final val bb: <ERROR TYPE REF: Unresolved name: foo> = <Unresolved name: B>#.<Unresolved name: foo>#()
+        public get(): <ERROR TYPE REF: Unresolved name: foo>

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -489,6 +489,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/expresssions/localInnerClass.kt");
         }
 
+        @TestMetadata("localObjects.kt")
+        public void testLocalObjects() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/expresssions/localObjects.kt");
+        }
+
         @TestMetadata("localScopes.kt")
         public void testLocalScopes() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/expresssions/localScopes.kt");

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.FirProvider
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.impl.FirProviderImpl
@@ -21,6 +22,7 @@ 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.KtFile
 
 class IdeFirProvider(
@@ -81,6 +83,20 @@ class IdeFirProvider(
         return cacheProvider.getFirClassifierContainerFileIfAny(fqName)
     }
 
+    override fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile {
+        return getFirClassifierContainerFileIfAny(symbol)
+            ?: error("Couldn't find container for ${symbol.classId}")
+    }
+
+    override fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? {
+        val psi = symbol.fir.source?.psi
+        if (psi is KtClassOrObject && psi.isLocal) {
+            val ktFile = psi.containingKtFile
+            return getOrBuildFile(ktFile)
+        }
+        return getFirClassifierContainerFileIfAny(symbol.classId)
+    }
+
     override fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile? {
         return cacheProvider.getFirCallableContainerFile(symbol)
     }