From 53bf5cbd4b456f059a70f7a1fddc5f0a21d42838 Mon Sep 17 00:00:00 2001
From: Cedric Beust <cedric@beust.com>
Date: Sat, 20 Feb 2016 13:02:47 -0800
Subject: [PATCH] Refactor source sets.

---
 .../main/kotlin/com/beust/kobalt/Variant.kt   |  4 +---
 .../kotlin/com/beust/kobalt/api/Project.kt    |  3 +++
 .../kobalt/internal/JvmCompilerPlugin.kt      | 17 +++------------
 .../com/beust/kobalt/internal/SourceSet.kt    | 21 ++++++++++++++++++-
 4 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/Variant.kt b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/Variant.kt
index e2ff55ca..4a2252d2 100644
--- a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/Variant.kt
+++ b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/Variant.kt
@@ -68,9 +68,7 @@ class Variant(val initialProductFlavor: ProductFlavorConfig? = null,
     private fun sourceDirectories(project: Project, suffix: String, variantFirst: Boolean, sourceSet: SourceSet)
             : List<File> {
         val result = arrayListOf<File>()
-        val sourceDirectories = (if (sourceSet == SourceSet.MAIN) project.sourceDirectories
-            else if (sourceSet == SourceSet.TEST) project.sourceDirectoriesTest
-            else throw KobaltException("Unknown source set: $sourceSet"))
+        val sourceDirectories = sourceSet.correctSourceSet(project)
                 .filter { File(project.directory, it).exists() }
                 .map { File(it) }
 
diff --git a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/api/Project.kt b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/api/Project.kt
index 6d4fe57b..ed372289 100644
--- a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/api/Project.kt
+++ b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/api/Project.kt
@@ -1,10 +1,13 @@
 package com.beust.kobalt.api
 
+import com.beust.kobalt.KobaltException
 import com.beust.kobalt.TestConfig
 import com.beust.kobalt.api.annotation.Directive
 import com.beust.kobalt.internal.JvmCompilerPlugin
+import com.beust.kobalt.internal.SourceSet
 import com.beust.kobalt.maven.dependency.MavenDependency
 import com.beust.kobalt.misc.KFiles
+import com.beust.kobalt.project
 import java.io.File
 import java.util.*
 
diff --git a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/JvmCompilerPlugin.kt b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/JvmCompilerPlugin.kt
index 9d6c3645..76ae6499 100644
--- a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/JvmCompilerPlugin.kt
+++ b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/JvmCompilerPlugin.kt
@@ -115,17 +115,7 @@ open class JvmCompilerPlugin @Inject constructor(
      * Copy the resources from a source directory to the build one
      */
     protected fun copyResources(project: Project, sourceSet: SourceSet) {
-        val sourceDirs: ArrayList<String> = arrayListOf()
-        var outputDir: String?
-        if (sourceSet == SourceSet.MAIN) {
-            sourceDirs.addAll(project.sourceDirectories.filter { it.contains("resources") })
-            outputDir = KFiles.CLASSES_DIR
-        } else if (sourceSet == SourceSet.TEST) {
-            sourceDirs.addAll(project.sourceDirectoriesTest.filter { it.contains("resources") })
-            outputDir = KFiles.TEST_CLASSES_DIR
-        } else {
-            throw IllegalArgumentException("Unknown source set: $sourceSet")
-        }
+        var outputDir = sourceSet.correctOutputDir(project)
 
         val variantSourceDirs = context.variant.resourceDirectories(project, sourceSet)
         if (variantSourceDirs.size > 0) {
@@ -172,8 +162,7 @@ open class JvmCompilerPlugin @Inject constructor(
 
     private fun doTaskCompile(project: Project, isTest: Boolean): TaskResult {
         // Set up the source files now that we have the variant
-        sourceDirectories.addAll(context.variant.sourceDirectories(project, context,
-                if (isTest) SourceSet.TEST else SourceSet.MAIN))
+        sourceDirectories.addAll(context.variant.sourceDirectories(project, context, SourceSet.of(isTest)))
 
         val sourceDirectory = context.variant.maybeGenerateBuildConfig(project, context)
         if (sourceDirectory != null) {
@@ -264,7 +253,7 @@ open class JvmCompilerPlugin @Inject constructor(
      */
     protected fun createCompilerActionInfo(project: Project, context: KobaltContext, isTest: Boolean,
             sourceDirectories: List<File>, sourceSuffixes: List<String>): CompilerActionInfo {
-        copyResources(project, if (isTest) SourceSet.TEST else SourceSet.MAIN)
+        copyResources(project, SourceSet.of(isTest))
 
         val fullClasspath = if (isTest) dependencyManager.testDependencies(project, context)
             else dependencyManager.dependencies(project, context)
diff --git a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/SourceSet.kt b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/SourceSet.kt
index 7899ccb3..c094b82c 100644
--- a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/SourceSet.kt
+++ b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/SourceSet.kt
@@ -1,7 +1,26 @@
 package com.beust.kobalt.internal
 
+import com.beust.kobalt.KobaltException
+import com.beust.kobalt.api.Project
+import com.beust.kobalt.misc.KFiles
+
 enum class SourceSet {
     MAIN,
-    TEST
+    TEST;
+
+    fun correctSourceSet(project: Project) =
+        if (this == SourceSet.MAIN) project.sourceDirectories
+        else if (this == SourceSet.TEST) project.sourceDirectoriesTest
+        else throw KobaltException("Unknown source set: $this")
+
+    fun correctOutputDir(project: Project) =
+        if (this == SourceSet.MAIN) KFiles.CLASSES_DIR
+        else if (this == SourceSet.TEST) KFiles.TEST_CLASSES_DIR
+        else throw IllegalArgumentException("Unknown source set: $this")
+
+    companion object {
+        fun of(isTest: Boolean) = if (isTest) TEST else MAIN
+
+    }
 }