Implement FirProvider consistency checker

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/impl/FirProviderImpl.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.impl
 
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
@@ -24,7 +25,7 @@ import org.jetbrains.kotlin.name.Name
 
 class FirProviderImpl(val session: FirSession) : FirProvider {
     override fun getFirCallableContainerFile(symbol: ConeCallableSymbol): FirFile? {
-        return callableContainerMap[symbol]
+        return state.callableContainerMap[symbol]
     }
 
     override fun getClassLikeSymbolByFqName(classId: ClassId): ConeClassLikeSymbol? {
@@ -32,39 +33,39 @@ class FirProviderImpl(val session: FirSession) : FirProvider {
     }
 
     override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<ConeCallableSymbol> {
-        return (callableMap[CallableId(packageFqName, null, name)] ?: emptyList())
+        return (state.callableMap[CallableId(packageFqName, null, name)] ?: emptyList())
     }
 
     override fun getClassDeclaredMemberScope(classId: ClassId) =
         (getFirClassifierByFqName(classId) as? FirRegularClass)?.let(::FirClassDeclaredMemberScope)
 
     override fun getFirClassifierContainerFile(fqName: ClassId): FirFile {
-        return classifierContainerFileMap[fqName] ?: error("Couldn't find container for $fqName")
+        return state.classifierContainerFileMap[fqName] ?: error("Couldn't find container for $fqName")
     }
 
     fun recordFile(file: FirFile) {
+        recordFile(file, state)
+    }
+
+    private fun recordFile(file: FirFile, state: State) = with(state) {
         val packageName = file.packageFqName
         fileMap.merge(packageName, listOf(file)) { a, b -> a + b }
 
         file.acceptChildren(object : FirVisitorVoid() {
             override fun visitElement(element: FirElement) {}
 
-            var containerFqName: FqName = FqName.ROOT
 
             override fun visitRegularClass(regularClass: FirRegularClass) {
-                val fqName = containerFqName.child(regularClass.name)
-                val classId = ClassId(packageName, fqName, false)
+                val classId = regularClass.symbol.classId
+
                 classifierMap[classId] = regularClass
                 classifierContainerFileMap[classId] = file
 
-                containerFqName = fqName
                 regularClass.acceptChildren(this)
-                containerFqName = fqName.parent()
             }
 
             override fun visitTypeAlias(typeAlias: FirTypeAlias) {
-                val fqName = containerFqName.child(typeAlias.name)
-                val classId = ClassId(packageName, fqName, false)
+                val classId = typeAlias.symbol.classId
                 classifierMap[classId] = typeAlias
                 classifierContainerFileMap[classId] = file
             }
@@ -90,18 +91,75 @@ class FirProviderImpl(val session: FirSession) : FirProvider {
         })
     }
 
-    private val fileMap = mutableMapOf<FqName, List<FirFile>>()
-    private val classifierMap = mutableMapOf<ClassId, FirClassLikeDeclaration>()
-    private val classifierContainerFileMap = mutableMapOf<ClassId, FirFile>()
-    private val callableMap = mutableMapOf<CallableId, List<ConeCallableSymbol>>()
-    private val callableContainerMap = mutableMapOf<ConeCallableSymbol, FirFile>()
+    private val state = State()
+
+    private class State {
+        val fileMap = mutableMapOf<FqName, List<FirFile>>()
+        val classifierMap = mutableMapOf<ClassId, FirClassLikeDeclaration>()
+        val classifierContainerFileMap = mutableMapOf<ClassId, FirFile>()
+        val callableMap = mutableMapOf<CallableId, List<ConeCallableSymbol>>()
+        val callableContainerMap = mutableMapOf<ConeCallableSymbol, FirFile>()
+    }
 
     override fun getFirFilesByPackage(fqName: FqName): List<FirFile> {
-        return fileMap[fqName].orEmpty()
+        return state.fileMap[fqName].orEmpty()
     }
 
     override fun getFirClassifierByFqName(fqName: ClassId): FirClassLikeDeclaration? {
-        return classifierMap[fqName]
+        return state.classifierMap[fqName]
+    }
+
+    @TestOnly
+    fun ensureConsistent(files: List<FirFile>) {
+        val newState = State()
+        files.forEach { recordFile(it, newState) }
+
+        fun <K, V> checkMapDiff(a: Map<K, V>, b: Map<K, V>, equal: (old: V?, new: V?) -> Boolean = { old, new -> old === new }) {
+            val unionKeys = a.keys + b.keys
+            val failures = mutableListOf<String>()
+            for ((key, aValue, bValue) in unionKeys.map { Triple(it, a[it], b[it]) }) {
+                if (!equal(aValue, bValue)) {
+                    failures += "diff at key = '$key': was: '$aValue', become: '$bValue'"
+                }
+            }
+
+            assert(failures.isEmpty()) {
+                failures.joinToString(separator = "\n")
+            }
+        }
+
+        fun <K, V> checkMMapDiff(a: Map<K, List<V>>, b: Map<K, List<V>>) {
+            val unionKeys = a.keys + b.keys
+            val failures = mutableListOf<String>()
+            for ((key, aValue, bValue) in unionKeys.map { Triple(it, a[it], b[it])}) {
+                if (aValue == null || bValue == null) {
+                    failures += "diff at key = '$key': was: $aValue, become: $bValue"
+                } else {
+                    val aSet = aValue.toSet()
+                    val bSet = bValue.toSet()
+
+                    val aLost = aSet - bSet
+                    val bNew = bSet - aSet
+                    if (aLost.isNotEmpty() || bNew.isNotEmpty()) {
+                        failures += "diff at key = '$key':"
+                        failures += "    Lost:"
+                        aLost.forEach { failures += "     $it" }
+                        failures += "    New:"
+                        bNew.forEach { failures += "     $it" }
+                    }
+                }
+            }
+
+            assert(failures.isEmpty()) {
+                failures.joinToString(separator = "\n")
+            }
+        }
+
+        checkMMapDiff(state.fileMap, newState.fileMap)
+        checkMapDiff(state.classifierMap, newState.classifierMap)
+        checkMapDiff(state.classifierContainerFileMap, newState.classifierContainerFileMap)
+        checkMMapDiff(state.callableMap, newState.callableMap)
+        checkMapDiff(state.callableContainerMap, newState.callableContainerMap)
     }
 
     override fun getClassUseSiteMemberScope(classId: ClassId, useSiteSession: FirSession): FirScope? {

compiler/tests-common/tests/org/jetbrains/kotlin/fir/FirResolveBench.kt

@@ -7,6 +7,8 @@ package org.jetbrains.kotlin.fir
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.FirProvider
+import org.jetbrains.kotlin.fir.resolve.impl.FirProviderImpl
 import org.jetbrains.kotlin.fir.types.ConeClassErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinErrorType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
@@ -19,6 +21,14 @@ import kotlin.math.max
 import kotlin.reflect.KClass
 import kotlin.system.measureNanoTime
 
+
+fun checkFirProvidersConsistency(firFiles: List<FirFile>) {
+    for ((session, files) in firFiles.groupBy { it.session }) {
+        val provider = session.service<FirProvider>() as FirProviderImpl
+        provider.ensureConsistent(files)
+    }
+}
+
 fun doFirResolveTestBench(
     firFiles: List<FirFile>,
     transformers: List<FirTransformer<Nothing?>>,
@@ -62,6 +72,7 @@ fun doFirResolveTestBench(
                 }
                 //totalLength += StringBuilder().apply { FirRenderer(this).visitFile(firFile) }.length
             }
+            checkFirProvidersConsistency(firFiles)
         }
 
         if (fails.none()) {