From 614033d64c1e489048c5cbd21677be5cff0db185 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Mon, 23 Apr 2018 23:56:05 +0300
Subject: [PATCH] Pill: Import all defined source sets, not only "src" and
 "test"

---
 buildSrc/src/main/kotlin/pill/parser.kt | 86 +++++++++++++------------
 1 file changed, 44 insertions(+), 42 deletions(-)

diff --git a/buildSrc/src/main/kotlin/pill/parser.kt b/buildSrc/src/main/kotlin/pill/parser.kt
index e36da5861d3..99ec1fa2136 100644
--- a/buildSrc/src/main/kotlin/pill/parser.kt
+++ b/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,47 +204,46 @@ private fun parseSourceRoots(project: Project): List<PSourceRoot> {
 
     val sourceRoots = mutableListOf<PSourceRoot>()
 
-    project.configure<JavaPluginConvention> {
-        for ((sourceSetName, kind) in SOURCE_SET_MAPPING) {
-            val sourceSet = sourceSets.findByName(sourceSetName) ?: continue
+    for (sourceSet in project.sourceSets) {
+        val sourceSetName = sourceSet.name
+        val kind = if (sourceSetName == SourceSet.TEST_SOURCE_SET_NAME) Kind.TEST else Kind.PRODUCTION
 
-            fun Any.getKotlin(): SourceDirectorySet {
-                val kotlinMethod = javaClass.getMethod("getKotlin")
-                val oldIsAccessible = kotlinMethod.isAccessible
-                try {
-                    kotlinMethod.isAccessible = true
-                    return kotlinMethod(this) as SourceDirectorySet
-                } finally {
-                    kotlinMethod.isAccessible = oldIsAccessible
-                }
+        fun Any.getKotlin(): SourceDirectorySet {
+            val kotlinMethod = javaClass.getMethod("getKotlin")
+            val oldIsAccessible = kotlinMethod.isAccessible
+            try {
+                kotlinMethod.isAccessible = true
+                return kotlinMethod(this) as SourceDirectorySet
+            } finally {
+                kotlinMethod.isAccessible = oldIsAccessible
+            }
+        }
+
+        val kotlinSourceDirectories = (sourceSet as HasConvention).convention
+            .plugins["kotlin"]?.getKotlin()?.srcDirs
+                ?: emptySet()
+
+        val directories = (sourceSet.java.sourceDirectories.files + kotlinSourceDirectories).toList()
+            .filter { it.exists() }
+            .takeIf { it.isNotEmpty() }
+                ?: continue
+
+        for (resourceRoot in sourceSet.resources.sourceDirectories.files) {
+            if (!resourceRoot.exists() || resourceRoot in directories) {
+                continue
             }
 
-            val kotlinSourceDirectories = (sourceSet as HasConvention).convention
-                .plugins["kotlin"]?.getKotlin()?.srcDirs
-                    ?: emptySet()
-
-            val directories = (sourceSet.java.sourceDirectories.files + kotlinSourceDirectories).toList()
-                .filter { it.exists() }
-                .takeIf { it.isNotEmpty() }
-                    ?: continue
-
-            for (resourceRoot in sourceSet.resources.sourceDirectories.files) {
-                if (!resourceRoot.exists() || resourceRoot in directories) {
-                    continue
-                }
-
-                val resourceRootKind = when (kind) {
-                    Kind.PRODUCTION -> Kind.RESOURCES
-                    Kind.TEST -> Kind.TEST_RESOURCES
-                    else -> error("Invalid source root kind $kind")
-                }
-
-                sourceRoots += PSourceRoot(resourceRoot, resourceRootKind)
+            val resourceRootKind = when (kind) {
+                Kind.PRODUCTION -> Kind.RESOURCES
+                Kind.TEST -> Kind.TEST_RESOURCES
+                else -> error("Invalid source root kind $kind")
             }
 
-            for (directory in directories) {
-                sourceRoots += PSourceRoot(directory, kind)
-            }
+            sourceRoots += PSourceRoot(resourceRoot, resourceRootKind)
+        }
+
+        for (directory in directories) {
+            sourceRoots += PSourceRoot(directory, kind)
         }
     }
 
@@ -420,4 +415,11 @@ fun List<CollectedConfiguration>.collectDependencies(): List<DependencyInfo> {
     }
 
     return dependencies
-}
\ No newline at end of file
+}
+
+private val Project.sourceSets: SourceSetContainer
+    get() {
+        lateinit var result: SourceSetContainer
+        project.configure<JavaPluginConvention> { result = sourceSets }
+        return result
+    }
\ No newline at end of file