Pill: Import all defined source sets, not only "src" and "test"

buildSrc/src/main/kotlin/pill/parser.kt

@@ -3,6 +3,7 @@ package org.jetbrains.kotlin.pill
 
 import org.gradle.api.Project
 import org.gradle.api.artifacts.*
+import org.gradle.api.tasks.*
 import org.gradle.api.plugins.JavaPlugin
 import org.gradle.api.plugins.JavaPluginConvention
 import org.gradle.kotlin.dsl.configure
@@ -112,11 +113,6 @@ private val TEST_CONFIGURATION_MAPPING = mapOf(
     listOf("jpsTest") to Scope.TEST
 )
 
-private val SOURCE_SET_MAPPING = mapOf(
-    "main" to Kind.PRODUCTION,
-    "test" to Kind.TEST
-)
-
 private fun ParserContext.parseModules(project: Project, excludedProjects: List<Project>): List<PModule> {
     val (productionContentRoots, testContentRoots) = parseContentRoots(project).partition { !it.forTests }
 
@@ -208,9 +204,9 @@ private fun parseSourceRoots(project: Project): List<PSourceRoot> {
 
     val sourceRoots = mutableListOf<PSourceRoot>()
 
-    project.configure<JavaPluginConvention> {
+    for (sourceSet in project.sourceSets) {
-        for ((sourceSetName, kind) in SOURCE_SET_MAPPING) {
+        val sourceSetName = sourceSet.name
-            val sourceSet = sourceSets.findByName(sourceSetName) ?: continue
+        val kind = if (sourceSetName == SourceSet.TEST_SOURCE_SET_NAME) Kind.TEST else Kind.PRODUCTION
 
         fun Any.getKotlin(): SourceDirectorySet {
             val kotlinMethod = javaClass.getMethod("getKotlin")
@@ -250,7 +246,6 @@ private fun parseSourceRoots(project: Project): List<PSourceRoot> {
             sourceRoots += PSourceRoot(directory, kind)
         }
     }
-    }
 
     return sourceRoots
 }
@@ -421,3 +416,10 @@ fun List<CollectedConfiguration>.collectDependencies(): List<DependencyInfo> {
 
     return dependencies
 }
+
+private val Project.sourceSets: SourceSetContainer
+    get() {
+        lateinit var result: SourceSetContainer
+        project.configure<JavaPluginConvention> { result = sourceSets }
+        return result
+    }