[Gradle] Remove kotlinx.coroutines from KGP compile classpath

- Replace CompletableDeferred with a simple 'Completable'
- Replace withContext usage for 'withRestrictedStages' with custom
implementation

^KT-58162 Verification Pending

libraries/tools/kotlin-gradle-plugin/build.gradle.kts

@@ -60,6 +60,7 @@ dependencies {
     }
     commonCompileOnly(project(":kotlin-tooling-metadata"))
 
+
     commonImplementation(project(":kotlin-gradle-plugin-idea"))
     commonImplementation(project(":kotlin-gradle-plugin-idea-proto"))
     commonImplementation(project(":kotlin-util-klib"))
@@ -110,6 +111,8 @@ dependencies {
     testImplementation(project(":kotlin-tooling-metadata"))
 }
 
+configurations.commonCompileClasspath.get().exclude("org.jetbrains.kotlinx", "kotlinx-coroutines-core")
+
 if (kotlinBuildProperties.isInJpsBuildIdeaSync) {
     configurations.commonApi.get().exclude("com.android.tools.external.com-intellij", "intellij-core")
 }

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt

@@ -5,8 +5,6 @@
 
 package org.jetbrains.kotlin.gradle.plugin
 
-import kotlinx.coroutines.currentCoroutineContext
-import kotlinx.coroutines.withContext
 import org.gradle.api.Project
 import org.gradle.api.provider.Property
 import org.jetbrains.kotlin.gradle.plugin.KotlinPluginLifecycle.*
@@ -127,7 +125,7 @@ internal fun Project.startKotlinPluginLifecycle() {
  * the currently running coroutine. Throws if this coroutine was not started using a [KotlinPluginLifecycle]
  */
 internal suspend fun currentKotlinPluginLifecycle(): KotlinPluginLifecycle {
-    return currentCoroutineContext()[KotlinPluginLifecycleCoroutineContextElement]?.lifecycle
+    return coroutineContext[KotlinPluginLifecycleCoroutineContextElement]?.lifecycle
         ?: error("Missing $KotlinPluginLifecycleCoroutineContextElement in currentCoroutineContext")
 }
 
@@ -234,8 +232,17 @@ internal suspend fun <T> requireCurrentStage(block: suspend () -> T): T {
  * ```
  */
 internal suspend fun <T> withRestrictedStages(allowed: Set<Stage>, block: suspend () -> T): T {
-    return withContext(RestrictedLifecycleStages(currentKotlinPluginLifecycle(), allowed)) {
+    val newCoroutineContext = coroutineContext + RestrictedLifecycleStages(currentKotlinPluginLifecycle(), allowed)
-        block()
+    return suspendCoroutine { continuation ->
+        val newContinuation = object : Continuation<T> {
+            override val context: CoroutineContext
+                get() = newCoroutineContext
+
+            override fun resumeWith(result: Result<T>) {
+                continuation.resumeWith(result)
+            }
+        }
+        block.startCoroutine(newContinuation)
     }
 }
 

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/utils/Future.kt

@@ -5,9 +5,6 @@
 
 package org.jetbrains.kotlin.gradle.utils
 
-import kotlinx.coroutines.CompletableDeferred
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.completeWith
 import org.gradle.api.Project
 import org.jetbrains.kotlin.gradle.plugin.HasProject
 import org.jetbrains.kotlin.gradle.plugin.KotlinPluginLifecycle
@@ -17,6 +14,9 @@ import org.jetbrains.kotlin.tooling.core.ExtrasLazyProperty
 import org.jetbrains.kotlin.tooling.core.HasMutableExtras
 import org.jetbrains.kotlin.tooling.core.extrasLazyProperty
 import java.io.Serializable
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.resume
+import kotlin.coroutines.suspendCoroutine
 
 /**
  * See [KotlinPluginLifecycle]:
@@ -95,9 +95,8 @@ internal fun <T> CompletableFuture(): CompletableFuture<T> {
     return FutureImpl()
 }
 
-@OptIn(ExperimentalCoroutinesApi::class)
 private class FutureImpl<T>(
-    private val deferred: CompletableDeferred<T> = CompletableDeferred(),
+    private val deferred: Completable<T> = Completable(),
     private val lifecycle: KotlinPluginLifecycle? = null
 ) : CompletableFuture<T>, Serializable {
     fun completeWith(result: Result<T>) = deferred.completeWith(result)
@@ -123,7 +122,7 @@ private class FutureImpl<T>(
 
     private class Surrogate<T>(private val value: T) : Serializable {
         private fun readResolve(): Any {
-            return FutureImpl(CompletableDeferred(value))
+            return FutureImpl(Completable(value))
         }
     }
 }
@@ -153,7 +152,7 @@ private class LenientFutureImpl<T>(
 
     private class Surrogate<T>(private val value: T) : Serializable {
         private fun readResolve(): Any {
-            return LenientFutureImpl(FutureImpl(CompletableDeferred(value)))
+            return LenientFutureImpl(FutureImpl(Completable(value)))
         }
     }
 }
@@ -173,7 +172,63 @@ private class LazyFutureImpl<T>(private val future: Lazy<Future<T>>) : Future<T>
 
     private class Surrogate<T>(private val value: T) : Serializable {
         private fun readResolve(): Any {
-            return FutureImpl(CompletableDeferred(value))
+            return FutureImpl(Completable(value))
         }
     }
+}
+
+/**
+ * Simple, Single Threaded, replacement for kotlinx.coroutines.CompletableDeferred.
+ */
+private class Completable<T>(
+    private var value: Result<T>? = null
+) {
+    constructor(value: T) : this(Result.success(value))
+
+    val isCompleted: Boolean get() = value != null
+
+    private val waitingContinuations = mutableListOf<Continuation<Result<T>>>()
+
+    fun completeWith(result: Result<T>) {
+        check(value == null) { "Already completed with $value" }
+        value = result
+
+        /* Capture and clear current waiting continuations */
+        val continuations = waitingContinuations.toList()
+        waitingContinuations.clear()
+
+        continuations.forEach { continuation ->
+            continuation.resume(result)
+        }
+
+        /*
+        Safety check:
+        We do not expect any coroutines waiting:
+        Any continuation that, during its above .resume, calls into '.await()' shall
+        directly resume and receive the value currently set.
+
+        If the waiting continuations are not empty, then those would be leaking.
+         */
+        assert(waitingContinuations.isEmpty())
+    }
+
+    fun complete(value: T) {
+        completeWith(Result.success(value))
+    }
+
+    fun getCompleted(): T {
+        val value = this.value ?: throw IllegalStateException("Not completed yet")
+        return value.getOrThrow()
+    }
+
+    suspend fun await(): T {
+        val value = this.value
+        if (value != null) {
+            return value.getOrThrow()
+        }
+
+        return suspendCoroutine<Result<T>> { continuation ->
+            waitingContinuations.add(continuation)
+        }.getOrThrow()
+    }
 }