From b4aa2aef002a8b7f1f066f18397df52cbbfbbf69 Mon Sep 17 00:00:00 2001
From: Vladimir Dolzhenko <vladimir.dolzhenko@jetbrains.com>
Date: Fri, 9 Dec 2022 12:04:50 +0000
Subject: [PATCH] Add ability to invalidate AbstractResolverForProject

#KTIJ-23785

Merge-request: KT-MR-7900
Merged-by: Vladimir Dolzhenko <Vladimir.Dolzhenko@jetbrains.com>
---
 .../analyzer/AbstractResolverForProject.kt    | 37 ++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/analyzer/AbstractResolverForProject.kt b/compiler/frontend/src/org/jetbrains/kotlin/analyzer/AbstractResolverForProject.kt
index 8a6b6db947b..5f4f6ab68c3 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/analyzer/AbstractResolverForProject.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/analyzer/AbstractResolverForProject.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analyzer
 
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.context.ProjectContext
@@ -22,7 +23,7 @@ abstract class AbstractResolverForProject<M : ModuleInfo>(
     protected val fallbackModificationTracker: ModificationTracker? = null,
     private val delegateResolver: ResolverForProject<M> = EmptyResolverForProject(),
     private val packageOracleFactory: PackageOracleFactory = PackageOracleFactory.OptimisticFactory
-) : ResolverForProject<M>() {
+) : ResolverForProject<M>(), Disposable {
 
     protected class ModuleData(
         val moduleDescriptor: ModuleDescriptorImpl,
@@ -36,6 +37,9 @@ abstract class AbstractResolverForProject<M : ModuleInfo>(
         }
     }
 
+    @Volatile
+    protected var disposed = false
+
     // Protected by ("projectContext.storageManager.lock")
     protected val descriptorByModule = hashMapOf<M, ModuleData>()
 
@@ -55,6 +59,7 @@ abstract class AbstractResolverForProject<M : ModuleInfo>(
     abstract fun builtInsForModule(module: M): KotlinBuiltIns
     abstract fun createResolverForModule(descriptor: ModuleDescriptor, moduleInfo: M): ResolverForModule
     override fun tryGetResolverForModule(moduleInfo: M): ResolverForModule? {
+        checkValid()
         if (!isCorrectModuleInfo(moduleInfo)) {
             return null
         }
@@ -62,6 +67,7 @@ abstract class AbstractResolverForProject<M : ModuleInfo>(
     }
 
     private fun setupModuleDescriptor(module: M, moduleDescriptor: ModuleDescriptorImpl) {
+        checkValid()
         moduleDescriptor.setDependencies(
             LazyModuleDependencies(
                 projectContext.storageManager,
@@ -124,7 +130,9 @@ abstract class AbstractResolverForProject<M : ModuleInfo>(
      */
     private fun resolverForModuleDescriptorImpl(descriptor: ModuleDescriptor): ResolverForModule? {
         return projectContext.storageManager.compute {
+            checkValid()
             descriptor.assertValid()
+
             val module = moduleInfoByDescriptor[descriptor]
             if (module == null) {
                 if (delegateResolver is EmptyResolverForProject<*>) {
@@ -148,11 +156,13 @@ abstract class AbstractResolverForProject<M : ModuleInfo>(
         }
 
     override fun descriptorForModule(moduleInfo: M): ModuleDescriptorImpl {
+        checkValid()
         checkModuleIsCorrect(moduleInfo)
         return doGetDescriptorForModule(moduleInfo)
     }
 
     override fun moduleInfoForModuleDescriptor(moduleDescriptor: ModuleDescriptor): M {
+        checkValid()
         return moduleInfoByDescriptor[moduleDescriptor] ?: delegateResolver.moduleInfoForModuleDescriptor(moduleDescriptor)
     }
 
@@ -215,6 +225,29 @@ abstract class AbstractResolverForProject<M : ModuleInfo>(
         return ModuleData(moduleDescriptor, modificationTracker)
     }
 
+    private fun checkValid() {
+        if (disposed) {
+            reportInvalidResolver()
+        }
+    }
+
+    protected open fun reportInvalidResolver() {
+        throw InvalidResolverException("$name is invalidated")
+    }
+
+    override fun dispose() {
+        projectContext.storageManager.compute {
+            disposed = true
+            descriptorByModule.values.forEach {
+                moduleInfoByDescriptor.remove(it.moduleDescriptor)
+                it.moduleDescriptor.isValid = false
+            }
+            descriptorByModule.clear()
+            moduleInfoByDescriptor.keys.forEach { it.isValid = false }
+            moduleInfoByDescriptor.clear()
+        }
+    }
+
     private fun renderResolversChainContents(): String {
         val resolversChain = generateSequence(this) { it.delegateResolver as? AbstractResolverForProject<M> }
 
@@ -334,3 +367,5 @@ private object DiagnoseUnknownModuleInfoReporter {
 
     private fun otherError(message: String) = KotlinExceptionWithAttachments(message)
 }
+
+class InvalidResolverException(message: String) : IllegalStateException(message)
\ No newline at end of file