From 54946a793a330074b639919fb006f65b60414125 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Thu, 9 Jul 2020 00:06:07 +0300
Subject: [PATCH] Fix `FirCompletionContext` for safe access elements

- `towerDataContextForStatement` may contain safe access (?.) fir
elements as keys; however, `state.getCachedMapping(current)` would
never return those since they are considered as sugar fir elements

Co-authored-by: Ilya Kirillov <ilya.kirillov@jetbrains.com>
---
 .../idea/fir/low/level/api/LowLevelFirApiFacade.kt     | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/LowLevelFirApiFacade.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/LowLevelFirApiFacade.kt
index bc32a7dc54a..23f365404d8 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/LowLevelFirApiFacade.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/LowLevelFirApiFacade.kt
@@ -9,10 +9,10 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtNamedFunction
@@ -36,7 +36,7 @@ object LowLevelFirApiFacade {
 
     class FirCompletionContext internal constructor(
         val session: FirSession,
-        private val towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>,
+        private val towerDataContextForStatement: Map<FirStatement, FirTowerDataContext>,
         private val state: FirModuleResolveState,
     ) {
         fun getTowerDataContext(element: KtElement): FirTowerDataContext {
@@ -46,6 +46,12 @@ object LowLevelFirApiFacade {
 
                 if (mappedFir is FirStatement) {
                     towerDataContextForStatement[mappedFir]?.let { return it }
+                    for ((key, value) in towerDataContextForStatement) {
+                        // TODO Rework that
+                        if (key.psi == mappedFir.psi) {
+                            return value
+                        }
+                    }
                 }
                 current = current.parent
             }