FIR IDE API: add preliminary control of file in-block modifications

In this commit, we build FirFile from the beginning if in-block modifications are detected.
However, better way would be to rebuild only part of FirFile in this case

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt

@@ -142,6 +142,7 @@ fun KtFile.getOrBuildFirWithDiagnostics(state: FirResolveState): FirFile {
         }
     }
 
+    ProgressIndicatorProvider.checkCanceled() // ???
     if (state.hasDiagnosticsForFile(this)) return firFile
 
     FirIdeDiagnosticsCollector(state).collectDiagnostics(firFile)

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolveState.kt

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.idea.fir
 
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.FirFile
@@ -75,7 +77,7 @@ class FirResolveStateImpl(override val sessionProvider: FirSessionProvider) : Fi
 
     private val diagnosticCache = mutableMapOf<KtElement, MutableList<Diagnostic>>()
 
-    private val diagnosedFiles = mutableSetOf<KtFile>()
+    private val diagnosedFiles = mutableMapOf<KtFile, Long>()
 
     override fun get(psi: KtElement): FirElement? = cache[psi]
 
@@ -84,7 +86,20 @@ class FirResolveStateImpl(override val sessionProvider: FirSessionProvider) : Fi
     }
 
     override fun hasDiagnosticsForFile(file: KtFile): Boolean {
-        return file in diagnosedFiles
+        val previousStamp = diagnosedFiles[file] ?: return false
+        if (file.modificationStamp == previousStamp) {
+            return true
+        }
+        diagnosedFiles.remove(file)
+        file.accept(object : PsiElementVisitor() {
+            override fun visitElement(element: PsiElement) {
+                cache.remove(element)
+                diagnosticCache.remove(element)
+                element.acceptChildren(this)
+                super.visitElement(element)
+            }
+        })
+        return false
     }
 
     override fun record(psi: KtElement, fir: FirElement) {
@@ -101,7 +116,7 @@ class FirResolveStateImpl(override val sessionProvider: FirSessionProvider) : Fi
             (diagnostic.source.psi as? KtElement)?.let { record(it, diagnostic.diagnostic) }
         }
 
-        diagnosedFiles += file
+        diagnosedFiles[file] = file.modificationStamp
     }
 }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/IdeFirProvider.kt

@@ -32,8 +32,10 @@ class IdeFirProvider(
     val session: FirSession
 ) : FirProvider() {
     private val cacheProvider = FirProviderImpl(session)
-    // TODO: invalidation?
+
-    private val files = mutableMapOf<KtFile, FirFile>()
+    data class FirFileWithStamp(val file: FirFile, val stamp: Long)
+
+    private val files = mutableMapOf<KtFile, FirFileWithStamp>()
 
     override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? {
         return cacheProvider.getFirClassifierByFqName(classId) ?: run {
@@ -54,17 +56,31 @@ class IdeFirProvider(
     }
 
     fun getOrBuildFile(ktFile: KtFile): FirFile {
-        files[ktFile]?.let { return it }
+        val modificationStamp = ktFile.modificationStamp
+        files[ktFile]?.let { (firFile, stamp) ->
+            if (stamp == modificationStamp) {
+                return firFile
+            }
+        }
         return synchronized(ktFile) {
-            files.getOrPut(ktFile) {
+            var fileWithStamp = files[ktFile]
+            if (fileWithStamp != null && fileWithStamp.stamp == modificationStamp) {
+                fileWithStamp.file
+            } else {
                 val file = RawFirBuilder(session, stubMode = false).buildFirFile(ktFile)
                 cacheProvider.recordFile(file)
+                fileWithStamp = FirFileWithStamp(file, modificationStamp)
+                files[ktFile] = fileWithStamp
                 file
             }
         }
     }
 
-    fun getFile(ktFile: KtFile): FirFile? = files[ktFile]
+    fun getFile(ktFile: KtFile): FirFile? {
+        val (firFile, stamp) = files[ktFile] ?: return null
+        if (stamp == ktFile.modificationStamp) return firFile
+        return null
+    }
 
     override fun getClassLikeSymbolByFqName(classId: ClassId): FirClassLikeSymbol<*>? {
         return getFirClassifierByFqName(classId)?.symbol