[Gradle] Implement lifecycle stage restrictions APIs

KT-34662

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

@@ -6,6 +6,7 @@
 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.KotlinMultiplatformPluginLifecycle.*
@@ -14,11 +15,7 @@ import java.lang.ref.WeakReference
 import java.util.*
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.collections.ArrayDeque
-import kotlin.collections.Map
-import kotlin.collections.associateWith
-import kotlin.collections.getValue
 import kotlin.collections.set
-import kotlin.collections.toList
 import kotlin.coroutines.*
 import kotlin.reflect.KProperty
 
@@ -30,13 +27,21 @@ internal fun Project.launch(block: suspend KotlinMultiplatformPluginLifecycle.()
     kotlinMultiplatformPluginLifecycle.launch(block)
 }
 
-internal fun Project.launch(stage: Stage, block: suspend KotlinMultiplatformPluginLifecycle.() -> Unit) {
+internal fun Project.launchInStage(stage: Stage, block: suspend KotlinMultiplatformPluginLifecycle.() -> Unit) {
     launch {
         await(stage)
         block()
     }
 }
 
+internal fun Project.launchInRequiredStage(stage: Stage, block: suspend KotlinMultiplatformPluginLifecycle.() -> Unit) {
+    launchInStage(stage) {
+        requiredStage(stage) {
+            block()
+        }
+    }
+}
+
 internal val Project.kotlinMultiplatformPluginLifecycle: KotlinMultiplatformPluginLifecycle
     get() = extraProperties.getOrPut(KotlinMultiplatformPluginLifecycle::class.java.name) {
         KotlinMultiplatformPluginLifecycleImpl(project)
@@ -56,7 +61,7 @@ internal suspend fun await(stage: Stage) {
 }
 
 internal inline fun <reified T : Any> Project.newLifecycleAwareProperty(
-    finaliseIn: Stage, initialValue: T? = null
+    finaliseIn: Stage = Stage.FinaliseDsl, initialValue: T? = null
 ): LifecycleAwareProperty<T> {
     return kotlinMultiplatformPluginLifecycle.newLifecycleAwareProperty(T::class.java, finaliseIn, initialValue)
 }
@@ -74,6 +79,20 @@ internal suspend fun Property<*>.isLifecycleAware(): Boolean {
     return findLifecycleAwareProperty() != null
 }
 
+internal suspend fun <T> requiredStage(stage: Stage, block: suspend () -> T): T {
+    return withRestrictedStages(hashSetOf(stage), block)
+}
+
+internal suspend fun <T> requireCurrentStage(block: suspend () -> T): T {
+    return requiredStage(currentMultiplatformPluginLifecycle().stage, block)
+}
+
+internal suspend fun <T> withRestrictedStages(allowed: Set<Stage>, block: suspend () -> T): T {
+    return withContext(RestrictedLifecycleStages(currentMultiplatformPluginLifecycle(), allowed)) {
+        block()
+    }
+}
+
 /*
 Definition of the Lifecycle and its stages
  */
@@ -83,6 +102,10 @@ internal interface KotlinMultiplatformPluginLifecycle {
         Configure,
         AfterEvaluate,
 
+        BeforeFinaliseDsl,
+        FinaliseDsl,
+        AfterFinaliseDsl,
+
         BeforeFinaliseRefinesEdges,
         FinaliseRefinesEdges,
         AfterFinaliseRefinesEdges,
@@ -91,14 +114,16 @@ internal interface KotlinMultiplatformPluginLifecycle {
         FinaliseCompilations,
         AfterFinaliseCompilations,
 
-        Finalised
-    }
+        Finalised;
 
-    interface LifecycleAwareProperty<T : Any> {
-        val finaliseIn: Stage
-        val property: Property<T>
-        suspend fun awaitFinalValue(): T?
-        operator fun getValue(thisRef: Any?, property: KProperty<*>): Property<T> = this.property
+        operator fun rangeTo(other: Stage): Set<Stage> {
+            return values.subList(this.ordinal, other.ordinal + 1).toSet()
+        }
+
+        companion object {
+            val values = values().toList()
+            fun upTo(stage: Stage): Set<Stage> = values.first()..stage
+        }
     }
 
     val project: Project
@@ -116,6 +141,13 @@ internal interface KotlinMultiplatformPluginLifecycle {
     ): LifecycleAwareProperty<T>
 
     class IllegalLifecycleException(message: String) : IllegalStateException(message)
+
+    interface LifecycleAwareProperty<T : Any> {
+        val finaliseIn: Stage
+        val property: Property<T>
+        suspend fun awaitFinalValue(): T?
+        operator fun getValue(thisRef: Any?, property: KProperty<*>): Property<T> = this.property
+    }
 }
 
 
@@ -124,7 +156,7 @@ Implementation
  */
 
 private class KotlinMultiplatformPluginLifecycleImpl(override val project: Project) : KotlinMultiplatformPluginLifecycle {
-    private val enqueuedStages: ArrayDeque<Stage> = ArrayDeque(Stage.values().toList())
+    private val enqueuedStages: ArrayDeque<Stage> = ArrayDeque(Stage.values)
     private val enqueuedActions: Map<Stage, ArrayDeque<KotlinMultiplatformPluginLifecycle.() -> Unit>> =
         Stage.values().associateWith { ArrayDeque() }
 
@@ -197,8 +229,8 @@ private class KotlinMultiplatformPluginLifecycleImpl(override val project: Proje
         val lifecycle = this
         check(isStarted.get()) { "Cannot launch when ${KotlinMultiplatformPluginLifecycle::class.simpleName} is not started" }
         check(!isFinished.get()) { "Cannot launch when ${KotlinMultiplatformPluginLifecycle::class.simpleName} is already finished" }
-        val coroutine = block.createCoroutine(this, object : Continuation<Unit> {
 
+        val coroutine = block.createCoroutine(this, object : Continuation<Unit> {
             override val context: CoroutineContext = EmptyCoroutineContext +
                     KotlinMultiplatformPluginLifecycleCoroutineContextElement(lifecycle)
 
@@ -251,3 +283,36 @@ private class KotlinMultiplatformPluginLifecycleCoroutineContextElement(
 
     override val key: CoroutineContext.Key<KotlinMultiplatformPluginLifecycleCoroutineContextElement> = Key
 }
+
+private class RestrictedLifecycleStages(
+    private val lifecycle: KotlinMultiplatformPluginLifecycle,
+    private val allowedStages: Set<Stage>,
+) : CoroutineContext.Element, ContinuationInterceptor {
+    @OptIn(ExperimentalStdlibApi::class)
+    companion object Key : AbstractCoroutineContextKey<ContinuationInterceptor, RestrictedLifecycleStages>(
+        ContinuationInterceptor, { it as? RestrictedLifecycleStages }
+    )
+
+    override val key: CoroutineContext.Key<*> = Key
+
+    override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> = object : Continuation<T> {
+        override val context: CoroutineContext
+            get() = continuation.context
+
+        override fun resumeWith(result: Result<T>) = when {
+            result.isFailure -> continuation.resumeWith(result)
+            lifecycle.stage !in allowedStages -> continuation.resumeWithException(
+                IllegalLifecycleException(
+                    "Required stage in '$allowedStages', but lifecycle switched to '${lifecycle.stage}'"
+                )
+            )
+            else -> continuation.resumeWith(result)
+        }
+    }
+
+    init {
+        if (lifecycle.stage !in allowedStages) {
+            throw IllegalLifecycleException("Required stage in '${allowedStages}' but lifecycle is currently in '${lifecycle.stage}'")
+        }
+    }
+}

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

@@ -172,9 +172,9 @@ class KotlinMultiplatformPlugin : Plugin<Project> {
         val production = sourceSets.create(KotlinSourceSet.COMMON_MAIN_SOURCE_SET_NAME)
         val test = sourceSets.create(KotlinSourceSet.COMMON_TEST_SOURCE_SET_NAME)
         targets.all { target ->
-            project.launch(KotlinMultiplatformPluginLifecycle.Stage.FinaliseRefinesEdges) {
+            project.launchInStage(KotlinMultiplatformPluginLifecycle.Stage.FinaliseRefinesEdges) {
                 /* Only setup default refines edges when no KotlinTargetHierarchy was applied */
-                if (project.multiplatformExtension.internalKotlinTargetHierarchy.appliedDescriptors.isNotEmpty()) return@launch
+                if (project.multiplatformExtension.internalKotlinTargetHierarchy.appliedDescriptors.isNotEmpty()) return@launchInStage
 
                 target.compilations.findByName(KotlinCompilation.MAIN_COMPILATION_NAME)?.let { mainCompilation ->
                     mainCompilation.defaultSourceSet.takeIf { it != production }?.dependsOn(production)

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/sources/android/configurator/MultiplatformLayoutV2DependsOnConfigurator.kt

@@ -11,10 +11,10 @@ import org.jetbrains.kotlin.gradle.dsl.kotlinExtension
 import org.jetbrains.kotlin.gradle.dsl.multiplatformExtensionOrNull
 import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle
 import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
-import org.jetbrains.kotlin.gradle.plugin.kotlinMultiplatformPluginLifecycle
-import org.jetbrains.kotlin.gradle.plugin.launch
+import org.jetbrains.kotlin.gradle.plugin.KotlinTargetHierarchy
+import org.jetbrains.kotlin.gradle.plugin.launchInStage
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinAndroidTarget
-import org.jetbrains.kotlin.gradle.plugin.mpp.targetHierarchy.KotlinTargetHierarchy
+import org.jetbrains.kotlin.gradle.plugin.mpp.targetHierarchy.orNull
 import org.jetbrains.kotlin.gradle.plugin.sources.android.AndroidBaseSourceSetName
 import org.jetbrains.kotlin.gradle.plugin.sources.android.AndroidVariantType
 import org.jetbrains.kotlin.gradle.plugin.sources.android.type
@@ -32,15 +32,15 @@ internal object MultiplatformLayoutV2DependsOnConfigurator : KotlinAndroidSource
     }
 
     private fun setDefaultDependsOn(target: KotlinAndroidTarget, kotlinSourceSet: KotlinSourceSet, variantType: AndroidVariantType) {
-        target.project.launch(KotlinMultiplatformPluginLifecycle.Stage.FinaliseRefinesEdges) {
+        target.project.launchInStage(KotlinMultiplatformPluginLifecycle.Stage.FinaliseRefinesEdges) {
             /* Only setup default if not KotlinTargetHierarchy was applied */
             if (target.project.multiplatformExtensionOrNull?.internalKotlinTargetHierarchy?.appliedDescriptors.orEmpty().isNotEmpty()) {
-                return@launch
+                return@launchInStage
             }
 
-            val module = KotlinTargetHierarchy.Module.orNull(target, variantType) ?: return@launch
+            val module = KotlinTargetHierarchy.ModuleName.orNull(target, variantType) ?: return@launchInStage
             val commonSourceSetName = lowerCamelCaseName("common", module.name)
-            val commonSourceSet = target.project.kotlinExtension.sourceSets.findByName(commonSourceSetName) ?: return@launch
+            val commonSourceSet = target.project.kotlinExtension.sourceSets.findByName(commonSourceSetName) ?: return@launchInStage
             kotlinSourceSet.dependsOn(commonSourceSet)
         }
     }

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/metadata/KotlinMetadataTargetConfigurator.kt

@@ -175,7 +175,7 @@ class KotlinMetadataTargetConfigurator :
     private fun createMetadataCompilationsForCommonSourceSets(
         target: KotlinMetadataTarget,
         allMetadataJar: TaskProvider<out Jar>
-    ) = target.project.launch(KotlinMultiplatformPluginLifecycle.Stage.FinaliseCompilations) {
+    ) = target.project.launchInStage(KotlinMultiplatformPluginLifecycle.Stage.FinaliseCompilations) {
         // Do this after all targets are configured by the user build script
 
         val publishedCommonSourceSets: Set<KotlinSourceSet> = getCommonSourceSetsForMetadataCompilation(project)

libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinMultiplatformPluginLifecycleTest.kt

@@ -8,11 +8,12 @@
 package org.jetbrains.kotlin.gradle.unitTests
 
 import org.gradle.api.ProjectConfigurationException
-import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle
+import org.jetbrains.kotlin.gradle.plugin.*
+import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle.IllegalLifecycleException
+import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle.Stage
 import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle.Stage.*
-import org.jetbrains.kotlin.gradle.plugin.kotlinMultiplatformPluginLifecycle
 import org.jetbrains.kotlin.gradle.util.buildProjectWithMPP
-import org.jetbrains.kotlin.tooling.core.linearClosure
+import org.jetbrains.kotlin.gradle.util.runLifecycleAwareTest
 import org.jetbrains.kotlin.tooling.core.withLinearClosure
 import org.junit.Test
 import java.util.concurrent.atomic.AtomicBoolean
@@ -67,10 +68,10 @@ class KotlinMultiplatformPluginLifecycleTest {
 
     @Test
     fun `test - all stages are executed in order`() {
-        val invocations = KotlinMultiplatformPluginLifecycle.Stage.values().associateWith { AtomicInteger(0) }
-        KotlinMultiplatformPluginLifecycle.Stage.values().toList().forEach { stage ->
+        val invocations = Stage.values().associateWith { AtomicInteger(0) }
+        Stage.values().toList().forEach { stage ->
             lifecycle.enqueue(stage) {
-                KotlinMultiplatformPluginLifecycle.Stage.values().forEach { otherStage ->
+                Stage.values().forEach { otherStage ->
                     when {
                         otherStage.ordinal < stage.ordinal -> assertEquals(1, invocations.getValue(otherStage).get())
                         otherStage.ordinal == stage.ordinal -> assertEquals(1, invocations.getValue(stage).incrementAndGet())
@@ -165,7 +166,7 @@ class KotlinMultiplatformPluginLifecycleTest {
     fun `test - enqueue of already executed stage - throws exception`() {
         val executed = AtomicBoolean(false)
         lifecycle.enqueue(Finalised) {
-            assertFailsWith<KotlinMultiplatformPluginLifecycle.IllegalLifecycleException> {
+            assertFailsWith<IllegalLifecycleException> {
                 lifecycle.enqueue(AfterEvaluate) { fail("This code shall not be executed!") }
             }
             assertFalse(executed.getAndSet(true))
@@ -177,7 +178,7 @@ class KotlinMultiplatformPluginLifecycleTest {
 
     @Test
     fun `test - stage property is correct`() {
-        KotlinMultiplatformPluginLifecycle.Stage.values().forEach { stage ->
+        Stage.values().forEach { stage ->
             lifecycle.enqueue(stage) {
                 assertEquals(lifecycle.stage, stage)
             }
@@ -304,4 +305,48 @@ class KotlinMultiplatformPluginLifecycleTest {
 
         assertTrue(testException in causes)
     }
+
+    @Test
+    fun `test - require current stage`() = project.runLifecycleAwareTest {
+        launchInStage(AfterEvaluate) {
+            requireCurrentStage { } // OK
+
+            requireCurrentStage {
+                /* Fails because of stage transition  using 'await' */
+                assertFailsWith<IllegalLifecycleException> { await(Finalised) }
+            }
+        }
+    }
+
+    @Test
+    fun `test - launch in required stage`() = project.runLifecycleAwareTest {
+        launchInRequiredStage(AfterEvaluate) {
+            assertEquals(AfterEvaluate, stage)
+            await(AfterEvaluate)
+            assertEquals(AfterEvaluate, stage)
+            assertFailsWith<IllegalLifecycleException> { await(Finalised) }
+        }
+    }
+
+    @Test
+    fun `test - withRestrictedStages`() = project.runLifecycleAwareTest {
+        launch {
+            withRestrictedStages(Stage.upTo(FinaliseRefinesEdges)) {
+                assertEquals(Configure, stage)
+
+                await(AfterEvaluate)
+                assertEquals(AfterEvaluate, stage)
+
+                await(BeforeFinaliseRefinesEdges)
+                assertEquals(BeforeFinaliseRefinesEdges, stage)
+
+                await(FinaliseRefinesEdges)
+                assertEquals(FinaliseRefinesEdges, stage)
+
+                assertFailsWith<IllegalLifecycleException> {
+                    await(AfterFinaliseRefinesEdges)
+                }
+            }
+        }
+    }
 }

libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/LifecycleAwarePropertyTest.kt

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.gradle.plugin.*
 import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle.Stage.*
 import org.jetbrains.kotlin.gradle.plugin.awaitFinalValue
 import org.jetbrains.kotlin.gradle.plugin.currentMultiplatformPluginLifecycle
-import org.jetbrains.kotlin.gradle.plugin.launch
+import org.jetbrains.kotlin.gradle.plugin.launchInStage
 import org.jetbrains.kotlin.gradle.plugin.newLifecycleAwareProperty
 import org.jetbrains.kotlin.gradle.util.buildProjectWithMPP
 import org.jetbrains.kotlin.gradle.util.runLifecycleAwareTest
@@ -30,11 +30,11 @@ class LifecycleAwarePropertyTest {
         val property by project.newLifecycleAwareProperty<Int>(AfterFinaliseRefinesEdges)
         assertTrue(property.isLifecycleAware())
 
-        launch(BeforeFinaliseRefinesEdges) {
+        launchInStage(BeforeFinaliseRefinesEdges) {
             property.set(1)
         }
 
-        launch(FinaliseRefinesEdges) {
+        launchInStage(FinaliseRefinesEdges) {
             assertEquals(1, property.get())
             property.set(2)
         }
@@ -56,7 +56,7 @@ class LifecycleAwarePropertyTest {
         val property by project.newLifecycleAwareProperty<Int>(AfterEvaluate)
         property.set(1)
 
-        launch(AfterEvaluate) {
+        launchInStage(AfterEvaluate) {
             assertFailsWith<IllegalStateException> { property.set(2) }
         }
     }

libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/sources/InternalKotlinSourceSetTest.kt

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.gradle.dsl.multiplatformExtension
 import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
 import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle
 import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
-import org.jetbrains.kotlin.gradle.plugin.launch
+import org.jetbrains.kotlin.gradle.plugin.launchInStage
 import org.jetbrains.kotlin.gradle.plugin.mpp.getHostSpecificMainSharedSourceSets
 import org.jetbrains.kotlin.gradle.plugin.sources.InternalKotlinSourceSet
 import org.jetbrains.kotlin.gradle.plugin.sources.internal
@@ -70,7 +70,7 @@ class InternalKotlinSourceSetTest {
             nativeMain.internal.compilations
         )
 
-        project.launch(KotlinMultiplatformPluginLifecycle.Stage.Finalised) {
+        project.launchInStage(KotlinMultiplatformPluginLifecycle.Stage.Finalised) {
             assertEquals<Set<KotlinCompilation<*>>>(
                 setOf(
                     metadataCompilation,