FIR resolution API: fix handling of expressions inside imports

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

@@ -97,6 +97,17 @@ fun KtClassOrObject.getOrBuildFir(
     return firClass
 }
 
+fun KtFile.getOrBuildFir(
+    state: FirResolveState,
+    phase: FirResolvePhase = FirResolvePhase.DECLARATIONS
+): FirFile {
+    val session = state.getSession(this)
+    val firProvider = FirProvider.getInstance(session) as IdeFirProvider
+    val firFile = firProvider.getOrBuildFile(this)
+    firFile.runResolve(firFile, firProvider, phase, state)
+    return firFile
+}
+
 private fun FirDeclaration.runResolve(
     file: FirFile,
     firProvider: IdeFirProvider,
@@ -108,21 +119,23 @@ private fun FirDeclaration.runResolve(
     if (toPhase > nonLazyPhase) {
         val designation = mutableListOf<FirElement>()
         designation += file
-        val id = when (this) {
+        if (this !is FirFile) {
-            is FirCallableDeclaration<*> -> {
+            val id = when (this) {
-                this.symbol.callableId.classId
+                is FirCallableDeclaration<*> -> {
+                    this.symbol.callableId.classId
+                }
+                is FirRegularClass -> {
+                    this.symbol.classId
+                }
+                else -> error("Unsupported: ${render()}")
             }
-            is FirRegularClass -> {
+            val outerClasses = generateSequence(id) { classId ->
-                this.symbol.classId
+                classId.outerClassId
+            }.mapTo(mutableListOf()) { firProvider.getFirClassifierByFqName(it)!! }
+            designation += outerClasses.asReversed()
+            if (this is FirCallableDeclaration<*>) {
+                designation += this
             }
-            else -> error("Unsupported: ${render()}")
-        }
-        val outerClasses = generateSequence(id) { classId ->
-            classId.outerClassId
-        }.mapTo(mutableListOf()) { firProvider.getFirClassifierByFqName(it)!! }
-        designation += outerClasses.asReversed()
-        if (this is FirCallableDeclaration<*>) {
-            designation += this
         }
         val transformer = FirDesignatedBodyResolveTransformer(
             designation.iterator(), state.getSession(psi as KtElement),
@@ -132,12 +145,12 @@ private fun FirDeclaration.runResolve(
     }
 }
 
-private fun KtElement.containingNonLocalDeclaration(): KtDeclaration {
+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")
+    return container
 }
 
 fun KtElement.getOrBuildFir(
@@ -148,6 +161,7 @@ fun KtElement.getOrBuildFir(
         when (val container = this.containingNonLocalDeclaration()) {
             is KtCallableDeclaration -> container.getOrBuildFir(state, phase)
             is KtClassOrObject -> container.getOrBuildFir(state, phase)
+            null -> containingKtFile.getOrBuildFir(state, phase)
             else -> error("Unsupported: ${container.text}")
         }
 

idea/testData/fir/lazyResolve/import/bar/bar.kt

@@ -0,0 +1,3 @@
+package bar
+
+fun barbar(): Int = 13

idea/testData/fir/lazyResolve/import/bar/bar.txt

@@ -0,0 +1,4 @@
+FILE: bar.kt
+    public final fun barbar(): R|kotlin/Int| {
+        ^barbar Int(13)
+    }

idea/testData/fir/lazyResolve/import/foo/main.kt

@@ -0,0 +1,7 @@
+package foo
+
+import bar.<caret>barbar
+
+fun foofoo(): Int {
+    return barbar()
+}

idea/testData/fir/lazyResolve/import/foo/main.txt

@@ -0,0 +1,4 @@
+FILE: main.kt
+    public final fun foofoo(): R|kotlin/Int| {
+        ^foofoo R|bar/barbar|()
+    }

idea/testData/fir/lazyResolve/import/import.test

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

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

@@ -0,0 +1 @@
+import bar.barbar

idea/tests/org/jetbrains/kotlin/idea/fir/AbstractFirLazyResolveTest.kt

@@ -13,6 +13,7 @@ import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.search.FileTypeIndex
 import com.intellij.testFramework.LightProjectDescriptor
 import junit.framework.TestCase
+import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.FirProvider
 import org.jetbrains.kotlin.idea.KotlinFileType
@@ -63,9 +64,23 @@ abstract class AbstractFirLazyResolveTest : KotlinLightCodeInsightFixtureTestCas
         } as KtExpression
 
         val resolveState = expressionToResolve.firResolveState()
-        val firExpression = expressionToResolve.getOrBuildFir(resolveState)
+        val resultsDump = when (val firElement = expressionToResolve.getOrBuildFir(resolveState)) {
-
+            is FirResolvedImport -> buildString {
-        val resultsDump = firExpression.render()
+                append("import ")
+                append(firElement.packageFqName)
+                val className = firElement.relativeClassName
+                if (className != null) {
+                    append("/")
+                    append(className)
+                }
+                val name = firElement.importedName
+                if (name != null) {
+                    append(".")
+                    append(name)
+                }
+            }
+            else -> firElement.render()
+        }
         KotlinTestUtils.assertEqualsToFile(File(testFile.parent, "results.txt"), resultsDump)
 
         fun expectedTxtPath(virtualFile: VirtualFile): String {
@@ -73,7 +88,7 @@ abstract class AbstractFirLazyResolveTest : KotlinLightCodeInsightFixtureTestCas
             var result: String? = null
             val root = File(path).parentFile
             for (file in root.walkTopDown()) {
-                if (!file.isDirectory && virtualPath in file.absolutePath) {
+                if (!file.isDirectory && virtualPath in file.absolutePath.replace("\\", "/")) {
                     result = file.absolutePath.replace(".kt", ".txt")
                 }
             }

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

@@ -34,6 +34,11 @@ public class FirLazyResolveTestGenerated extends AbstractFirLazyResolveTest {
         runTest("idea/testData/fir/lazyResolve/elvis/elvis.test");
     }
 
+    @TestMetadata("import/import.test")
+    public void testImport_Import() throws Exception {
+        runTest("idea/testData/fir/lazyResolve/import/import.test");
+    }
+
     @TestMetadata("inInit/inInit.test")
     public void testInInit_InInit() throws Exception {
         runTest("idea/testData/fir/lazyResolve/inInit/inInit.test");