From 913bfc0838d28cfc414cdecd9906242c4c70c48b Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Mon, 26 Sep 2022 15:31:14 +0200
Subject: [PATCH] [Low Level FIR] fix one case in CME when traversing FIR tree

Before extra child nodes of classes were visited during declaration by ClassId search.
If some other thread was mutating this tree, CME could happen

this is a fix of one case from ^KTIJ-21791
---
 .../errors/ExceptionAttachmentBuilder.kt      | 13 ++++++
 .../api/fir/providers/LLFirProviderHelper.kt  |  6 +--
 .../util/FirNonLocalDeclarationSearcher.kt    | 42 ++++++++++++++++++-
 3 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/errors/ExceptionAttachmentBuilder.kt b/analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/errors/ExceptionAttachmentBuilder.kt
index 8c5d8e5688e..a6a513219b5 100644
--- a/analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/errors/ExceptionAttachmentBuilder.kt
+++ b/analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/errors/ExceptionAttachmentBuilder.kt
@@ -81,3 +81,16 @@ public inline fun checkWithAttachmentBuilder(
         buildErrorWithAttachment(message(), buildAttachment = buildAttachment)
     }
 }
+
+@OptIn(ExperimentalContracts::class)
+public inline fun requireWithAttachmentBuilder(
+    condition: Boolean,
+    message: () -> String,
+    buildAttachment: ExceptionAttachmentBuilder.() -> Unit = {}
+) {
+    contract { returns() implies (condition) }
+
+    if (!condition) {
+        buildErrorWithAttachment(message(), buildAttachment = buildAttachment)
+    }
+}
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt
index 630c49812cf..9abe6251b75 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt
@@ -42,11 +42,7 @@ internal class LLFirProviderHelper(
             ?: return@createCache null
         if (ktClass.getClassId() == null) return@createCache null
         val firFile = firFileBuilder.buildRawFirFileWithCaching(ktClass.containingKtFile)
-        FirElementFinder.findElementIn<FirClassLikeDeclaration>(
-            firFile,
-            canGoInside = { it is FirRegularClass },
-            predicate = { it.symbol.classId == classId },
-        )
+        FirElementFinder.findClassifierWithClassId(firFile, classId)
             ?: error("Classifier $classId was found in file ${ktClass.containingKtFile.virtualFilePath} but was not found in FirFile")
     }
 
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt
index 32b1c60bfb0..c4aeb7f912e 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt
@@ -5,15 +5,55 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 
+import org.jetbrains.kotlin.analysis.utils.errors.requireWithAttachmentBuilder
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.utils.classId
+import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
-import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.name.ClassId
 
 object FirElementFinder {
+    fun findClassifierWithClassId(firFile: FirFile, classId: ClassId): FirClassLikeDeclaration? {
+        requireWithAttachmentBuilder(!classId.isLocal, { "ClassId should not be local" }) {
+            withEntry("classId", classId) { it.asString() }
+        }
+        requireWithAttachmentBuilder(
+            firFile.packageFqName == classId.packageFqName,
+            { "ClassId should not be local" }
+        ) {
+            withEntry("FirFile.packageName", firFile.packageFqName) { it.asString() }
+            withEntry("ClassId.packageName", classId.packageFqName) { it.asString() }
+        }
+
+        val classIdPathSegment = classId.relativeClassName.pathSegments()
+        var result: FirClassLikeDeclaration? = null
+
+        fun find(declarations: List<FirDeclaration>, classIdPathIndex: Int) {
+            if (result != null) return
+            val currentClassSegment = classIdPathSegment[classIdPathIndex]
+
+            for (subDeclaration in declarations) {
+                if (subDeclaration is FirClassLikeDeclaration && currentClassSegment == subDeclaration.symbol.name) {
+                    if (classIdPathIndex == classIdPathSegment.lastIndex) {
+                        result = subDeclaration
+                        return
+                    }
+                    if (subDeclaration is FirRegularClass) {
+                        find(subDeclaration.declarations, classIdPathIndex + 1)
+                    }
+                }
+            }
+        }
+
+        find(firFile.declarations, classIdPathIndex = 0)
+        return result
+    }
+
     inline fun <reified E : FirElement> findElementIn(
         container: FirElement,
         crossinline canGoInside: (E) -> Boolean = { true },