From 831f05a802b31a99b62fdb79085f9b5d7dd0d1ab Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Mon, 12 Jul 2021 16:47:40 +0200
Subject: [PATCH] LL API: fix kt -> fir mapping for KtImportList

---
 .../fir/low/level/api/FirModuleResolveStateDepended.kt |  7 +++++--
 .../low/level/api/element/builder/FirElementBuilder.kt | 10 ++++++++++
 .../testdata/getOrBuildFir/inImport/importList.kt      |  6 ++++++
 .../testdata/getOrBuildFir/inImport/importList.txt     |  5 +++++
 .../fir/low/level/api/GetOrBuildFirTestGenerated.java  |  6 ++++++
 5 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100644 idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.kt
 create mode 100644 idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.txt

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateDepended.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateDepended.kt
index ebbf519cea3..3fe2c97e0b9 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateDepended.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateDepended.kt
@@ -42,11 +42,14 @@ internal class FirModuleResolveStateDepended(
         originalState.getSessionFor(moduleInfo)
 
     override fun getOrBuildFirFor(element: KtElement): FirElement? {
-        val psi = originalState.elementBuilder.getPsiAsFirElementSource(element) ?: return null
+        val elementBuilder = originalState.elementBuilder
+
+        val psi = elementBuilder.getPsiAsFirElementSource(element) ?: return null
+        if (!elementBuilder.doKtElementHasCorrespondingFirElement(element)) return null
 
         ktToFirMapping?.getFirOfClosestParent(psi, this)?.let { return it }
 
-        return originalState.elementBuilder.getOrBuildFirFor(
+        return elementBuilder.getOrBuildFirFor(
             element = element,
             firFileBuilder = originalState.firFileBuilder,
             moduleFileCache = originalState.rootModuleSession.cache,
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt
index a7211c81a79..5880bfd2cd0 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt
@@ -52,6 +52,11 @@ internal class FirElementBuilder {
         }
     }
 
+    fun doKtElementHasCorrespondingFirElement(ktElement: KtElement): Boolean = when (ktElement) {
+        is KtImportList -> false
+        else -> true
+    }
+
     fun getOrBuildFirFor(
         element: KtElement,
         firFileBuilder: FirFileBuilder,
@@ -88,6 +93,11 @@ internal class FirElementBuilder {
         state: FirModuleResolveState,
     ): FirElement? {
         require(element !is KtFile)
+
+        if (!doKtElementHasCorrespondingFirElement(element)) {
+            return null
+        }
+
         val firFile = element.containingKtFile
         val fileStructure = fileStructureCache.getFileStructure(firFile, moduleFileCache)
 
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.kt
new file mode 100644
index 00000000000..8af65f0c191
--- /dev/null
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.kt
@@ -0,0 +1,6 @@
+// LOOK_UP_FOR_ELEMENT_OF_TYPE: org.jetbrains.kotlin.psi.KtImportList
+
+<expr>
+import a.b.c
+import b.c.d
+</expr>
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.txt b/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.txt
new file mode 100644
index 00000000000..efa858f1dd9
--- /dev/null
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.txt
@@ -0,0 +1,5 @@
+KT element: KtImportList
+FIR element: null
+
+FIR element rendered:
+null
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/GetOrBuildFirTestGenerated.java b/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/GetOrBuildFirTestGenerated.java
index 83c5feae100..093b34fa2d3 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/GetOrBuildFirTestGenerated.java
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/GetOrBuildFirTestGenerated.java
@@ -101,6 +101,12 @@ public class GetOrBuildFirTestGenerated extends AbstractGetOrBuildFirTest {
             runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/firstImportNamePart.kt");
         }
 
+        @Test
+        @TestMetadata("importList.kt")
+        public void testImportList() throws Exception {
+            runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/inImport/importList.kt");
+        }
+
         @Test
         @TestMetadata("middleImportNamePart.kt")
         public void testMiddleImportNamePart() throws Exception {