fix: Include the main source set output

Author: bric3

buildSrc/src/main/kotlin/datadog/gradle/plugin/csi/CallSiteInstrumentationExtension.kt

@@ -5,7 +5,6 @@ import org.gradle.api.artifacts.Configuration
 import org.gradle.api.file.DirectoryProperty
 import org.gradle.api.file.ProjectLayout
 import org.gradle.api.model.ObjectFactory
-import org.gradle.api.plugins.JavaPlugin
 import org.gradle.api.plugins.JavaPlugin.COMPILE_CLASSPATH_CONFIGURATION_NAME
 import org.gradle.api.plugins.JavaPlugin.RUNTIME_CLASSPATH_CONFIGURATION_NAME
 import org.gradle.api.provider.ListProperty
@@ -16,6 +15,7 @@ import org.gradle.jvm.toolchain.JavaLanguageVersion
 import org.gradle.kotlin.dsl.listProperty
 import org.gradle.kotlin.dsl.property
 import java.io.File
+import java.util.Locale
 import javax.inject.Inject
 
 
@@ -83,13 +83,35 @@ abstract class CallSiteInstrumentationExtension @Inject constructor(
    * By default, includes all `main*` source sets, but only the `test`
    * (as wee don't want other test configurations by default).
    */
-  val configurations = objectFactory.listProperty<Configuration>().convention(
+  val configurations: ListProperty<Configuration> = objectFactory.listProperty<Configuration>().convention(
     project.provider {
       project.configurations.matching {
         // Includes all main* source sets, but only the test (as wee don;t want other )
-        (it.name.startsWith(MAIN_SOURCE_SET_NAME) || it.name == TEST_SOURCE_SET_NAME)
-            && (it.name.endsWith(RUNTIME_CLASSPATH_CONFIGURATION_NAME, ignoreCase = true) || it.name.endsWith(COMPILE_CLASSPATH_CONFIGURATION_NAME, ignoreCase = true))
+        // * For main => runtimeClasspath, compileClasspath
+        // * For test => testRuntimeClasspath, testCompileClasspath
+        // * For other main* => "main_javaXXRuntimeClasspath", "main_javaXXCompileClasspath"
+
+        when (it.name) {
+          // Regular main and test source sets
+          RUNTIME_CLASSPATH_CONFIGURATION_NAME,
+          COMPILE_CLASSPATH_CONFIGURATION_NAME,
+          TEST_SOURCE_SET_NAME + RUNTIME_CLASSPATH_CONFIGURATION_NAME.capitalize(),
+          TEST_SOURCE_SET_NAME + COMPILE_CLASSPATH_CONFIGURATION_NAME.capitalize() -> true
+
+          // Other main_javaXX source sets
+          else -> {
+            it.name.startsWith(MAIN_SOURCE_SET_NAME) &&
+                (it.name.endsWith(RUNTIME_CLASSPATH_CONFIGURATION_NAME, ignoreCase = true) ||
+                    it.name.endsWith(COMPILE_CLASSPATH_CONFIGURATION_NAME, ignoreCase = true))
+          }
+        }
       }
     }
   )
+
+  private fun String.capitalize(): String = replaceFirstChar {
+    if (it.isLowerCase()) it.titlecase(
+      Locale.getDefault()
+    ) else it.toString()
+  }
 }

buildSrc/src/main/kotlin/datadog/gradle/plugin/csi/CallSiteInstrumentationPlugin.kt

@@ -42,7 +42,7 @@ abstract class CallSiteInstrumentationPlugin : Plugin<Project>{
   private fun configureSourceSets(project: Project, extension: CallSiteInstrumentationExtension) {
     // create a new source set for the csi files
     val targetFolder = newBuildFolder(project, extension.targetFolder.get().asFile.toString())
-    val sourceSets = getSourceSets(project)
+    val sourceSets = project.sourceSets
     val mainSourceSet = sourceSets.named(MAIN_SOURCE_SET_NAME).get()
     val csiSourceSet = sourceSets.create("csi") {
       compileClasspath += mainSourceSet.output // mainly needed for the plugin tests
@@ -91,10 +91,6 @@ abstract class CallSiteInstrumentationPlugin : Plugin<Project>{
     return file
   }
 
-  private fun getSourceSets(project: Project): SourceSetContainer {
-    return project.extensions.getByType<JavaPluginExtension>().sourceSets
-  }
-
   private fun createTasks(project: Project, extension: CallSiteInstrumentationExtension) {
     registerGenerateCallSiteTask(project, extension, "compileJava")
     val targetFolder = extension.targetFolder.get().asFile
@@ -152,17 +148,20 @@ abstract class CallSiteInstrumentationPlugin : Plugin<Project>{
 
       // Write the call site instrumenter arguments into a temporary file
       doFirst {
+        val sourceSetOutput = project.sourceSets.matching { it.name.startsWith(MAIN_SOURCE_SET_NAME) }.flatMap {
+          it.output.classesDirs.files
+        }
         val programClassPath = csiExtension.configurations.get().flatMap {
           it.files
-        }.map { it.toString() }
+        }
 
         val arguments = listOf(
           csiExtension.srcFolder.get().asFile.toString(),
           inputProvider.get().asFile.toString(),
           output.get().asFile.toString(),
           csiExtension.suffix.get(),
           csiExtension.reporters.get().joinToString(",")
-        ) + programClassPath
+        ) + (sourceSetOutput + programClassPath).map { it.toString() }
 
         val argumentFile = newTempFile(temporaryDir, "call-site-arguments")
         Files.write(argumentFile.toPath(), arguments)
@@ -174,7 +173,7 @@ abstract class CallSiteInstrumentationPlugin : Plugin<Project>{
     }
 
     // make all sourcesets' class tasks depend on call site generator
-    val sourceSets = getSourceSets(project)
+    val sourceSets = project.sourceSets
     sourceSets.named(MAIN_SOURCE_SET_NAME) {
       project.tasks.named(classesTaskName) {
         dependsOn(callSiteGeneratorTask)
@@ -188,4 +187,7 @@ abstract class CallSiteInstrumentationPlugin : Plugin<Project>{
       }
     }
   }
+
+  private val Project.sourceSets: SourceSetContainer
+    get() = project.extensions.getByType<JavaPluginExtension>().sourceSets
 }