KT-50736: Fix the implementation of whenEvaluated

Re-implement `whenEvaluated` using a constant number of calls to
`withPlugin` per project.

The approach with `plugins.all { ... }` introduced by the fix
of KT-50509 had an issue with Android plugins applied last in the build:
in those cases that implementation didn't dispatch the actions after
Android and instead fired them earlier in `afterEvaluate`, because
`hasPlugin` doesn't return true in the callback when that exact plugin
is applied.

This change preserves the performance similar to that of the previous
implementation.

Issue #KT-50736

Merge-request: KT-MR-5418

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/AbstractKotlinAndroidGradleTests.kt

@@ -779,6 +779,60 @@ abstract class KotlinAndroid3GradleIT : AbstractKotlinAndroidGradleTests() {
             assertTasksExecuted(*kotlinTaskNames.toTypedArray())
         }
     }
+
+    @Test
+    fun testAfterEvaluateOrdering() = with(Project("AndroidProject")) {
+        setupWorkingDir()
+
+        gradleBuildScript("Lib").writeText(
+            """
+            buildscript {
+                repositories {
+                    mavenLocal()
+                    google()
+                    gradlePluginPortal()
+                }
+                dependencies {
+                    classpath "com.android.tools.build:gradle:${'$'}android_tools_version"
+                    classpath "org.jetbrains.kotlin:kotlin-gradle-plugin:${'$'}kotlin_version"
+                }
+            }
+            
+            plugins {
+                id 'org.jetbrains.kotlin.multiplatform'
+            }
+            
+            class MyAction implements kotlin.jvm.functions.Function1<Project, Void> {
+                Void invoke(Project p) {
+                    println("compilations: " + p.kotlin.targets.getByName("android").compilations.names)
+                }
+            }
+            
+            org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginKt.whenEvaluated(project, new MyAction())
+            
+            apply plugin: "android-library"
+            
+            android {
+                compileSdkVersion 22
+            }
+            
+            kotlin { android("android") { } }
+        """.trimIndent())
+
+        build("help") {
+            assertSuccessful()
+            val reportedCompilations = output.lines()
+                .single { it.contains("compilations: ") }
+                .substringAfter("compilations: ")
+                .removeSurrounding("[", "]")
+                .split(", ")
+                .toSet()
+            assertEquals(
+                setOf("debug", "debugAndroidTest", "debugUnitTest", "release", "releaseUnitTest"),
+                reportedCompilations
+            )
+        }
+    }
 }
 
 abstract class AbstractKotlinAndroidGradleTests : BaseGradleIT() {

libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/WhenEvaluatedAndroidOrderingTest.kt

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+
+/* Associate compilations are not yet supported by the IDE. KT-34102 */
+@file:Suppress("invisible_reference", "invisible_member", "FunctionName", "DuplicatedCode")
+
+import com.android.build.gradle.LibraryExtension
+import org.gradle.api.Project
+import org.gradle.api.internal.project.ProjectInternal
+import org.gradle.api.plugins.HelpTasksPlugin
+import org.gradle.buildinit.plugins.BuildInitPlugin
+import org.gradle.buildinit.plugins.WrapperPlugin
+import org.gradle.testfixtures.ProjectBuilder
+import org.jetbrains.kotlin.gradle.applyMultiplatformPlugin
+import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
+import org.jetbrains.kotlin.gradle.plugin.whenEvaluated
+import org.junit.Test
+import kotlin.test.*
+
+class WhenEvaluatedAndroidOrderingTest {
+
+    private lateinit var project: ProjectInternal
+
+    @BeforeTest
+    fun setup() {
+        project = ProjectBuilder.builder().build() as ProjectInternal
+    }
+
+    /**
+     * Check that the `whenEvaluated` actions that are scheduled before the Android plugin is applied get triggered only after the actions
+     * done in the Android plugin's afterEvaluate phase
+     */
+    @Test
+    fun `test Android compilations visible in whenEvaluated`() {
+        project.applyGradleBuiltInPlugins()
+
+        val kotlin = project.applyMultiplatformPlugin()
+
+        var triggered = false
+
+        project.whenEvaluated {
+            /** These are created by the Kotlin plugin immediately on Android plugin's afterEvaluate actions */
+            val androidCompilations = kotlin.targets.getByName("android").compilations
+            assertTrue { androidCompilations.isNotEmpty() }
+
+            assertFalse(triggered, "whenEvaluated should call the function only once")
+            triggered = true
+        }
+
+        project.applyAndroidLibraryPlugin()
+        kotlin.android()
+
+        project.evaluate()
+
+        assertTrue { triggered }
+    }
+
+    // Apply these built-in plugins, so that Gradle doesn't apply them in `project.evaluate()`
+    // below and trigger the plugin application callbacks after Android is applied
+    private fun Project.applyGradleBuiltInPlugins() {
+        plugins.apply(org.gradle.api.plugins.HelpTasksPlugin::class.java)
+        plugins.apply(org.gradle.buildinit.plugins.BuildInitPlugin::class.java)
+        plugins.apply(org.gradle.buildinit.plugins.WrapperPlugin::class.java)
+    }
+
+    private fun Project.applyAndroidLibraryPlugin() {
+        project.plugins.apply("android-library")
+        val android = project.extensions.getByName("android") as LibraryExtension
+        android.compileSdkVersion(30)
+    }
+}

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

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.gradle.logging.kotlinWarn
 import org.jetbrains.kotlin.gradle.tasks.AbstractKotlinCompile
 import org.jetbrains.kotlin.gradle.utils.SingleWarningPerBuild
 import org.jetbrains.kotlin.gradle.utils.androidPluginIds
+import org.jetbrains.kotlin.gradle.utils.getOrPut
 import java.util.concurrent.atomic.AtomicBoolean
 
 abstract class KotlinPlatformPluginBase(protected val platformName: String) : Plugin<Project> {
@@ -215,27 +216,37 @@ internal fun <T> Project.whenEvaluated(fn: Project.() -> T) {
         return
     }
 
-    /* Make sure that all afterEvaluate blocks from the AndroidPlugin get scheduled first */
+    /** If there's already an Android plugin applied, just dispatch the action to `afterEvaluate`, it gets executed after AGP's actions */
-    val isDispatched = AtomicBoolean(false)
+    if (androidPluginIds.any { pluginManager.hasPlugin(it) }) {
+        afterEvaluate { fn() }
+        return
+    }
 
-    fun isAndroidPluginApplied(): Boolean =
+    val isDispatchedAfterAndroid = AtomicBoolean(false)
-        androidPluginIds.any { pluginManager.hasPlugin(it) }
 
-    if (isAndroidPluginApplied()) {
+    /**
-        if (!isDispatched.getAndSet(true)) {
+     * This queue holds all actions submitted to `whenEvaluated` in this project, waiting for one of the Android plugins to be applied.
-            afterEvaluate { fn() }
+     * After (and if) an Android plugin gets applied, we dispatch all the actions in the queue to `afterEvaluate`, so that they are
+     * executed after what AGP scheduled to `afterEvaluate`. There are different Android plugins, so actions in the queue also need to check
+     * if it's the first Android plugin, using `isDispatched` (each has its own instance).
+     */
+    val afterAndroidDispatchQueue = project.extensions.extraProperties.getOrPut("org.jetbrains.kotlin.whenEvaluated") {
+        val queue = mutableListOf<() -> Unit>()
+        // Trigger the actions on any plugin applied; the actions themselves ensure that they only dispatch the fn once.
+        androidPluginIds.forEach { id ->
+            pluginManager.withPlugin(id) { queue.forEach { it() } }
         }
-    } else {
+        queue
-        plugins.all {
+    }
-            if (!isDispatched.get() && isAndroidPluginApplied() && !isDispatched.getAndSet(true)) {
+    afterAndroidDispatchQueue.add {
-                afterEvaluate { fn() }
+        if (!isDispatchedAfterAndroid.getAndSet(true)) {
-            }
+            afterEvaluate { fn() }
         }
     }
 
     afterEvaluate {
-        /* If no Android plugin was loaded, then the action was not dispatched and we can freely execute it now */
+        /** If no Android plugin was loaded, then the action was not dispatched, and we can freely execute it now */
-        if (!isDispatched.getAndSet(true)) {
+        if (!isDispatchedAfterAndroid.getAndSet(true)) {
             fn()
         }
     }