Resolve extra classpath lazily when configuring Protobuf generation tasks

shared-src/gradle-plugin/protokt/v1/gradle/ProtobufBuild.kt

@@ -21,7 +21,8 @@ import com.google.protobuf.gradle.ProtobufExtract
 import com.google.protobuf.gradle.ProtobufPlugin
 import com.google.protobuf.gradle.id
 import org.gradle.api.Project
-import org.gradle.api.file.FileCollection
+import org.gradle.api.artifacts.Configuration
+import org.gradle.api.provider.Provider
 import org.gradle.jvm.tasks.Jar
 import org.gradle.kotlin.dsl.apply
 import org.gradle.kotlin.dsl.configure
@@ -52,51 +53,90 @@ internal fun configureProtobufPlugin(
         }
 
         generateProtoTasks {
+            val mainExtractProtoAdditions = mutableListOf<TaskInputFiles>()
+            val testExtractProtoAdditions = mutableListOf<TaskInputFiles?>()
+
             for (task in all()) {
                 if (disableJava) {
                     task.builtins {
                         findByName("java")?.run(::remove)
                     }
                 }
 
+                val extensions = resolveExtensions(project, task)
+
+                if (task.isTestTask()) {
+                    testExtractProtoAdditions.add(extensions.test?.let { TaskInputFiles(extensions.taskName, it) })
+                } else {
+                    mainExtractProtoAdditions.add(TaskInputFiles(extensions.taskName, extensions.main))
+                }
+
                 task.plugins {
                     id(target.protocPluginName) {
-                        project.afterEvaluate {
-                            option("$KOTLIN_EXTRA_CLASSPATH=${extraClasspath(project, task)}")
-                            option("$GENERATE_TYPES=${ext.generate.types}")
-                            option("$GENERATE_DESCRIPTORS=${ext.generate.descriptors}")
-                            option("$GENERATE_GRPC_DESCRIPTORS=${ext.generate.grpcDescriptors}")
-                            option("$GENERATE_GRPC_KOTLIN_STUBS=${ext.generate.grpcKotlinStubs}")
-                            option("$FORMAT_OUTPUT=${ext.formatOutput}")
-                            option("$KOTLIN_TARGET=$target")
-                        }
+                        option("$KOTLIN_EXTRA_CLASSPATH=${extraClasspath(project, extensions)}")
+                        option("$GENERATE_TYPES=${ext.generate.types}")
+                        option("$GENERATE_DESCRIPTORS=${ext.generate.descriptors}")
+                        option("$GENERATE_GRPC_DESCRIPTORS=${ext.generate.grpcDescriptors}")
+                        option("$GENERATE_GRPC_KOTLIN_STUBS=${ext.generate.grpcKotlinStubs}")
+                        option("$FORMAT_OUTPUT=${ext.formatOutput}")
+                        option("$KOTLIN_TARGET=$target")
                     }
                 }
             }
+
+            project.handleExtraInputFiles(mainExtractProtoAdditions, testExtractProtoAdditions)
         }
     }
 }
 
-private fun extraClasspath(project: Project, task: GenerateProtoTask): String {
-    var extensions: FileCollection = project.configurations.getByName(EXTENSIONS)
-
-    if (task.isTest) {
-        extensions += project.configurations.getByName(TEST_EXTENSIONS)
+private fun Project.handleExtraInputFiles(main: List<TaskInputFiles>, test: List<TaskInputFiles?>) {
+    afterEvaluate {
+        handleExtractProtoAdditions(main, false)
+        handleExtractProtoAdditions(test, true)
     }
+}
 
-    // Must explicitly register input files here; if any extensions dependencies are project dependencies then Gradle
-    // won't pick them up as dependencies unless we do this. There may be a better way to do this but for now just
-    // manually do what protobuf-gradle-plugin used to do.
-    // https://github.com/google/protobuf-gradle-plugin/commit/0521fe707ccedee7a0b4ce0fb88409eefb04e59d
-    project.tasks.withType<ProtobufExtract> {
-        if (name.startsWith("extractInclude")) {
-            inputFiles.from(extensions)
+private fun Project.handleExtractProtoAdditions(additions: List<TaskInputFiles?>, test: Boolean) {
+    additions.filterNotNull().forEach {
+        // there is more than one task in multiplatform projects: e.g. extractIncludeProto and extractIncludeJvmMainProto
+        extractIncludeProtoTasks().forEach { task ->
+            val isTest = task.isTestTask()
+            if ((!test && !isTest) || (test && isTest)) {
+                logger.log(DEBUG_LOG_LEVEL, "Adding input files to task ${task.name} from ${it.taskName}")
+                task.inputFiles.from(it.inputFiles)
+            }
         }
     }
+}
 
-    return extensions.joinToString(";") { URLEncoder.encode(it.path, "UTF-8") }
+private class TaskInputFiles(
+    val taskName: String,
+    val inputFiles: Any
+)
+
+private fun extraClasspath(project: Project, extensions: ExtensionsConfigurations) =
+    project.objects.fileCollection().from(extensions.asList()).files.joinToString(";") { URLEncoder.encode(it.path, "UTF-8") }
+
+private class ExtensionsConfigurations(
+    val taskName: String,
+    val main: Provider<Configuration>,
+    val test: Provider<Configuration>?
+) {
+    fun asList() =
+        listOfNotNull(main, test)
 }
 
+private fun resolveExtensions(project: Project, task: GenerateProtoTask) =
+    ExtensionsConfigurations(
+        task.name,
+        project.configurations.named(EXTENSIONS),
+        if (task.isTestTask()) {
+            project.configurations.named(TEST_EXTENSIONS)
+        } else {
+            null
+        }
+    )
+
 private fun configureSources(project: Project) {
     project.afterEvaluate {
         if (project.tasks.findByName("sourcesJar") != null) {
@@ -114,3 +154,15 @@ private fun normalizePath(binaryPath: String) =
     } else {
         binaryPath
     }
+
+private fun GenerateProtoTask.isTestTask() =
+    name.isTestTask()
+
+private fun ProtobufExtract.isTestTask() =
+    name.isTestTask()
+
+private fun String.isTestTask() =
+    endsWith("TestProto")
+
+private fun Project.extractIncludeProtoTasks() =
+    project.tasks.withType<ProtobufExtract>().filter { it.name.startsWith("extractInclude") }

shared-src/gradle-plugin/protokt/v1/gradle/ProtoktBuild.kt

@@ -60,7 +60,6 @@ internal fun configureProtokt(
     injectKotlinPluginsIntoProtobufGradle()
 
     val config = project.createExtensionConfigurations()
-    project.configureProtobuf(disableJava, config, binary)
 
     // must wait for extension to resolve
     project.afterEvaluate {
@@ -69,6 +68,8 @@ internal fun configureProtokt(
             project.resolveProtoktGrpcDep(protoktVersion)
         ).forEach(config.extensions.dependencies::add)
     }
+
+    project.configureProtobuf(disableJava, config, binary)
 }
 
 private fun injectKotlinPluginsIntoProtobufGradle() {