[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

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)
+    }
+}

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>(
+        FirElementFinder.findClassifierWithClassId(firFile, classId)
-            firFile,
-            canGoInside = { it is FirRegularClass },
-            predicate = { it.symbol.classId == classId },
-        )
             ?: error("Classifier $classId was found in file ${ktClass.containingKtFile.virtualFilePath} but was not found in FirFile")
     }
 

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 },