Revise per module out-of-code-block (KT-32364)

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/IdeaModuleInfos.kt

@@ -31,11 +31,11 @@ import org.jetbrains.kotlin.config.SourceKotlinRootType
 import org.jetbrains.kotlin.config.TestSourceKotlinRootType
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.util.enlargedSearchScope
+import org.jetbrains.kotlin.idea.caches.trackers.KotlinModuleOutOfCodeBlockModificationTracker
 import org.jetbrains.kotlin.idea.configuration.BuildSystemType
 import org.jetbrains.kotlin.idea.configuration.getBuildSystemType
 import org.jetbrains.kotlin.idea.core.isInTestSourceContentKotlinAware
 import org.jetbrains.kotlin.idea.framework.getLibraryPlatform
-import org.jetbrains.kotlin.idea.caches.trackers.KotlinModuleModificationTracker
 import org.jetbrains.kotlin.idea.project.TargetPlatformDetector
 import org.jetbrains.kotlin.idea.project.findAnalyzerServices
 import org.jetbrains.kotlin.idea.project.getStableName
@@ -44,15 +44,15 @@ import org.jetbrains.kotlin.idea.util.isInSourceContentWithoutInjected
 import org.jetbrains.kotlin.idea.util.rootManager
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.DefaultIdeTargetPlatformKindProvider
-import org.jetbrains.kotlin.platform.idePlatformKind
-import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
-import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.compat.toOldPlatform
+import org.jetbrains.kotlin.platform.idePlatformKind
 import org.jetbrains.kotlin.platform.isCommon
 import org.jetbrains.kotlin.platform.js.isJs
+import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.platform.jvm.isJvm
 import org.jetbrains.kotlin.platform.konan.isNative
+import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
 import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.*
