From 7fe29a7e55f1ec0c78628b81705ce8b15f9b785d Mon Sep 17 00:00:00 2001
From: Alexey Tsvetkov <alexey.tsvetkov@jetbrains.com>
Date: Mon, 25 Jul 2016 16:25:22 +0300
Subject: [PATCH] Refactoring: simplify source roots search

---
 .../jetbrains/kotlin/gradle/tasks/Tasks.kt    | 64 ++++++++-----------
 .../android/AndroidGradleWrapper.groovy       | 17 ++---
 2 files changed, 37 insertions(+), 44 deletions(-)

diff --git a/libraries/tools/kotlin-gradle-plugin-core/src/main/kotlin/org/jetbrains/kotlin/gradle/tasks/Tasks.kt b/libraries/tools/kotlin-gradle-plugin-core/src/main/kotlin/org/jetbrains/kotlin/gradle/tasks/Tasks.kt
index f96553d7923..51aa35661b3 100644
--- a/libraries/tools/kotlin-gradle-plugin-core/src/main/kotlin/org/jetbrains/kotlin/gradle/tasks/Tasks.kt
+++ b/libraries/tools/kotlin-gradle-plugin-core/src/main/kotlin/org/jetbrains/kotlin/gradle/tasks/Tasks.kt
@@ -141,8 +141,7 @@ open class KotlinCompile() : AbstractKotlinCompile<K2JVMCompilerArguments>() {
     override val compiler = K2JVMCompiler()
     override fun createBlankArgs(): K2JVMCompilerArguments = K2JVMCompilerArguments()
 
-    // Should be SourceDirectorySet or File
-    val srcDirsSources = HashSet<Any>()
+    private val sourceRoots = HashSet<File>()
 
     // lazy because name is probably not available when constructor is called
     private val taskBuildDirectory: File by lazy { File(File(project.buildDir, KOTLIN_BUILD_DIR_NAME), name) }
@@ -493,55 +492,48 @@ open class KotlinCompile() : AbstractKotlinCompile<K2JVMCompilerArguments>() {
     }
 
     private fun getJavaSourceRoots(): Set<File> =
-            getSource()
-            .filter { it.isJavaFile() }
-            .map { findSrcDirRoot(it) }
-            .filterNotNull()
-            .toSet()
+            findRootsForSources(getSource().filter { it.isJavaFile() })
 
-    private fun File.isJavaFile() = extension.equals(JavaFileType.INSTANCE.defaultExtension, ignoreCase = true)
+    private fun File.isJavaFile() =
+            extension.equals(JavaFileType.INSTANCE.defaultExtension, ignoreCase = true)
 
     // override setSource to track source directory sets and files (for generated android folders)
     override fun setSource(source: Any?) {
-        srcDirsSources.clear()
-        if (source is SourceDirectorySet) {
-            srcDirsSources.add(source)
-        }
-        else if (source is File) {
-            srcDirsSources.add(source)
-        }
+        sourceRoots.clear()
+        addSourceRoot(source)
         super.setSource(source)
     }
 
     // override source to track source directory sets and files (for generated android folders)
     override fun source(vararg sources: Any?): SourceTask? {
-        for (source in sources) {
-            if (source is SourceDirectorySet) {
-                srcDirsSources.add(source)
-            }
-            else if (source is File) {
-                srcDirsSources.add(source)
-            }
-        }
+        sources.forEach { addSourceRoot(it) }
         return super.source(sources)
     }
 
-    fun findSrcDirRoot(file: File): File? {
-        for (source in srcDirsSources) {
-            if (source is SourceDirectorySet) {
-                for (root in source.srcDirs) {
-                    if (FileUtil.isAncestor(root, file, false)) {
-                        return root
-                    }
-                }
-            }
-            else if (source is File) {
-                if (FileUtil.isAncestor(source, file, false)) {
-                    return source
+    fun findRootsForSources(sources: Iterable<File>): Set<File> {
+        val resultRoots = HashSet<File>()
+        val sourceDirs = sources.mapTo(HashSet()) { it.parentFile }
+
+        for (sourceDir in sourceDirs) {
+            for (sourceRoot in sourceRoots) {
+                if (FileUtil.isAncestor(sourceRoot, sourceDir, /* strict = */false)) {
+                    resultRoots.add(sourceRoot)
                 }
             }
         }
-        return null
+
+        return resultRoots
+    }
+
+    private fun addSourceRoot(source: Any?) {
+        when (source) {
+            is SourceDirectorySet -> {
+                sourceRoots.addAll(source.srcDirs)
+            }
+            is File -> {
+                sourceRoots.add(source)
+            }
+        }
     }
 }
 
diff --git a/libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/android/AndroidGradleWrapper.groovy b/libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/android/AndroidGradleWrapper.groovy
index bb0b0dd399a..70a4fb93105 100644
--- a/libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/android/AndroidGradleWrapper.groovy
+++ b/libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/android/AndroidGradleWrapper.groovy
@@ -103,8 +103,15 @@ class AndroidGradleWrapper {
   }
 
   static def List<File> getJavaSources(BaseVariantData variantData) {
-    def result = new ArrayList<File>()
+    def result = new LinkedHashSet<File>()
 
+    // user sources
+    List<SourceProvider> providers = variantData.variantConfiguration.getSortedSourceProviders();
+    for (SourceProvider provider : providers) {
+      result.addAll((provider as AndroidSourceSet).getJava().getSrcDirs());
+    }
+
+    // generated sources
     def getJavaSourcesMethod = variantData.getMetaClass().getMetaMethod("getJavaSources")
     if (getJavaSourcesMethod.returnType.metaClass == Object[].metaClass) {
       result.addAll(variantData.getJavaSources().findAll { it instanceof File })
@@ -114,12 +121,6 @@ class AndroidGradleWrapper {
       result.addAll(fileTrees.collect { it.getDir() })
     }
     else {
-      // Old impl copied from android tools source. Delete?
-      List<SourceProvider> providers = variantData.variantConfiguration.getSortedSourceProviders();
-      for (SourceProvider provider : providers) {
-        result.addAll((provider as AndroidSourceSet).getJava().getSourceDirectoryTrees());
-      }
-
       if (variantData.scope.getGenerateRClassTask() != null) {
         result.add(variantData.scope.getRClassSourceOutputDir());
       }
@@ -150,6 +151,6 @@ class AndroidGradleWrapper {
       }
     }
 
-    return result
+    return result.toList()
   }
 }