Also register shaders for kotlin source sets

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

@@ -11,6 +11,7 @@ import org.gradle.api.internal.project.ProjectInternal
 import org.gradle.testfixtures.ProjectBuilder
 import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
 import org.jetbrains.kotlin.gradle.dsl.multiplatformExtension
+import org.jetbrains.kotlin.gradle.plugin.mpp.kotlinSourceSet
 import kotlin.test.*
 
 class SyncKotlinAndAndroidSourceSetsTest {
@@ -136,6 +137,20 @@ class SyncKotlinAndAndroidSourceSetsTest {
             "Expected custom configured source directories being present on android source set after evaluation"
         )
     }
+
+    @Test
+    fun `AndroidSourceSet#kotlinSourceSet extension`() {
+        kotlin.android()
+
+        val main = android.sourceSets.getByName("main")
+        assertSame(kotlin.sourceSets.getByName("androidMain"), main.kotlinSourceSet)
+
+        val test = android.sourceSets.getByName("test")
+        assertSame(kotlin.sourceSets.getByName("androidTest"), test.kotlinSourceSet)
+
+        val androidTest = android.sourceSets.getByName("androidTest")
+        assertSame(kotlin.sourceSets.getByName("androidAndroidTest"), androidTest.kotlinSourceSet)
+    }
 }
 
 private fun <T> Set<T>.allPairs(): Sequence<Pair<T, T>> {

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

@@ -22,10 +22,11 @@ internal fun syncKotlinAndAndroidSourceSets(target: KotlinAndroidTarget) {
     android.sourceSets.all { androidSourceSet ->
         val kotlinSourceSetName = kotlinSourceSetNameForAndroidSourceSet(target, androidSourceSet.name)
         val kotlinSourceSet = project.kotlinExtension.sourceSets.maybeCreate(kotlinSourceSetName)
+        androidSourceSet.kotlinSourceSet = kotlinSourceSet
+
         createDefaultDependsOnEdges(target, kotlinSourceSet, androidSourceSet)
         syncKotlinAndAndroidSourceDirs(target, kotlinSourceSet, androidSourceSet)
         syncKotlinAndAndroidResources(target, kotlinSourceSet, androidSourceSet)
-        androidSourceSet.addConvention(KOTLIN_DSL_NAME, kotlinSourceSet)
 
         ifKaptEnabled(project) {
             Kapt3GradleSubplugin.createAptConfigurationIfNeeded(project, androidSourceSet.name)
@@ -33,6 +34,12 @@ internal fun syncKotlinAndAndroidSourceSets(target: KotlinAndroidTarget) {
     }
 }
 
+internal var AndroidSourceSet.kotlinSourceSet: KotlinSourceSet
+    get() = getConvention(KOTLIN_DSL_NAME) as KotlinSourceSet
+    private set(value) {
+        addConvention(KOTLIN_DSL_NAME, value)
+    }
+
 private fun createDefaultDependsOnEdges(
     target: KotlinAndroidTarget,
     kotlinSourceSet: KotlinSourceSet,
@@ -85,6 +92,14 @@ private fun syncKotlinAndAndroidResources(
     kotlinSourceSet: KotlinSourceSet,
     androidSourceSet: AndroidSourceSet
 ) {
+    /*
+    Non-MPP projects will not have a different name for kotlin source sets vs Android source sets
+    Registering resource folders is unnecessary therefore.
+     */
+    if (kotlinSourceSet.name == androidSourceSet.name) {
+        return
+    }
+
     val project = target.project
 
     androidSourceSet.resources.srcDirs(*kotlinSourceSet.resources.toList().toTypedArray())
@@ -116,6 +131,10 @@ private fun syncKotlinAndAndroidResources(
     if (androidSourceSet.jniLibs.srcDirs.isNotEmpty()) {
         androidSourceSet.jniLibs.srcDir(kotlinSourceSet.sourceFolderFor(project, "jniLibs"))
     }
+
+    if (androidSourceSet.shaders.srcDirs.isNotEmpty()) {
+        androidSourceSet.shaders.srcDir(kotlinSourceSet.sourceFolderFor(project, "shaders"))
+    }
 }
 
 private fun KotlinSourceSet.sourceFolderFor(project: Project, type: String): File {