Get sources from compile tasks

instead of default source set and its dependencies.

There are 3 ways (that I know of) to register extra sources to
compilation tasks:

  1. Add to default source sets.
  2. Add to dependencies of default source sets.
  3. Add to compile task directly, e.g., via AbstractCompile.source

KSP picked up from the first 2 only. Now it also does third.
diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt
index 7ef59a8..30bec0d 100644
--- a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt
+++ b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt
@@ -297,7 +297,13 @@
                 classOutputDir,
                 resourceOutputDir
             )
-            kotlinCompilation.allKotlinSourceSets.forEach { sourceSet -> kspTask.source(sourceSet.kotlin) }
+            kspTask.source(kotlinCompileTask.source)
+            if (kotlinCompileTask is AbstractKotlinCompile<*>) {
+                val sourceRoots = kotlinCompileTask.getSourceRoots()
+                if (sourceRoots is SourceRoots.ForJvm) {
+                    kspTask.source(sourceRoots.javaSourceRoots)
+                }
+            }
 
             // Don't support binary generation for K/N yet.
             // FIXME: figure out how to add user generated libraries.
diff --git a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/PlaygroundIT.kt b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/PlaygroundIT.kt
index 8611a0f..9076d2e 100644
--- a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/PlaygroundIT.kt
+++ b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/PlaygroundIT.kt
@@ -25,6 +25,8 @@
         JarFile(artifact).use { jarFile ->
             Assert.assertTrue(jarFile.getEntry("TestProcessor.log").size > 0)
             Assert.assertTrue(jarFile.getEntry("hello/HELLO.class").size > 0)
+            Assert.assertTrue(jarFile.getEntry("g/G.class").size > 0)
+            Assert.assertTrue(jarFile.getEntry("g/GBuilder.class").size > 0)
             Assert.assertTrue(jarFile.getEntry("com/example/AClassBuilder.class").size > 0)
         }
 
diff --git a/integration-tests/src/test/resources/playground/workload/G.kt b/integration-tests/src/test/resources/playground/workload/G.kt
new file mode 100644
index 0000000..f2f80ff
--- /dev/null
+++ b/integration-tests/src/test/resources/playground/workload/G.kt
@@ -0,0 +1,6 @@
+package g
+
+import com.example.annotation.Builder
+
+@Builder
+class G
diff --git a/integration-tests/src/test/resources/playground/workload/build.gradle.kts b/integration-tests/src/test/resources/playground/workload/build.gradle.kts
index 6b283b4..3c44d53 100644
--- a/integration-tests/src/test/resources/playground/workload/build.gradle.kts
+++ b/integration-tests/src/test/resources/playground/workload/build.gradle.kts
@@ -23,3 +23,12 @@
     arg("option1", "value1")
     arg("option2", "value2")
 }
+
+val compileKotlin: AbstractCompile by tasks
+tasks.register<Copy>("copyG") {
+    from("G.kt")
+    into(File(buildDir, "generatedSources").apply { mkdirs() })
+}.let {
+    // Magic. `map` creates a provider to propagate task dependency.
+    compileKotlin.source(it.map { it.destinationDir })
+}