Refactor source configuration for android

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/internal/AnnotationProcessingManager.kt

@@ -54,6 +54,8 @@ fun Project.initKapt(
         kotlinTask.destinationDir = stubsDir
         kotlinTask.kaptOptions.generateStubs = true
 
+        kotlinAfterJavaTask.source(kaptManager.generatedKotlinSourceDir)
+        kotlinAfterJavaTask.source(kaptManager.aptOutputDir)
         val javaDestinationDir = project.files(javaTask.destinationDir)
         javaTask.doLast {
             kotlinAfterJavaTask.source(kotlinTask.source)
@@ -80,7 +82,6 @@ fun Project.initKapt(
 
     kotlinTask.doFirst {
         kaptManager.generateJavaHackFile()
-        kotlinAfterJavaTask?.source(kaptManager.getGeneratedKotlinSourceDir())
     }
 
     var originalJavaCompilerArgs: List<String>? = null
@@ -88,7 +89,7 @@ fun Project.initKapt(
         originalJavaCompilerArgs = (javaTask as JavaCompile).options.compilerArgs
         kaptManager.setupKapt()
         kaptManager.generateJavaHackFile()
-        kotlinAfterJavaTask?.source(kaptManager.getGeneratedKotlinSourceDir())
+        kotlinAfterJavaTask?.source(kaptManager.generatedKotlinSourceDir)
     }
 
     javaTask.doLast {
@@ -148,11 +149,12 @@ public class AnnotationProcessingManager(
         return File(wrappersDirectory, "annotations.$taskQualifier.txt")
     }
 
-    fun getGeneratedKotlinSourceDir(): File {
-        val kotlinGeneratedDir = File(aptWorkingDir, "kotlinGenerated")
-        if (!kotlinGeneratedDir.exists()) kotlinGeneratedDir.mkdirs()
-        return kotlinGeneratedDir
-    }
+    val generatedKotlinSourceDir: File
+        get() {
+            val kotlinGeneratedDir = File(aptWorkingDir, "kotlinGenerated")
+            if (!kotlinGeneratedDir.exists()) kotlinGeneratedDir.mkdirs()
+            return kotlinGeneratedDir
+        }
 
     fun setupKapt() {
         if (aptFiles.isEmpty()) return
@@ -207,7 +209,7 @@ public class AnnotationProcessingManager(
     private fun appendAnnotationsArguments() {
         javaTask.modifyCompilerArguments { list ->
             list.add("-Akapt.annotations=" + getAnnotationFile())
-            list.add("-Akapt.kotlin.generated=" + getGeneratedKotlinSourceDir())
+            list.add("-Akapt.kotlin.generated=" + generatedKotlinSourceDir)
         }
     }
 

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPlugin.kt

@@ -372,25 +372,9 @@ open class KotlinAndroidPlugin @Inject constructor(val scriptHandler: ScriptHand
             kotlinTask.description = "Compiles the ${variantDataName} kotlin."
             kotlinTask.setDependsOn(javaTask.dependsOn)
 
-            fun SourceDirectorySet.addSourceDirectories(additionalSourceFiles: Collection<File>) {
-                for (dir in additionalSourceFiles) {
-                    this.srcDir(dir)
-                    logger.kotlinDebug("Source directory ${dir.absolutePath} was added to kotlin source for $kotlinTaskName")
-                }
-            }
-
             val aptFiles = arrayListOf<File>()
 
-            // getSortedSourceProviders should return only actual java sources, generated sources should be collected earlier
-            val providers = variantData.variantConfiguration.sortedSourceProviders
-            for (provider in providers) {
-                val javaSrcDirs = AndroidGradleWrapper.getJavaSrcDirs(provider as AndroidSourceSet)
-                val kotlinSourceSet = getExtension<KotlinSourceSet>(provider, "kotlin")
-                val kotlinSourceDirectorySet = kotlinSourceSet.getKotlin()
-                kotlinTask.source(kotlinSourceDirectorySet)
-
-                kotlinSourceDirectorySet.addSourceDirectories(javaSrcDirs)
-
+            for (provider in variantData.sourceProviders) {
                 val aptConfiguration = aptConfigurations[(provider as AndroidSourceSet).name]
                 // Ignore if there's only an annotation processor wrapper in dependencies (added by default)
                 if (aptConfiguration != null && aptConfiguration.dependencies.size > 1) {
@@ -399,15 +383,6 @@ open class KotlinAndroidPlugin @Inject constructor(val scriptHandler: ScriptHand
                 }
             }
 
-            // getJavaSources should return the Java sources used for compilation
-            // We want to collect only generated files, like R-class output dir
-            // Actual java sources will be collected later
-            val additionalSourceFiles = AndroidGradleWrapper.getGeneratedSourceDirs(variantData)
-            for (file in additionalSourceFiles) {
-                kotlinTask.source(file)
-                logger.kotlinDebug("Source directory with generated files ${file.absolutePath} was added to kotlin source for $kotlinTaskName")
-            }
-
             subpluginEnvironment.addSubpluginArguments(project, kotlinTask)
 
             // should not be evaluated until right before compileKotlin evaluation since android can change
@@ -435,11 +410,34 @@ open class KotlinAndroidPlugin @Inject constructor(val scriptHandler: ScriptHand
                 kotlinAfterJavaTask = project.initKapt(kotlinTask, javaTask, kaptManager, variantDataName, kotlinOptions, subpluginEnvironment, tasksProvider)
             }
 
+            configureSources(kotlinTask, variantData)
+            if (kotlinAfterJavaTask != null) {
+                configureSources(kotlinAfterJavaTask, variantData)
+            }
+
             configureJavaTask(kotlinTask, javaTask, kotlinAfterJavaTask, logger)
             createSyncOutputTask(project, kotlinTask, javaTask, kotlinAfterJavaTask, variantDataName)
         }
     }
 
+    private fun configureSources(kotlinTask: KotlinCompile, variantData: BaseVariantData<out BaseVariantOutputData>) {
+        val logger = kotlinTask.project.logger
+
+        for (provider in variantData.sourceProviders) {
+            val kotlinSourceSet = getExtension<KotlinSourceSet>(provider, "kotlin")
+            val kotlinSourceDirectorySet = kotlinSourceSet.getKotlin()
+            kotlinTask.source(kotlinSourceDirectorySet)
+        }
+
+        for (javaSrcDir in AndroidGradleWrapper.getJavaSources(variantData)) {
+            kotlinTask.source(javaSrcDir)
+            logger.kotlinDebug("Source directory $javaSrcDir was added to kotlin source for ${kotlinTask.name}")
+        }
+    }
+
+    private val BaseVariantData<*>.sourceProviders: List<SourceProvider>
+        get() = variantConfiguration.sortedSourceProviders
+
     fun <T> getExtension(obj: Any, extensionName: String): T {
         if (obj is ExtensionAware) {
             val result = obj.extensions.findByName(extensionName)

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/android/AndroidGradleWrapper.groovy

@@ -8,6 +8,8 @@ import com.android.build.gradle.api.BaseVariant
 import com.android.build.gradle.api.TestVariant
 import com.android.build.gradle.internal.VariantManager
 import com.android.build.gradle.internal.variant.BaseVariantData
+import com.android.builder.model.SourceProvider
+import org.gradle.api.file.ConfigurableFileTree
 import org.gradle.api.internal.DefaultDomainObjectSet
 import org.gradle.api.tasks.compile.AbstractCompile
 import org.gradle.api.tasks.util.PatternFilterable
@@ -100,14 +102,24 @@ class AndroidGradleWrapper {
     return plugin.getVariantManager()
   }
 
-  static def List<File> getGeneratedSourceDirs(BaseVariantData variantData) {
+  static def List<File> getJavaSources(BaseVariantData variantData) {
     def result = new ArrayList<File>()
 
     def getJavaSourcesMethod = variantData.getMetaClass().getMetaMethod("getJavaSources")
     if (getJavaSourcesMethod.returnType.metaClass == Object[].metaClass) {
       result.addAll(variantData.getJavaSources().findAll { it instanceof File })
     }
+    else if (getJavaSourcesMethod.returnType.metaClass == List.metaClass) {
+      def fileTrees = variantData.getJavaSources().findAll { it instanceof ConfigurableFileTree }
+      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());
       }