From 04fca93b1e00c56354a1559dc254c5b22a11cf8d Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Wed, 16 Sep 2020 13:06:18 +0300
Subject: [PATCH] FIR IDE: fix completion on parenthesized expression

---
 .../api/FirModuleResolveStateForCompletion.kt |  2 +-
 .../api/element/builder/FirElementBuilder.kt  | 29 ++++++++++++-------
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt
index 59ae6ddd4c8..9063ab47590 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt
@@ -32,7 +32,7 @@ internal class FirModuleResolveStateForCompletion(
         originalState.getSessionFor(moduleInfo)
 
     override fun getOrBuildFirFor(element: KtElement, toPhase: FirResolvePhase): FirElement {
-        completionMapping[element]?.let { return it }
+        completionMapping[originalState.elementBuilder.getPsiAsFirElementSource(element)]?.let { return it }
         return originalState.elementBuilder.getOrBuildFirFor(
             element,
             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 0f2a4cb8875..f20d11fb637 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
@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureC
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import org.jetbrains.kotlin.psi2ir.deparenthesize
 
 /**
  * Maps [KtElement] to [FirElement]
@@ -23,6 +24,22 @@ import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
  */
 @ThreadSafe
 internal class FirElementBuilder {
+    fun getPsiAsFirElementSource(element: KtElement): KtElement {
+        val deparenthesized = if (element is KtPropertyDelegate) element.deparenthesize() else element
+        return when {
+            deparenthesized is KtParenthesizedExpression -> deparenthesized.deparenthesize()
+            deparenthesized is KtPropertyDelegate -> deparenthesized.expression ?: element
+            deparenthesized is KtQualifiedExpression && deparenthesized.selectorExpression is KtCallExpression -> {
+                /*
+                 KtQualifiedExpression with KtCallExpression in selector transformed in FIR to FirFunctionCall expression
+                 Which will have a receiver as qualifier
+                 */
+                deparenthesized.selectorExpression ?: error("Incomplete code:\n${element.getElementTextInContext()}")
+            }
+            else -> deparenthesized
+        }
+    }
+
     fun getOrBuildFirFor(
         element: KtElement,
         moduleFileCache: ModuleFileCache,
@@ -31,17 +48,7 @@ internal class FirElementBuilder {
         val fileStructure = fileStructureCache.getFileStructure(element.containingKtFile, moduleFileCache)
         val mappings = fileStructure.getStructureElementFor(element).mappings
 
-        val psi = when {
-            element is KtPropertyDelegate -> element.expression ?: element
-            element is KtQualifiedExpression && element.selectorExpression is KtCallExpression -> {
-                /*
-                 KtQualifiedExpression with KtCallExpression in selector transformed in FIR to FirFunctionCall expression
-                 Which will have a receiver as qualifier
-                 */
-                element.selectorExpression ?: error("Incomplete code:\n${element.getElementTextInContext()}")
-            }
-            else -> element
-        }
+        val psi = getPsiAsFirElementSource(element)
         mappings[psi]?.let { return it }
         return psi.getFirOfClosestParent(mappings)?.second
             ?: error("FirElement is not found for:\n${element.getElementTextInContext()}")