[LL FIR] add cache restoring in some situations

In some caches, we assume that the resulting value must be not null if
we provide a context, so we can fix the cache inconsistency in this case.
This is error situations that shouldn't happen, but currently we have
problems with script dependencies (KT-61267), so such checks will help
us to distinguish some LL FIR problems from dependency problems.

^KT-61331 Fixed
^KTIJ-26670

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/FirCacheWithInvalidation.kt

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.caches
+
+import org.jetbrains.kotlin.fir.caches.FirCache
+
+internal interface FirCacheWithInvalidation<K : Any, V, CONTEXT> {
+    /**
+     * Drops the incorrect value from the cache and add a new value instead.
+     */
+    fun fixInconsistentValue(
+        key: K,
+        context: CONTEXT & Any,
+        inconsistencyMessage: String,
+        mapping: (oldValue: V, newValue: V & Any) -> V & Any,
+    ): V & Any
+}
+
+/**
+ * Return cached value or created a new one from [context].
+ * This method assumes that we can't return null for not-null context.
+ * Logs inconsistency error if it is present.
+ *
+ * @return not-null [VALUE] in case of [FirCacheWithInvalidation] cache.
+ */
+internal fun <KEY : Any, VALUE, CONTEXT> FirCache<KEY, VALUE, CONTEXT>.getNotNullValueForNotNullContext(
+    key: KEY,
+    context: CONTEXT,
+): VALUE {
+    val value = getValue(key, context)
+    @Suppress("CANNOT_CHECK_FOR_ERASED")
+    return if (value != null ||
+        context == null ||
+        this !is FirCacheWithInvalidation<KEY, VALUE, CONTEXT>
+    ) {
+        value
+    } else {
+        fixInconsistentValue(
+            key = key,
+            context = context,
+            inconsistencyMessage = "Inconsistency in the cache. Someone without context put a null value in the cache",
+            mapping = { old, new -> old ?: new },
+        )
+    }
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/FirThreadSafeCache.kt

@@ -1,23 +1,76 @@
 /*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.fir.caches
 
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.low.level.api.fir.caches.FirCacheWithInvalidation
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
+import org.jetbrains.kotlin.fir.utils.exceptions.withFirSymbolEntry
+import org.jetbrains.kotlin.utils.exceptions.ExceptionAttachmentBuilder
+import org.jetbrains.kotlin.utils.exceptions.checkWithAttachment
+import org.jetbrains.kotlin.utils.exceptions.logErrorWithAttachment
+import org.jetbrains.kotlin.utils.exceptions.withPsiEntry
 import java.util.concurrent.ConcurrentHashMap
 
 internal class FirThreadSafeCache<K : Any, V, CONTEXT>(
     private val map: ConcurrentHashMap<K, Any> = ConcurrentHashMap<K, Any>(),
-    private val createValue: (K, CONTEXT) -> V
-) : FirCache<K, V, CONTEXT>() {
+    private val createValue: (K, CONTEXT) -> V,
+) : FirCache<K, V, CONTEXT>(), FirCacheWithInvalidation<K, V, CONTEXT> {
+    override fun getValue(key: K, context: CONTEXT): V = map.getOrPutWithNullableValue(key) {
+        createValue(it, context)
+    }
 
-    override fun getValue(key: K, context: CONTEXT): V =
-        map.getOrPutWithNullableValue(key) {
-            createValue(it, context)
+    override fun getValueIfComputed(key: K): V? = map[key]?.nullValueToNull()
+
+    override fun fixInconsistentValue(
+        key: K,
+        context: CONTEXT & Any,
+        inconsistencyMessage: String,
+        mapping: (oldValue: V, newValue: V & Any) -> V & Any,
+    ): V & Any {
+        val newValue = createValue(key, context)
+        checkWithAttachment(
+            newValue != null,
+            message = { "A value for requested key & context must not be null due to the contract" },
+        ) {
+            buildAttachments(key, context, newValue)
         }
 
-    override fun getValueIfComputed(key: K): V? =
-        map[key]?.nullValueToNull()
+        LOG.logErrorWithAttachment(inconsistencyMessage) {
+            buildAttachments(key, context, newValue)
+        }
+
+        val result = map.merge(key, newValue) { old, _ ->
+            mapping(old.nullValueToNull(), newValue)
+        }
+
+        @Suppress("UNCHECKED_CAST")
+        return result as (V & Any)
+    }
+
+    private fun ExceptionAttachmentBuilder.buildAttachments(key: K, context: CONTEXT, value: V) {
+        withEntry("key", key.toString())
+
+        if (context is PsiElement) {
+            withPsiEntry("context", context)
+        } else {
+            withEntry("context", context.toString())
+        }
+
+        val unwrapped = (value as? Collection<*>)?.singleOrNull() ?: value
+        when (unwrapped) {
+            is FirElement -> withFirEntry("value", unwrapped)
+            is FirBasedSymbol<*> -> withFirSymbolEntry("value", unwrapped)
+            else -> withEntry("value", unwrapped.toString())
+        }
+    }
 }
+
+private val LOG = Logger.getInstance(FirThreadSafeCache::class.java)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt

@@ -1,11 +1,12 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
 
 import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.low.level.api.fir.caches.getNotNullValueForNotNullContext
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.CompositeKotlinPackageProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirResolveExtensionTool
@@ -115,7 +116,7 @@ internal class LLFirProviderHelper(
     ): FirClassLikeDeclaration? {
         if (classId.isLocal) return null
         if (!allowKotlinPackage && classId.isKotlinPackage()) return null
-        return classifierByClassId.getValue(classId, classLikeDeclaration)
+        return classifierByClassId.getNotNullValueForNotNullContext(classId, classLikeDeclaration)
     }
 
     fun getTopLevelClassNamesInPackage(packageFqName: FqName): Set<Name> {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirDeserializedSymbolProvider.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -7,18 +7,16 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.stubBased.deserializatio
 
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.low.level.api.fir.caches.getNotNullValueForNotNullContext
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirKotlinSymbolProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirKotlinSymbolNamesProvider
-import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createPackageProvider
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
-import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.deserialization.SingleModuleDataProvider
-import org.jetbrains.kotlin.fir.java.deserialization.JvmClassFileBasedSymbolProvider
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinBuiltins
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolNamesProvider
@@ -29,9 +27,8 @@ import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.stubs.impl.*
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
-import org.jetbrains.kotlin.serialization.deserialization.MetadataPackageFragment
 import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
-import java.util.ArrayList
+import org.jetbrains.kotlin.fir.caches.getValue
 
 typealias DeserializedTypeAliasPostProcessor = (FirTypeAliasSymbol) -> Unit
 
@@ -197,12 +194,12 @@ internal open class StubBasedFirDeserializedSymbolProvider(
         }
     }
 
-    private fun getClass(
-        classId: ClassId,
-        parentContext: StubBasedFirDeserializationContext? = null
-    ): FirRegularClassSymbol? {
-        return classCache.getValue(classId, parentContext)
-    }
+    private fun getClass(classId: ClassId, parentContext: StubBasedFirDeserializationContext? = null): FirRegularClassSymbol? =
+        if (parentContext?.classLikeDeclaration != null) {
+            classCache.getNotNullValueForNotNullContext(classId, parentContext)
+        } else {
+            classCache.getValue(classId, parentContext)
+        }
 
     private fun getTypeAlias(classId: ClassId, context: StubBasedFirDeserializationContext? = null): FirTypeAliasSymbol? {
         if (!classId.relativeClassName.isOneSegmentFQN()) return null
@@ -294,15 +291,11 @@ internal open class StubBasedFirDeserializedSymbolProvider(
             outerClassSymbol = null,
             outerTypeParameters = emptyList(),
             initialOrigin,
-            classLikeDeclaration
+            classLikeDeclaration,
         )
-        if (classLikeDeclaration is KtClassOrObject) {
-            return classCache.getValue(
-                classId,
-                deserializationContext
-            )
-        }
-        return typeAliasCache.getValue(classId, deserializationContext)
+
+        val cache = if (classLikeDeclaration is KtClassOrObject) classCache else typeAliasCache
+        return cache.getNotNullValueForNotNullContext(classId, deserializationContext)
     }
 
     fun getTopLevelCallableSymbol(