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 1ca707ca..31f13e08 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
@@ -8,21 +8,21 @@ enum class SourceSet {
     MAIN,
     TEST;
 
-    private fun unknown(sourceSet: SourceSet) = throw KobaltException("Unknown source set: $this")
+    fun correctSourceSet(project: Project) = when(this) {
+        SourceSet.MAIN -> project.sourceDirectories
+        SourceSet.TEST -> project.sourceDirectoriesTest
+        else -> unknown(this)
+    }
 
-    fun correctSourceSet(project: Project) =
-        if (this == SourceSet.MAIN) project.sourceDirectories
-        else if (this == SourceSet.TEST) project.sourceDirectoriesTest
-        else unknown(this)
-
-    fun correctOutputDir(project: Project) =
-        if (this == SourceSet.MAIN) KFiles.CLASSES_DIR
-        else if (this == SourceSet.TEST) KFiles.TEST_CLASSES_DIR
-        else unknown(this)
+    fun correctOutputDir(project: Project) = when(this) {
+        SourceSet.MAIN -> KFiles.CLASSES_DIR
+        SourceSet.TEST -> KFiles.TEST_CLASSES_DIR
+        else -> unknown(this)
+    }
 
     companion object {
         fun of(isTest: Boolean) = if (isTest) TEST else MAIN
-
+        private fun unknown(sourceSet: SourceSet) = throw KobaltException("Unknown source set: $sourceSet")
     }
 }