[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

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 ->
+        val elementFromCache = structureElements[declaration]
-            when {
+        if (elementFromCache == null) {
-                structureElement == null -> createStructureElement(declaration)
+            val newElement = createStructureElement(declaration)
-                structureElement is ReanalyzableStructureElement<*, *> && !structureElement.isUpToDate() -> {
+            return structureElements.putIfAbsent(declaration, newElement) ?: newElement
-                    structureElement.reanalyze()
-                }
-
-                else -> structureElement
-            }
         }
 
-        return structureElement ?: errorWithFirSpecificEntries(
+        if (elementFromCache !is ReanalyzableStructureElement<*, *> || elementFromCache.isUpToDate()) {
-            "FileStructureElement for was not defined for ${declaration::class.simpleName}",
+            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)
         }