Derive the java class name for CelLiteDescriptor

PiperOrigin-RevId: 755424324

Author: l46kok

java_lite_proto_cel_library.bzl

@@ -19,15 +19,17 @@ load("@com_google_protobuf//bazel:java_lite_proto_library.bzl", "java_lite_proto
 
 def java_lite_proto_cel_library(
         name,
-        java_descriptor_class_name,
         proto_src,
+        java_descriptor_class_name = None,
         debug = False):
     """Generates a CelLiteDescriptor
 
     Args:
        name: name of this target.
-       java_descriptor_class_name: Name of the generated descriptor java class.
        proto_src: Name of the proto_library target.
+       java_descriptor_class_name (optional): Java class name for the generated CEL lite descriptor.
+                                  By default, CEL will use the first encountered message name in proto_src with "CelLiteDescriptor"
+                                  suffixed as the class name. Use this field to override this name.
        debug: (optional) If true, prints additional information during codegen for debugging purposes.
     """
     java_proto_library_dep = name + "_java_lite_proto_dep"
@@ -37,9 +39,9 @@ def java_lite_proto_cel_library(
     )
 
     java_lite_proto_cel_library_impl(
-        name,
-        java_descriptor_class_name,
-        proto_src,
-        java_proto_library_dep,
-        debug,
+        name = name,
+        proto_src = proto_src,
+        java_descriptor_class_name = java_descriptor_class_name,
+        java_proto_library_dep = java_proto_library_dep,
+        debug = debug,
     )

java_lite_proto_cel_library_impl.bzl

@@ -18,39 +18,39 @@ This is an implementation detail. Clients should use 'java_lite_proto_cel_librar
 """
 
 load("@rules_java//java:defs.bzl", "java_library")
-load("@rules_proto//proto:defs.bzl", "proto_descriptor_set")
 load("//publish:cel_version.bzl", "CEL_VERSION")
 load("@com_google_protobuf//bazel:java_lite_proto_library.bzl", "java_lite_proto_library")
+load("@rules_proto//proto:defs.bzl", "ProtoInfo")
 
 def java_lite_proto_cel_library_impl(
         name,
-        java_descriptor_class_name,
         proto_src,
         java_proto_library_dep,
+        java_descriptor_class_name = None,
         debug = False):
     """Generates a CelLiteDescriptor
 
     Args:
        name: name of this target.
-       java_descriptor_class_name: Name of the generated descriptor java class.
        proto_src: Name of the proto_library target.
-       java_proto_library_dep: (optional) Uses the provided java_lite_proto_library or java_proto_library to generate the lite descriptors. If none is provided, java_lite_proto_library is used by default behind the scenes. Most use cases should not need to provide this.
+       java_descriptor_class_name (optional): Java class name for the generated CEL lite descriptor.
+                                   By default, CEL will use the first encountered message name in proto_src with "CelLiteDescriptor"
+                                   suffixed as the class name. Use this field to override this name.
+       java_proto_library_dep: (optional) Uses the provided java_lite_proto_library or java_proto_library to generate the lite descriptors.
+                                If none is provided, java_lite_proto_library is used by default behind the scenes. Most use cases should not need to provide this.
        debug: (optional) If true, prints additional information during codegen for debugging purposes.
     """
     if not name:
         fail("You must provide a name.")
 
-    if not java_descriptor_class_name:
-        fail("You must provide a descriptor_class_prefix.")
-
     if not proto_src:
         fail("You must provide a proto_library dependency.")
 
-    _generate_cel_lite_descriptor_class(
-        name,
-        java_descriptor_class_name,
-        proto_src,
-        debug,
+    generated = name + "_cel_lite_descriptor"
+    java_lite_proto_cel_library_rule(
+        name = generated,
+        descriptor = proto_src,
+        java_descriptor_class_name = java_descriptor_class_name,
     )
 
     if not java_proto_library_dep:
@@ -67,44 +67,52 @@ def java_lite_proto_cel_library_impl(
 
     java_library(
         name = name,
-        srcs = [":" + name + "_cel_lite_descriptor"],
+        srcs = [":" + generated],
         deps = descriptor_codegen_deps,
     )
 
-def _generate_cel_lite_descriptor_class(
-        name,
-        descriptor_class_name,
-        proto_src,
-        debug):
-    outfile = "%s.java" % descriptor_class_name
-
-    transitive_descriptor_set_name = "%s_transitive_descriptor_set" % name
-    proto_descriptor_set(
-        name = transitive_descriptor_set_name,
-        deps = [proto_src],
+def _generate_cel_lite_descriptor_class(ctx):
+    srcjar_output = ctx.actions.declare_file(ctx.attr.name + ".srcjar")
+    java_file_path = srcjar_output.path
+
+    proto_info = ctx.attr.descriptor[ProtoInfo]
+    transitive_descriptors = proto_info.transitive_descriptor_sets
+
+    args = ctx.actions.args()
+    args.add("--version", CEL_VERSION)
+    args.add("--descriptor", proto_info.direct_descriptor_set)
+    args.add_joined("--transitive_descriptor_set", transitive_descriptors, join_with = ",")
+    args.add("--out", java_file_path)
+
+    if ctx.attr.java_descriptor_class_name:
+        args.add("--overridden_descriptor_class_name", ctx.attr.java_descriptor_class_name)
+    if ctx.attr.debug:
+        args.add("--debug")
+
+    ctx.actions.run(
+        mnemonic = "CelLiteDescriptorGenerator",
+        arguments = [args],
+        inputs = transitive_descriptors,
+        outputs = [srcjar_output],
+        progress_message = "Generating CelLiteDescriptor for: " + ctx.attr.name,
+        executable = ctx.executable._tool,
     )
 
-    direct_descriptor_set_name = proto_src
-
-    debug_flag = "--debug" if debug else ""
-
-    cmd = (
-        "$(location //protobuf:cel_lite_descriptor_generator) " +
-        "--descriptor $(location %s) " % direct_descriptor_set_name +
-        "--transitive_descriptor_set $(location %s) " % transitive_descriptor_set_name +
-        "--descriptor_class_name %s " % descriptor_class_name +
-        "--out $(location %s) " % outfile +
-        "--version %s " % CEL_VERSION +
-        debug_flag
-    )
-
-    native.genrule(
-        name = name + "_cel_lite_descriptor",
-        srcs = [
-            transitive_descriptor_set_name,
-            direct_descriptor_set_name,
-        ],
-        cmd = cmd,
-        outs = [outfile],
-        tools = ["//protobuf:cel_lite_descriptor_generator"],
-    )
+    return [DefaultInfo(files = depset([srcjar_output]))]
+
+java_lite_proto_cel_library_rule = rule(
+    implementation = _generate_cel_lite_descriptor_class,
+    attrs = {
+        "java_descriptor_class_name": attr.string(),
+        "descriptor": attr.label(
+            providers = [ProtoInfo],
+        ),
+        "debug": attr.bool(),
+        "_tool": attr.label(
+            executable = True,
+            cfg = "exec",
+            allow_files = True,
+            default = Label("//protobuf:cel_lite_descriptor_generator"),
+        ),
+    },
+)

protobuf/src/main/java/dev/cel/protobuf/CelLiteDescriptorGenerator.java

@@ -25,13 +25,17 @@
 import dev.cel.protobuf.JavaFileGenerator.JavaFileGeneratorOption;
 import java.io.File;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.Callable;
 import picocli.CommandLine;
 import picocli.CommandLine.Model.OptionSpec;
 import picocli.CommandLine.Option;
 
 final class CelLiteDescriptorGenerator implements Callable<Integer> {
 
+  private static final String DEFAULT_CEL_LITE_DESCRIPTOR_CLASS_SUFFIX = "CelLiteDescriptor";
+
   @Option(
       names = {"--out"},
       description = "Outpath for the CelLiteDescriptor")
@@ -46,13 +50,14 @@ final class CelLiteDescriptorGenerator implements Callable<Integer> {
 
   @Option(
       names = {"--transitive_descriptor_set"},
+      split = ",",
       description = "Path to the transitive set of descriptors")
-  private String transitiveDescriptorSetPath = "";
+  private List<String> transitiveDescriptorSetPath = new ArrayList<>();
 
   @Option(
-      names = {"--descriptor_class_name"},
-      description = "Class name for the CelLiteDescriptor")
-  private String descriptorClassName = "";
+      names = {"--overridden_descriptor_class_name"},
+      description = "Java class name for the CelLiteDescriptor")
+  private String overriddenDescriptorClassName = "";
 
   @Option(
       names = {"--version"},
@@ -70,14 +75,13 @@ final class CelLiteDescriptorGenerator implements Callable<Integer> {
   public Integer call() throws Exception {
     String targetDescriptorProtoPath = extractProtoPath(targetDescriptorPath);
     debugPrinter.print("Target descriptor proto path: " + targetDescriptorProtoPath);
+    FileDescriptorSet transitiveDescriptorSet = combineFileDescriptors(transitiveDescriptorSetPath);
 
     FileDescriptor targetFileDescriptor = null;
     ImmutableSet<FileDescriptor> transitiveFileDescriptors =
-        CelDescriptorUtil.getFileDescriptorsFromFileDescriptorSet(
-            load(transitiveDescriptorSetPath));
+        CelDescriptorUtil.getFileDescriptorsFromFileDescriptorSet(transitiveDescriptorSet);
     for (FileDescriptor fd : transitiveFileDescriptors) {
       if (fd.getFullName().equals(targetDescriptorProtoPath)) {
-        debugPrinter.print("Transitive Descriptor Path: " + fd.getFullName());
         targetFileDescriptor = fd;
         break;
       }
@@ -97,17 +101,34 @@ public Integer call() throws Exception {
 
   private void codegenCelLiteDescriptor(FileDescriptor targetFileDescriptor) throws Exception {
     String javaPackageName = ProtoJavaQualifiedNames.getJavaPackageName(targetFileDescriptor);
+    String javaClassName = overriddenDescriptorClassName;
+    if (javaClassName.isEmpty()) {
+      // Derive the java class name. Use first encountered message/enum in the FDS as a default,
+      // with a suffix applied for uniqueness (we don't want to collide with java protoc default
+      // generated class name).
+      if (!targetFileDescriptor.getMessageTypes().isEmpty()) {
+        javaClassName = targetFileDescriptor.getMessageTypes().get(0).getName();
+      } else if (!targetFileDescriptor.getEnumTypes().isEmpty()) {
+        javaClassName = targetFileDescriptor.getEnumTypes().get(0).getName();
+      } else {
+        throw new IllegalArgumentException(
+            "File descriptor does not contain any messages or enums!");
+      }
+
+      javaClassName += DEFAULT_CEL_LITE_DESCRIPTOR_CLASS_SUFFIX;
+    }
     ProtoDescriptorCollector descriptorCollector =
         ProtoDescriptorCollector.newInstance(debugPrinter);
 
     debugPrinter.print(
-        String.format("Descriptor Java class name: %s.%s", javaPackageName, descriptorClassName));
+        String.format(
+            "Fully qualified descriptor java class name: %s.%s", javaPackageName, javaClassName));
 
     JavaFileGenerator.createFile(
         outPath,
         JavaFileGeneratorOption.newBuilder()
             .setVersion(version)
-            .setDescriptorClassName(descriptorClassName)
+            .setDescriptorClassName(javaClassName)
             .setPackageName(javaPackageName)
             .setDescriptorMetadataList(
                 descriptorCollector.collectCodegenMetadata(targetFileDescriptor))
@@ -127,7 +148,18 @@ private String extractProtoPath(String descriptorPath) {
     return fileDescriptorProto.getName();
   }
 
-  private FileDescriptorSet load(String descriptorSetPath) {
+  private FileDescriptorSet combineFileDescriptors(List<String> descriptorPaths) {
+    FileDescriptorSet.Builder combinedDescriptorBuilder = FileDescriptorSet.newBuilder();
+
+    for (String descriptorPath : descriptorPaths) {
+      FileDescriptorSet loadedFds = load(descriptorPath);
+      combinedDescriptorBuilder.addAllFile(loadedFds.getFileList());
+    }
+
+    return combinedDescriptorBuilder.build();
+  }
+
+  private static FileDescriptorSet load(String descriptorSetPath) {
     try {
       byte[] descriptorBytes = Files.toByteArray(new File(descriptorSetPath));
       return FileDescriptorSet.parseFrom(descriptorBytes, ExtensionRegistry.getEmptyRegistry());

protobuf/src/main/java/dev/cel/protobuf/JavaFileGenerator.java

@@ -19,17 +19,22 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.io.Files;
+import com.google.common.io.ByteStreams;
 // CEL-Internal-5
 import freemarker.template.Configuration;
 import freemarker.template.DefaultObjectWrapperBuilder;
 import freemarker.template.Template;
 import freemarker.template.TemplateException;
 import freemarker.template.Version;
-import java.io.File;
+import java.io.ByteArrayInputStream;
+import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.StringWriter;
 import java.io.Writer;
+import java.util.Locale;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
 
 final class JavaFileGenerator {
 
@@ -50,10 +55,27 @@ public static void createFile(String filePath, JavaFileGeneratorOption option)
 
     Template template = cfg.getTemplate(HELPER_CLASS_TEMPLATE_FILE);
     Writer out = new StringWriter();
-
     template.process(option.getTemplateMap(), out);
 
-    Files.asCharSink(new File(filePath), UTF_8).write(out.toString());
+    writeSrcJar(filePath, option.descriptorClassName(), out.toString());
+  }
+
+  private static void writeSrcJar(
+      String srcjarFilePath, String javaClassName, String javaClassContent) throws IOException {
+    if (!srcjarFilePath.toLowerCase(Locale.getDefault()).endsWith(".srcjar")) {
+      throw new IllegalArgumentException("File must end with .srcjar, provided: " + srcjarFilePath);
+    }
+    try (FileOutputStream fos = new FileOutputStream(srcjarFilePath);
+        ZipOutputStream zos = new ZipOutputStream(fos)) {
+      ZipEntry entry = new ZipEntry(javaClassName + ".java");
+      zos.putNextEntry(entry);
+
+      try (InputStream inputStream = new ByteArrayInputStream(javaClassContent.getBytes(UTF_8))) {
+        ByteStreams.copy(inputStream, zos);
+      }
+
+      zos.closeEntry();
+    }
   }
 
   @AutoValue

protobuf/src/test/java/dev/cel/protobuf/CelLiteDescriptorTest.java

@@ -17,7 +17,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
-import dev.cel.expr.conformance.proto3.TestAllTypesProto3LiteCelDescriptor;
+import dev.cel.expr.conformance.proto3.TestAllTypesCelLiteDescriptor;
 import dev.cel.protobuf.CelLiteDescriptor.FieldLiteDescriptor;
 import dev.cel.protobuf.CelLiteDescriptor.FieldLiteDescriptor.EncodingType;
 import dev.cel.protobuf.CelLiteDescriptor.FieldLiteDescriptor.JavaType;
@@ -29,8 +29,8 @@
 @RunWith(TestParameterInjector.class)
 public class CelLiteDescriptorTest {
 
-  private static final TestAllTypesProto3LiteCelDescriptor TEST_ALL_TYPES_CEL_LITE_DESCRIPTOR =
-      TestAllTypesProto3LiteCelDescriptor.getDescriptor();
+  private static final TestAllTypesCelLiteDescriptor TEST_ALL_TYPES_CEL_LITE_DESCRIPTOR =
+      TestAllTypesCelLiteDescriptor.getDescriptor();
 
   @Test
   public void getProtoTypeNamesToDescriptors_containsAllMessages() {