Generate KSP classes jar as output (#1329)

smocherla-brex · web-flow · commit 17862f9eaa64 · 2025-08-12T18:34:57.000-07:00
* Generate KSP classes jar as output

* Revert unrelated changes

* rename var for consistency

* Attempt to fix tests

* lint

* Add tests
diff --git a/kotlin/internal/jvm/compile.bzl b/kotlin/internal/jvm/compile.bzl
@@ -419,6 +419,7 @@ def _run_ksp_builder_actions(
         A struct containing KSP outputs
     """
     ksp_generated_java_srcjar = ctx.actions.declare_file(ctx.label.name + "-ksp-kt-gensrc.jar")
+    ksp_generated_classes_jar = ctx.actions.declare_file(ctx.label.name + "-ksp-kt-genclasses.jar")
 
     _run_kt_builder_action(
         ctx = ctx,
@@ -433,12 +434,13 @@ def _run_ksp_builder_actions(
         plugins = plugins,
         outputs = {
             "ksp_generated_java_srcjar": ksp_generated_java_srcjar,
+            "ksp_generated_classes_jar": ksp_generated_classes_jar,
         },
         build_kotlin = False,
         mnemonic = "KotlinKsp",
     )
 
-    return struct(ksp_generated_class_jar = ksp_generated_java_srcjar)
+    return struct(ksp_generated_class_jar = ksp_generated_classes_jar, ksp_generated_src_jar = ksp_generated_java_srcjar)
 
 def _run_kt_builder_action(
         ctx,
@@ -804,6 +806,7 @@ def _run_kt_java_builder_actions(
 
     # Run KSP
     ksp_generated_class_jar = None
+    ksp_generated_src_jar = None
     if has_kt_sources and ksp_annotation_processors:
         ksp_outputs = _run_ksp_builder_actions(
             ctx,
@@ -817,7 +820,9 @@ def _run_kt_java_builder_actions(
             plugins = plugins,
         )
         ksp_generated_class_jar = ksp_outputs.ksp_generated_class_jar
-        generated_ksp_src_jars.append(ksp_generated_class_jar)
+        output_jars.append(ksp_generated_class_jar)
+        ksp_generated_src_jar = ksp_outputs.ksp_generated_src_jar
+        generated_ksp_src_jars.append(ksp_generated_src_jar)
 
     java_infos = []
 
@@ -942,7 +947,7 @@ def _run_kt_java_builder_actions(
     if annotation_processors or ksp_annotation_processors:
         is_ksp = (ksp_annotation_processors != None)
         processor = ksp_annotation_processors if is_ksp else annotation_processors
-        gen_jar = ksp_generated_class_jar if is_ksp else ap_generated_src_jar
+        gen_jar = ksp_generated_src_jar if is_ksp else ap_generated_src_jar
         outputs_list = [java_info.outputs for java_info in java_infos]
         annotation_processing = _create_annotation_processing(
             annotation_processors = processor,
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/KotlinBuilder.kt b/src/main/kotlin/io/bazel/kotlin/builder/tasks/KotlinBuilder.kt
@@ -88,6 +88,7 @@ class KotlinBuilder
         REDUCED_CLASSPATH_MODE("--reduced_classpath_mode"),
         INSTRUMENT_COVERAGE("--instrument_coverage"),
         KSP_GENERATED_JAVA_SRCJAR("--ksp_generated_java_srcjar"),
+        KSP_GENERATED_CLASSES_JAR("--ksp_generated_classes_jar"),
         BUILD_TOOLS_API("--build_tools_api"),
       }
     }
@@ -222,6 +223,9 @@ class KotlinBuilder
           argMap.optionalSingle(KotlinBuilderFlags.KSP_GENERATED_JAVA_SRCJAR)?.let {
             generatedKspSrcJar = it
           }
+          argMap.optionalSingle(KotlinBuilderFlags.KSP_GENERATED_CLASSES_JAR)?.let {
+            generatedKspClassesJar = it
+          }
         }
 
         with(root.directoriesBuilder) {
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/CompilationTask.kt b/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/CompilationTask.kt
@@ -422,6 +422,21 @@ internal fun JvmCompilationTask.createGeneratedKspKotlinSrcJar() {
   }
 }
 
+/**
+ * Produce a jar of classes generated by KSP.
+ */
+internal fun JvmCompilationTask.createdGeneratedKspClassesJar() {
+  JarCreator(
+    path = Paths.get(outputs.generatedKspClassesJar),
+    normalize = true,
+    verbose = false,
+  ).also {
+    it.addDirectory(Paths.get(directories.generatedClasses))
+    it.setJarOwner(info.label, info.bazelRuleKind)
+    it.execute()
+  }
+}
+
 /**
  * Compiles Kotlin sources to classes. Does not compile Java sources.
  */
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/KotlinJvmTaskExecutor.kt b/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/KotlinJvmTaskExecutor.kt
@@ -150,6 +150,9 @@ class KotlinJvmTaskExecutor
           if (outputs.generatedKspSrcJar.isNotEmpty()) {
             context.execute("creating KSP generated src jar", ::createGeneratedKspKotlinSrcJar)
           }
+          if (outputs.generatedKspClassesJar.isNotEmpty()) {
+            context.execute("creating KSP generated classes jar", ::createdGeneratedKspClassesJar)
+          }
         }
       }
     }
diff --git a/src/main/protobuf/kotlin_model.proto b/src/main/protobuf/kotlin_model.proto
@@ -131,6 +131,8 @@ message JvmCompilationTask {
     string generated_class_jar = 7;
     // Source jar containing the generated KSP sources.
     string generated_ksp_src_jar = 8;
+    // The path to the jar containing the generated KSP classes
+    string generated_ksp_classes_jar = 9;
   }
 
   message Inputs {
diff --git a/src/test/data/jvm/ksp/BUILD b/src/test/data/jvm/ksp/BUILD
@@ -48,6 +48,14 @@ kt_ksp_plugin(
     ],
 )
 
+kt_ksp_plugin(
+    name = "generate_bytecode_plugin",
+    processor_class = "com.example.BytecodeGeneratorProcessor",
+    deps = [
+        "//src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/com/example:processor",
+    ],
+)
+
 kt_jvm_library(
     name = "ksp_kotlin_resources",
     srcs = ["CoffeeAppModel.kt"],
@@ -140,9 +148,40 @@ kt_jvm_library(
     ],
 )
 
+kt_jvm_library(
+    name = "ksp_generate_bytecode",
+    srcs = ["BytecodeExample.kt"],
+    plugins = [":generate_bytecode_plugin"],
+    deps = [
+        "//src/test/data/jvm/ksp/bytecodegenerator/annotation",
+    ],
+)
+
+kt_jvm_library(
+    name = "ksp_bytecode_plugin_generates_no_classes_with_other_plugins",
+    # these files don't have relevant annotations, and don't generate the bytecode/generated classes
+    # from generate_bytecode_plugin
+    srcs = [
+        "CoffeeApp.kt",
+        "CoffeeBean.java",
+        "CoffeeMaker.kt",
+        "DripCoffeeModule.kt",
+    ],
+    plugins = [
+        ":generate_bytecode_plugin",
+        ":dagger",
+    ],
+    deps = [
+        "@kotlin_rules_maven//:com_google_dagger_dagger",
+        "@kotlin_rules_maven//:com_google_dagger_dagger_compiler",
+    ],
+)
+
 filegroup(
     name = "ksp",
     srcs = [
+        ":ksp_bytecode_plugin_generates_no_classes_with_other_plugins.jar",
+        ":ksp_generate_bytecode.jar",
         ":ksp_kotlin_resources.jar",
         ":ksp_kotlin_resources_missing_plugin.jar",
         ":ksp_kotlin_resources_multiple_plugins.jar",
diff --git a/src/test/data/jvm/ksp/BytecodeExample.kt b/src/test/data/jvm/ksp/BytecodeExample.kt
@@ -0,0 +1,6 @@
+package src.test.data.jvm.ksp
+
+import src.test.data.jvm.ksp.bytecodegenerator.annotation.GenerateBytecode
+
+@GenerateBytecode
+class BytecodeExample
diff --git a/src/test/data/jvm/ksp/bytecodegenerator/annotation/BUILD.bazel b/src/test/data/jvm/ksp/bytecodegenerator/annotation/BUILD.bazel
@@ -0,0 +1,22 @@
+load("//kotlin:jvm.bzl", "kt_jvm_library")
+
+# Copyright 2018 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+package(default_visibility = ["//visibility:private"])
+
+kt_jvm_library(
+    name = "annotation",
+    srcs = ["GenerateBytecode.kt"],
+    visibility = ["//src/test:__subpackages__"],
+)
diff --git a/src/test/data/jvm/ksp/bytecodegenerator/annotation/GenerateBytecode.kt b/src/test/data/jvm/ksp/bytecodegenerator/annotation/GenerateBytecode.kt
@@ -0,0 +1,5 @@
+package src.test.data.jvm.ksp.bytecodegenerator.annotation
+
+@Target(AnnotationTarget.CLASS)
+@Retention(AnnotationRetention.SOURCE)
+annotation class GenerateBytecode
diff --git a/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/com/example/BUILD.bazel b/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/com/example/BUILD.bazel
@@ -0,0 +1,33 @@
+load("//kotlin:jvm.bzl", "kt_jvm_library")
+
+# Copyright 2018 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+package(default_visibility = ["//visibility:private"])
+
+kt_jvm_library(
+    name = "processor",
+    srcs = [
+        "BytecodeGeneratorProcessor.kt",
+        "BytecodeGeneratorProcessorProvider.kt",
+    ],
+    resource_strip_prefix = "src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/resources",
+    resources = ["//src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/resources"],
+    visibility = ["//src/test:__subpackages__"],
+    deps = [
+        "//kotlin/compiler:symbol-processing",
+        "//kotlin/compiler:symbol-processing-api",
+        "//src/test/data/jvm/ksp/bytecodegenerator/annotation",
+        "@kotlin_rules_maven//:org_ow2_asm_asm",
+    ],
+)
diff --git a/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/com/example/BytecodeGeneratorProcessor.kt b/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/com/example/BytecodeGeneratorProcessor.kt
@@ -0,0 +1,92 @@
+package com.example
+
+import com.google.devtools.ksp.processing.*
+import com.google.devtools.ksp.symbol.KSAnnotated
+import com.google.devtools.ksp.symbol.KSClassDeclaration
+import org.objectweb.asm.ClassWriter
+import org.objectweb.asm.Opcodes
+import src.test.data.jvm.ksp.bytecodegenerator.annotation.GenerateBytecode
+
+/**
+ * An example KSP processor that generates bytecode to be used only for tests
+ */
+class BytecodeGeneratorProcessor(
+    private val codeGen: CodeGenerator,
+    private val logger: KSPLogger
+) : SymbolProcessor {
+  override fun process(resolver: Resolver): List<KSAnnotated> {
+    val annotated = resolver
+      .getSymbolsWithAnnotation(GenerateBytecode::class.qualifiedName!!)
+      .filterIsInstance<KSClassDeclaration>()
+      .toList()
+    if (annotated.isEmpty()) return emptyList()
+
+    annotated.forEach { cls ->
+      val origName = cls.simpleName.asString()
+      val genName  = "${origName}\$GeneratedDefinition\$"
+      val pkg       = cls.packageName.asString()
+      val internal = pkg.replace('.', '/') + "/" + genName
+
+      // Build class with ASM
+      val cw = ClassWriter(ClassWriter.COMPUTE_FRAMES or ClassWriter.COMPUTE_MAXS)
+      cw.visit(
+          Opcodes.V1_8,
+        Opcodes.ACC_PUBLIC or Opcodes.ACC_SUPER,
+        internal,
+        null,
+        "java/lang/Object",
+        null
+      )
+
+      // Constructor
+      cw.visitMethod(Opcodes.ACC_PUBLIC, "<init>", "()V", null, null).apply {
+        visitCode()
+        visitVarInsn(Opcodes.ALOAD, 0)
+        visitMethodInsn(
+            Opcodes.INVOKESPECIAL,
+          "java/lang/Object", "<init>", "()V", false
+        )
+        visitInsn(Opcodes.RETURN)
+        visitMaxs(0, 0)
+        visitEnd()
+      }
+
+      // greet() method
+      cw.visitMethod(
+          Opcodes.ACC_PUBLIC,
+        "greet",
+        "()Ljava/lang/String;",
+        null,
+        null
+      ).apply {
+        visitCode()
+        visitLdcInsn("Hello, $origName!")
+        visitInsn(Opcodes.ARETURN)
+        visitMaxs(0, 0)
+        visitEnd()
+      }
+
+      cw.visitEnd()
+      val bytecode = cw.toByteArray()
+
+      // Write out .class file
+      val deps = Dependencies(false, *annotated.mapNotNull { it.containingFile }.toTypedArray())
+      val out = codeGen.createNewFile(dependencies = deps, packageName = pkg, fileName = genName, extensionName =  "class")
+      out.use { it.write(bytecode) }
+
+      // Generate service file
+      val fqcn = "$pkg.$genName"
+      val serviceFile = codeGen.createNewFile(
+        dependencies = Dependencies(false),
+        packageName = "META-INF.services",
+        fileName = fqcn,
+        extensionName = ""
+      )
+      serviceFile.bufferedWriter().use { it.appendLine(fqcn) }
+
+      logger.info("Generated $pkg.$genName")
+    }
+
+    return emptyList()
+  }
+}
diff --git a/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/com/example/BytecodeGeneratorProcessorProvider.kt b/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/com/example/BytecodeGeneratorProcessorProvider.kt
@@ -0,0 +1,11 @@
+package com.example
+
+import com.google.devtools.ksp.processing.SymbolProcessor
+import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
+import com.google.devtools.ksp.processing.SymbolProcessorProvider
+
+class BytecodeGeneratorProcessorProvider : SymbolProcessorProvider {
+  override fun create(environment: SymbolProcessorEnvironment): SymbolProcessor {
+    return BytecodeGeneratorProcessor(environment.codeGenerator, environment.logger)
+  }
+}
diff --git a/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/resources/BUILD.bazel b/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/resources/BUILD.bazel
@@ -0,0 +1,20 @@
+# Copyright 2018 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+package(default_visibility = ["//visibility:private"])
+
+filegroup(
+    name = "resources",
+    srcs = ["META-INF/services/com.google.devtools.ksp.processing.SymbolProcessorProvider"],
+    visibility = ["//src/test:__subpackages__"],
+)
diff --git a/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/resources/META-INF/services/com.google.devtools.ksp.processing.SymbolProcessorProvider b/src/test/data/jvm/ksp/bytecodegenerator/processor/src/main/resources/META-INF/services/com.google.devtools.ksp.processing.SymbolProcessorProvider
@@ -0,0 +1 @@
+com.example.BytecodeGeneratorProcessorProvider
diff --git a/src/test/kotlin/io/bazel/kotlin/KotlinJvmKspAssertionTest.kt b/src/test/kotlin/io/bazel/kotlin/KotlinJvmKspAssertionTest.kt

Original file line number	Diff line number	Diff line change
`@@ -88,6 +88,7 @@ class KotlinBuilder`
`88`	`88`	`REDUCED_CLASSPATH_MODE("--reduced_classpath_mode"),`
`89`	`89`	`INSTRUMENT_COVERAGE("--instrument_coverage"),`
`90`	`90`	`KSP_GENERATED_JAVA_SRCJAR("--ksp_generated_java_srcjar"),`
	`91`	`+ KSP_GENERATED_CLASSES_JAR("--ksp_generated_classes_jar"),`
`91`	`92`	`BUILD_TOOLS_API("--build_tools_api"),`
`92`	`93`	`}`
`93`	`94`	`}`
`@@ -222,6 +223,9 @@ class KotlinBuilder`
`222`	`223`	`argMap.optionalSingle(KotlinBuilderFlags.KSP_GENERATED_JAVA_SRCJAR)?.let {`
`223`	`224`	`generatedKspSrcJar = it`
`224`	`225`	`}`
	`226`	`+ argMap.optionalSingle(KotlinBuilderFlags.KSP_GENERATED_CLASSES_JAR)?.let {`
	`227`	`+ generatedKspClassesJar = it`
	`228`	`+ }`
`225`	`229`	`}`
`226`	`230`
`227`	`231`	`with(root.directoriesBuilder) {`
Original file line number	Diff line number	Diff line change
`@@ -150,6 +150,9 @@ class KotlinJvmTaskExecutor`
`150`	`150`	`if (outputs.generatedKspSrcJar.isNotEmpty()) {`
`151`	`151`	`context.execute("creating KSP generated src jar", ::createGeneratedKspKotlinSrcJar)`
`152`	`152`	`}`
	`153`	`+ if (outputs.generatedKspClassesJar.isNotEmpty()) {`
	`154`	`+ context.execute("creating KSP generated classes jar", ::createdGeneratedKspClassesJar)`
	`155`	`+ }`
`153`	`156`	`}`
`154`	`157`	`}`
`155`	`158`	`}`
Original file line number	Diff line number	Diff line change
`@@ -131,6 +131,8 @@ message JvmCompilationTask {`
`131`	`131`	`string generated_class_jar = 7;`
`132`	`132`	`// Source jar containing the generated KSP sources.`
`133`	`133`	`string generated_ksp_src_jar = 8;`
	`134`	`+ // The path to the jar containing the generated KSP classes`
	`135`	`+ string generated_ksp_classes_jar = 9;`
`134`	`136`	`}`
`135`	`137`
`136`	`138`	`message Inputs {`