@@ -167,7 +167,7 @@ interface ModuleSourceInfo : IdeaModuleInfo, TrackableModuleInfo {
         get() = platform.findAnalyzerServices
 
     override fun createModificationTracker(): ModificationTracker =
-        KotlinModuleModificationTracker(module)
+        KotlinModuleOutOfCodeBlockModificationTracker(module)
 }
 
 sealed class ModuleSourceInfoWithExpectedBy(private val forProduction: Boolean) : ModuleSourceInfo {

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/trackers/KotlinCodeBlockModificationListener.kt

@@ -5,8 +5,6 @@
 
 package org.jetbrains.kotlin.idea.caches.trackers
 
-import com.intellij.openapi.module.Module
-import com.intellij.openapi.module.ModuleUtil
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.ModificationTracker
@@ -44,15 +42,7 @@ class KotlinCodeBlockModificationListener(
     project: Project,
     private val treeAspect: TreeAspect
 ): PsiTreeChangePreprocessor {
-    private val perModuleModCount = mutableMapOf<Module, Long>()
-
     private val modificationTrackerImpl = modificationTracker as PsiModificationTrackerImpl
-    private var lastAffectedModule: Module? = null
-
-    private var lastAffectedModuleModCount = -1L
-    // All modifications since that count are known to be single-module modifications reflected in
-    // perModuleModCount map
-    private var perModuleChangesHighWatermark: Long? = null
 
     @Volatile
     private var kotlinModificationTracker: Long = 0
@@ -65,11 +55,7 @@ class KotlinCodeBlockModificationListener(
 
     val kotlinOutOfCodeBlockTracker: ModificationTracker = kotlinOutOfCodeBlockTrackerImpl
 
-    fun getModificationCount(module: Module): Long {
-        return perModuleModCount[module] ?: perModuleChangesHighWatermark ?: kotlinOutOfCodeBlockTracker.modificationCount
-    }
-
-    fun hasPerModuleModificationCounts() = perModuleChangesHighWatermark != null
+    internal val perModuleOutOfCodeBlockTrackerUpdater = KotlinModuleOutOfCodeBlockModificationTracker.Updater(project)
 
     init {
         val model = PomManager.getModel(project)
@@ -94,10 +80,7 @@ class KotlinCodeBlockModificationListener(
 
                     if (ktFile.isPhysical && !isReplLine(ktFile.virtualFile)) {
                         kotlinOutOfCodeBlockTrackerImpl.incModificationCount()
-                        lastAffectedModule = ModuleUtil.findModuleForPsiElement(ktFile)
-                        lastAffectedModuleModCount = kotlinOutOfCodeBlockTracker.modificationCount
-                        // TODO: Revise perModuleModCount
-//                        modificationTrackerImpl.incCounter()
+                        perModuleOutOfCodeBlockTrackerUpdater.onKotlinPhysicalFileOutOfBlockChange(ktFile)
                     }
 
                     incOutOfBlockModificationCount(ktFile)
@@ -115,17 +98,7 @@ class KotlinCodeBlockModificationListener(
                 kotlinModificationTracker = kotlinTrackerInternalIDECount
             }
 
-            val newModCount = kotlinOutOfCodeBlockTracker.modificationCount
-            val affectedModule = lastAffectedModule
-            if (affectedModule != null && newModCount == lastAffectedModuleModCount + 1) {
-                if (perModuleChangesHighWatermark == null) {
-                    perModuleChangesHighWatermark = lastAffectedModuleModCount
-                }
-                perModuleModCount[affectedModule] = newModCount
-            } else {
-                perModuleChangesHighWatermark = null
-                perModuleModCount.clear()
-            }
+            perModuleOutOfCodeBlockTrackerUpdater.onPsiModificationTrackerUpdate()
         })
     }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/trackers/KotlinModuleOutOfCodeBlockModificationTracker.kt

@@ -6,23 +6,34 @@
 package org.jetbrains.kotlin.idea.caches.trackers
 
 import com.intellij.openapi.module.Module
+import com.intellij.openapi.module.ModuleUtil
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ModuleRootManager
 import com.intellij.openapi.roots.ProjectRootModificationTracker
 import com.intellij.openapi.util.ModificationTracker
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.util.CommonProcessors
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.caches.project.cached
+import org.jetbrains.kotlin.psi.KtFile
 
-class KotlinModuleModificationTracker(val module: Module) : ModificationTracker {
-    private val kotlinModCountListener =
-        KotlinCodeBlockModificationListener.getInstance(module.project)
+class KotlinModuleOutOfCodeBlockModificationTracker private constructor(private val module: Module, private val updater: Updater) :
+    ModificationTracker {
+
+    constructor(module: Module) :
+            this(module, KotlinCodeBlockModificationListener.getInstance(module.project).perModuleOutOfCodeBlockTrackerUpdater)
+
+    private val kotlinOutOfCodeBlockTracker = KotlinCodeBlockModificationListener.getInstance(module.project).kotlinOutOfCodeBlockTracker
 
     private val dependencies by lazy {
+        // Avoid implicit capturing for this to make CachedValueStabilityChecker happy
+        val module = module
+
         module.cached(CachedValueProvider {
             CachedValueProvider.Result.create(
-                HashSet<Module>().apply {
+                HashSet<Module>().also { resultModuleSet ->
                     ModuleRootManager.getInstance(module).orderEntries().recursively().forEachModule(
-                        CommonProcessors.CollectProcessor(this)
+                        CommonProcessors.CollectProcessor(resultModuleSet)
                     )
                 },
                 ProjectRootModificationTracker.getInstance(module.project)
@@ -31,15 +42,15 @@ class KotlinModuleModificationTracker(val module: Module) : ModificationTracker
     }
 
     override fun getModificationCount(): Long {
-        val currentGlobalCount = kotlinModCountListener.kotlinOutOfCodeBlockTracker.modificationCount
+        val currentGlobalCount = kotlinOutOfCodeBlockTracker.modificationCount
 
-        if (kotlinModCountListener.hasPerModuleModificationCounts()) {
-            val selfCount = kotlinModCountListener.getModificationCount(module)
+        if (updater.hasPerModuleModificationCounts()) {
+            val selfCount = updater.getModificationCount(module)
             if (selfCount == currentGlobalCount) return selfCount
 
             var maxCount = selfCount
             for (dependency in dependencies) {
-                val depCount = kotlinModCountListener.getModificationCount(dependency)
+                val depCount = updater.getModificationCount(dependency)
                 if (depCount == currentGlobalCount) return currentGlobalCount
                 if (depCount > maxCount) maxCount = depCount
             }
@@ -48,4 +59,61 @@ class KotlinModuleModificationTracker(val module: Module) : ModificationTracker
 
         return currentGlobalCount
     }
+
+    companion object {
+        @TestOnly
+        fun getModificationCount(module: Module): Long {
+            val updater = KotlinCodeBlockModificationListener.getInstance(module.project).perModuleOutOfCodeBlockTrackerUpdater
+            return updater.getModificationCount(module)
+        }
+    }
+
+    internal class Updater(project: Project) {
+        private val kotlinOfOfCodeBlockTracker by lazy {
+            KotlinCodeBlockModificationListener.getInstance(project).kotlinOutOfCodeBlockTracker
+        }
+
+        private val perModuleModCount = mutableMapOf<Module, Long>()
+
+        private var lastAffectedModule: Module? = null
+
+        private var lastAffectedModuleModCount = -1L
+
+        // All modifications since that count are known to be single-module modifications reflected in
+        // perModuleModCount map
+        private var perModuleChangesHighWatermark: Long? = null
+
+        internal fun getModificationCount(module: Module): Long {
+            return perModuleModCount[module] ?: perModuleChangesHighWatermark ?: kotlinOfOfCodeBlockTracker.modificationCount
+        }
+
+        internal fun hasPerModuleModificationCounts() = perModuleChangesHighWatermark != null
+
+        internal fun onKotlinPhysicalFileOutOfBlockChange(ktFile: KtFile) {
+            lastAffectedModule = ModuleUtil.findModuleForPsiElement(ktFile)
+            lastAffectedModuleModCount = kotlinOfOfCodeBlockTracker.modificationCount
+
+            onPsiModificationTrackerUpdate()
+        }
+
+        internal fun onPsiModificationTrackerUpdate() {
+            val newModCount = kotlinOfOfCodeBlockTracker.modificationCount
+            val affectedModule = lastAffectedModule
+            if (affectedModule != null && newModCount == lastAffectedModuleModCount) {
+                if (perModuleChangesHighWatermark == null) {
+                    perModuleChangesHighWatermark = lastAffectedModuleModCount
+                }
+                perModuleModCount[affectedModule] = newModCount
+            } else {
+                // Some updates were not processed in our code so they probably came from other languages. Invalidate all.
+                clean()
+            }
+        }
+
+        private fun clean() {
+            perModuleChangesHighWatermark = null
+            lastAffectedModule = null
+            perModuleModCount.clear()
+        }
+    }
 }