From dc56c5cf9e47ac86ba739194710b52df02f68924 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Thu, 20 Jul 2023 12:44:02 +0200
Subject: [PATCH] [LL FIR] FileStructure: avoid synchronizations on cache
 access

From `ConcurrentMap#compute`
>The entire method invocation is performed atomically.
>Some attempted update operations on this map by other threads may
>be blocked while computation is in progress, so the computation
>should be short and simple

And we can call resolution (`reanalyze()`) under this synchronized
block that can take unpredictable time.

This fix drops all heavy operations from synchronization
---
 .../api/fir/file/structure/FileStructure.kt   | 30 +++++++++++--------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructure.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructure.kt
index 37737ef4fa6..791c050bbbe 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructure.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructure.kt
@@ -87,26 +87,32 @@ internal class FileStructure private constructor(
     }
 
     private fun getStructureElementForDeclaration(declaration: KtElement): FileStructureElement {
-        val structureElement = structureElements.compute(declaration) { _, structureElement ->
-            when {
-                structureElement == null -> createStructureElement(declaration)
-                structureElement is ReanalyzableStructureElement<*, *> && !structureElement.isUpToDate() -> {
-                    structureElement.reanalyze()
-                }
-
-                else -> structureElement
-            }
+        val elementFromCache = structureElements[declaration]
+        if (elementFromCache == null) {
+            val newElement = createStructureElement(declaration)
+            return structureElements.putIfAbsent(declaration, newElement) ?: newElement
         }
 
-        return structureElement ?: errorWithFirSpecificEntries(
-            "FileStructureElement for was not defined for ${declaration::class.simpleName}",
+        if (elementFromCache !is ReanalyzableStructureElement<*, *> || elementFromCache.isUpToDate()) {
+            return elementFromCache
+        }
+
+        val reanalyzedElement = elementFromCache.reanalyze()
+        return structureElements.merge(declaration, reanalyzedElement) { _, oldElement ->
+            if (oldElement === elementFromCache) {
+                reanalyzedElement
+            } else {
+                // Another thread already reanalyzed the declaration
+                oldElement
+            }
+        } ?: errorWithFirSpecificEntries(
+            "${reanalyzedElement::class.simpleName} for ${declaration::class.simpleName} is gone",
             psi = declaration,
         )
     }
 
     fun getAllDiagnosticsForFile(diagnosticCheckerFilter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> {
         val structureElements = getAllStructureElements()
-
         return buildList {
             collectDiagnosticsFromStructureElements(structureElements, diagnosticCheckerFilter)
         }