FIR IDE: add KotlinOutOfBlockPsiTreeChangePreprocessor

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/trackers/KotlinFirOutOfBlockModificationTracker.kt

@@ -18,7 +18,6 @@ import com.intellij.pom.event.PomModelEvent
 import com.intellij.pom.event.PomModelListener
 import com.intellij.pom.tree.TreeAspect
 import com.intellij.pom.tree.events.TreeChangeEvent
-import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingInBodyDeclarationWith
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileElementFactory
@@ -36,17 +35,12 @@ internal class KotlinFirModificationTrackerService(project: Project) : Disposabl
     fun getOutOfBlockModificationCountForModules(module: Module): Long =
         moduleModificationsState.getModificationsCountForModule(module)
 
-    @TestOnly
-    fun incrementModificationsCount() {
-        increaseModificationCountForAllModules()
-    }
-
     private val moduleModificationsState = ModuleModificationsState()
     private val treeAspect = TreeAspect.getInstance(project)
 
     override fun dispose() {}
 
-    private fun increaseModificationCountForAllModules() {
+    fun increaseModificationCountForAllModules() {
         projectGlobalOutOfBlockInKotlinFilesModificationCount++
         moduleModificationsState.increaseModificationCountForAllModules()
     }
@@ -130,7 +124,7 @@ private class ModuleModificationsState {
         modificationCountForModule.compute(module) { _, modifications ->
             when (modifications) {
                 null -> ModuleModifications(0, state)
-                else -> ModuleModifications(ModuleModifications(0, state).modificationsCount + 1, state)
+                else -> ModuleModifications(modifications.modificationsCount + 1, state)
             }
         }
     }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/trackers/KotlinFirOutOfBlockModificationTrackerFactory.kt

@@ -20,7 +20,7 @@ class KotlinFirOutOfBlockModificationTrackerFactory(private val project: Project
 
     @TestOnly
     fun incrementModificationsCount() {
-        project.service<KotlinFirModificationTrackerService>().incrementModificationsCount()
+        project.service<KotlinFirModificationTrackerService>().increaseModificationCountForAllModules()
     }
 }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/trackers/KotlinOutOfBlockPsiTreeChangePreprocessor.kt

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2020 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.idea.fir.low.level.api.trackers
+
+import com.intellij.openapi.components.service
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiDirectory
+import com.intellij.psi.PsiTreeChangeEvent
+import com.intellij.psi.impl.PsiModificationTrackerImpl
+import com.intellij.psi.impl.PsiTreeChangeEventImpl
+import com.intellij.psi.impl.PsiTreeChangePreprocessor
+
+class KotlinOutOfBlockPsiTreeChangePreprocessor(private val project: Project) : PsiTreeChangePreprocessor {
+    override fun treeChanged(event: PsiTreeChangeEventImpl) {
+        if (!PsiModificationTrackerImpl.canAffectPsi(event)) return
+        if (event.isOutOfBlockChange()) {
+            incrementModificationsCount()
+        }
+    }
+
+    private fun incrementModificationsCount() {
+        project.service<KotlinFirModificationTrackerService>().increaseModificationCountForAllModules()
+    }
+
+    // Copy logic from PsiModificationTrackerImpl.treeChanged(). Some out-of-code-block events are written to language modification
+    // tracker in PsiModificationTrackerImpl but don't have correspondent PomModelEvent. Increase kotlinOutOfCodeBlockTracker
+    // manually if needed.
+    private fun PsiTreeChangeEventImpl.isOutOfBlockChange() = when (code) {
+        PsiTreeChangeEventImpl.PsiEventType.PROPERTY_CHANGED ->
+            propertyName === PsiTreeChangeEvent.PROP_UNLOADED_PSI || propertyName === PsiTreeChangeEvent.PROP_ROOTS
+        PsiTreeChangeEventImpl.PsiEventType.CHILD_MOVED -> oldParent is PsiDirectory || newParent is PsiDirectory
+        else -> parent is PsiDirectory
+    }
+}

idea/resources-fir/META-INF/plugin.xml

@@ -113,6 +113,7 @@ The Kotlin FIR plugin provides language support in IntelliJ IDEA and Android Stu
 
     <projectService serviceImplementation="org.jetbrains.kotlin.idea.fir.low.level.api.trackers.KotlinFirModificationTrackerService"/>
     <projectService serviceImplementation="org.jetbrains.kotlin.idea.fir.low.level.api.trackers.KotlinFirOutOfBlockModificationTrackerFactory"/>
+    <psi.treeChangePreprocessor implementation="org.jetbrains.kotlin.idea.fir.low.level.api.trackers.KotlinOutOfBlockPsiTreeChangePreprocessor"/>
   </extensions>
 
   <!--  scripts  -->