From a5522388740721be5e329185c95ce12181ec39af Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Mon, 11 Mar 2024 11:19:00 +0200
Subject: [PATCH] [Plugin API] Provide a way to register disposable callbacks
 in CompilerPluginRegistrar

^KT-59555 Fixed
---
 .../plugin/CompilerPluginRegistrar.kt         | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/compiler/plugin-api/src/org/jetbrains/kotlin/compiler/plugin/CompilerPluginRegistrar.kt b/compiler/plugin-api/src/org/jetbrains/kotlin/compiler/plugin/CompilerPluginRegistrar.kt
index 60dc8ae1113..0537417f3c0 100644
--- a/compiler/plugin-api/src/org/jetbrains/kotlin/compiler/plugin/CompilerPluginRegistrar.kt
+++ b/compiler/plugin-api/src/org/jetbrains/kotlin/compiler/plugin/CompilerPluginRegistrar.kt
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.compiler.plugin
 
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.CompilerConfigurationKey
@@ -25,9 +26,24 @@ abstract class CompilerPluginRegistrar {
         val registeredExtensions: Map<ProjectExtensionDescriptor<*>, List<Any>>
             get() = _registeredExtensions
 
+        private val _disposables = mutableListOf<PluginDisposable>()
+        val disposables: List<PluginDisposable>
+            get() = _disposables
+
         fun <T : Any> ProjectExtensionDescriptor<T>.registerExtension(extension: T) {
             _registeredExtensions.getOrPut(this, ::mutableListOf).add(extension)
         }
+
+        /**
+         * Passed [disposable] will be called at the end of compilation
+         */
+        fun registerDisposable(disposable: PluginDisposable) {
+            _disposables += disposable
+        }
+    }
+
+    fun interface PluginDisposable {
+        fun dispose()
     }
 
     abstract val supportsK2: Boolean
@@ -47,6 +63,9 @@ fun CompilerPluginRegistrar.ExtensionStorage.registerInProject(
             }
         }
     }
+    for (disposable in disposables) {
+        Disposer.register(project) { disposable.dispose() }
+    }
 }
 
 private fun ProjectExtensionDescriptor<Any>.registerExtensionUnsafe(project: Project, extension: Any) {