[Gradle][MPP] KotlinWithJavaCompilationFactory: Support recursive call from .javaSourceSets.all

^KT-54783 Verification Pending

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

@@ -173,12 +173,8 @@ internal abstract class AbstractKotlinPlugin(
                 else -> KOTLIN_DSL_NAME
             }
 
-            // Workaround for indirect mutual recursion between the two `all { ... }` handlers:
-            val compilationsUnderConstruction = mutableMapOf<String, KotlinCompilation<*>>()
-
             javaSourceSets.all { javaSourceSet ->
-                val kotlinCompilation =
+                val kotlinCompilation = kotlinTarget.compilations.maybeCreate(javaSourceSet.name)
-                    compilationsUnderConstruction[javaSourceSet.name] ?: kotlinTarget.compilations.maybeCreate(javaSourceSet.name)
 
                 if (duplicateJavaSourceSetsAsKotlinSourceSets) {
                     val kotlinSourceSet = project.kotlinExtension.sourceSets.maybeCreate(kotlinCompilation.name)
@@ -195,12 +191,7 @@ internal abstract class AbstractKotlinPlugin(
             }
 
             kotlinTarget.compilations.all { kotlinCompilation ->
-                val sourceSetName = kotlinCompilation.name
+                kotlinCompilation.source(kotlinCompilation.defaultSourceSet)
-                compilationsUnderConstruction[sourceSetName] = kotlinCompilation
-
-                // Another Kotlin source set following the other convention, named according to the compilation, not the Java source set:
-                val kotlinSourceSet = project.kotlinExtension.sourceSets.maybeCreate(kotlinCompilation.defaultSourceSetName)
-                kotlinCompilation.source(kotlinSourceSet)
             }
 
             kotlinTarget.compilations.run {

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/jvm/KotlinWithJavaCompilationFactory.kt

@@ -26,7 +26,15 @@ class KotlinWithJavaCompilationFactory<KotlinOptionsType : KotlinCommonOptions,
 
     @Suppress("UNCHECKED_CAST")
     override fun create(name: String): KotlinWithJavaCompilation<KotlinOptionsType, CO> {
-        val javaSourceSet = target.javaSourceSets.maybeCreate(name)
+        val javaSourceSet = target.javaSourceSets.findByName(name) ?: run {
+            /*
+            Creating the java SourceSet first here:
+            After the javaSourceSet is created, another .all hook will call into this factory creating the KotlinCompilation.
+            This call will just return this instance instead eagerly
+             */
+            target.javaSourceSets.create(name)
+            return target.compilations.getByName(name)
+        }
 
         val compilationImplFactory = KotlinCompilationImplFactory(
             compilerOptionsFactory = { _, _ ->

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

@@ -0,0 +1,31 @@
+/*
+ * 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.
+ */
+
+@file:Suppress("FunctionName")
+
+package org.jetbrains.kotlin.gradle
+
+import org.gradle.kotlin.dsl.getByType
+import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
+import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinWithJavaCompilation
+import kotlin.test.Test
+import kotlin.test.assertSame
+
+class KT54783JvmWithJavaCompilationFactoryTest {
+    @Test
+    fun `test - create custom jvm compilation`() {
+        val project = buildProjectWithJvm()
+        val kotlin = project.extensions.getByType<KotlinJvmProjectExtension>()
+        var customCompilationInstanceInConfigureBlock: KotlinWithJavaCompilation<*, *>? = null
+        kotlin.target.compilations.create("custom") { customCompilationInstance ->
+            customCompilationInstanceInConfigureBlock = customCompilationInstance
+        }
+
+        assertSame(
+            kotlin.target.compilations.getByName("custom"), customCompilationInstanceInConfigureBlock,
+            "Expected same compilation instance in configure block as well as present in target.compilations"
+        )
+    }
+}