From d3d6cdcef7b51fd9fe6ddf606161069edb2e805a Mon Sep 17 00:00:00 2001
From: Sebastian Sellmair <sebastian.sellmair@jetbrains.com>
Date: Tue, 20 Jun 2023 10:58:42 +0200
Subject: [PATCH] [Gradle] Implement Future<T>.map(transform)

^KT-59446 In Progress
---
 .../jetbrains/kotlin/gradle/utils/Future.kt   | 36 +++++++++++++++++++
 .../kotlin/gradle/unitTests/FutureTest.kt     | 36 +++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/utils/Future.kt b/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/utils/Future.kt
index 6c1b7d0b225df..aa29237c6583e 100644
--- a/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/utils/Future.kt
+++ b/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/utils/Future.kt
@@ -55,6 +55,10 @@ internal interface CompletableFuture<T> : Future<T> {
     fun complete(value: T)
 }
 
+internal fun <T, R> Future<T>.map(transform: (T) -> R): Future<R> {
+    return MappedFutureImpl(this, transform)
+}
+
 internal fun CompletableFuture<Unit>.complete() = complete(Unit)
 
 /**
@@ -127,6 +131,38 @@ private class FutureImpl<T>(
     }
 }
 
+private class MappedFutureImpl<T, R>(
+    private val future: Future<T>,
+    private var transform: (T) -> R,
+) : Future<R>, Serializable {
+
+    private val value = Completable<R>()
+
+    override suspend fun await(): R {
+        if (value.isCompleted) return value.getCompleted()
+        value.complete(transform(future.await()))
+        transform = { throw IllegalStateException("Unexpected 'transform' in future") }
+        return value.getCompleted()
+    }
+
+    override fun getOrThrow(): R {
+        if (value.isCompleted) return value.getCompleted()
+        value.complete(transform(future.getOrThrow()))
+        transform = { throw IllegalStateException("Unexpected 'transform' in future") }
+        return value.getCompleted()
+    }
+
+    private fun writeReplace(): Any {
+        return Surrogate(getOrThrow())
+    }
+
+    private class Surrogate<T>(private val value: T) : Serializable {
+        private fun readResolve(): Any {
+            return FutureImpl(Completable(value))
+        }
+    }
+}
+
 private class LenientFutureImpl<T>(
     private val future: Future<T>,
 ) : LenientFuture<T>, Serializable {
diff --git a/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/FutureTest.kt b/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/FutureTest.kt
index c4646a5a3bb16..144d4130856ec 100644
--- a/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/FutureTest.kt
+++ b/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/FutureTest.kt
@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.gradle.idea.testFixtures.utils.serialize
 import org.jetbrains.kotlin.gradle.plugin.*
 import org.jetbrains.kotlin.gradle.plugin.KotlinPluginLifecycle.IllegalLifecycleException
 import org.jetbrains.kotlin.gradle.plugin.KotlinPluginLifecycle.Stage.FinaliseDsl
+import org.jetbrains.kotlin.gradle.plugin.KotlinPluginLifecycle.Stage.ReadyForExecution
 import org.jetbrains.kotlin.gradle.util.buildProject
 import org.jetbrains.kotlin.gradle.util.runLifecycleAwareTest
 import org.jetbrains.kotlin.gradle.utils.*
@@ -114,4 +115,39 @@ class FutureTest {
             assertEquals(1, futureInvocations.get())
         }
     }
+
+    @Test
+    fun `test - future map`() = project.runLifecycleAwareTest {
+        val future = project.future { "1" }
+        val transformInvocations = AtomicInteger(0)
+
+        val mappedFuture = future.map { value ->
+            transformInvocations.incrementAndGet()
+            value.toInt()
+        }
+
+        assertEquals(1, mappedFuture.await())
+        assertEquals(1, mappedFuture.getOrThrow())
+        assertEquals(1, transformInvocations.get())
+
+        val deserializedMappedFuture = mappedFuture.serialize().deserialize() as Future<*>
+        assertEquals(1, deserializedMappedFuture.await())
+        assertEquals(1, deserializedMappedFuture.getOrThrow())
+    }
+
+    @Test
+    fun `test - future map - getOrThrow - exception`() = project.runLifecycleAwareTest {
+        val future = project.future {
+            ReadyForExecution.await()
+            "1"
+        }
+
+        val mappedFuture = future.map {
+            it.toInt()
+        }
+
+        assertFailsWith<IllegalLifecycleException> { mappedFuture.getOrThrow() }
+        future.await()
+        assertEquals(1, mappedFuture.getOrThrow())
+    }
 }