Split the code generators for .upb and .upbdefs.

bazel/upb_proto_library.bzl

@@ -153,26 +153,29 @@ _UpbDefsWrappedCcInfo = provider(fields = ["cc_info"])
 _WrappedGeneratedSrcsInfo = provider(fields = ["srcs"])
 _WrappedDefsGeneratedSrcsInfo = provider(fields = ["srcs"])
 
-def _compile_upb_protos(ctx, proto_info, proto_sources, ext):
+def _compile_upb_protos(ctx, generator, proto_info, proto_sources):
     if len(proto_sources) == 0:
         return GeneratedSrcsInfo(srcs = [], hdrs = [])
 
+    ext = "." + generator
+    tool = getattr(ctx.executable, "_gen_" + generator)
     srcs = [_generate_output_file(ctx, name, ext + ".c") for name in proto_sources]
     hdrs = [_generate_output_file(ctx, name, ext + ".h") for name in proto_sources]
     transitive_sets = proto_info.transitive_descriptor_sets.to_list()
-    fasttable_enabled = ctx.attr._fasttable_enabled[_FastTableEnabled].enabled
+    fasttable_enabled = (hasattr(ctx.attr, "_fasttable_enabled") and
+                         ctx.attr._fasttable_enabled[_FastTableEnabled].enabled)
     codegen_params = "fasttable:" if fasttable_enabled else ""
     ctx.actions.run(
         inputs = depset(
             direct = [proto_info.direct_descriptor_set],
             transitive = [proto_info.transitive_descriptor_sets],
         ),
-        tools = [ctx.executable._upbc],
+        tools = [tool],
         outputs = srcs + hdrs,
         executable = ctx.executable._protoc,
         arguments = [
-                        "--upb_out=" + codegen_params + _get_real_root(srcs[0]),
-                        "--plugin=protoc-gen-upb=" + ctx.executable._upbc.path,
+                        "--" + generator + "_out=" + codegen_params + _get_real_root(srcs[0]),
+                        "--plugin=protoc-gen-" + generator + "=" + tool.path,
                         "--descriptor_set_in=" + ctx.configuration.host_path_separator.join([f.path for f in transitive_sets]),
                     ] +
                     [_get_real_short_path(file) for file in proto_sources],
@@ -213,22 +216,20 @@ def _upb_proto_rule_impl(ctx):
         cc_info,
     ]
 
-def _upb_proto_aspect_impl(target, ctx, cc_provider, file_provider):
+def _upb_proto_aspect_impl(target, ctx, generator, cc_provider, file_provider):
     proto_info = target[ProtoInfo]
-    files = _compile_upb_protos(ctx, proto_info, proto_info.direct_sources, ctx.attr._ext)
-    deps = ctx.rule.attr.deps + ctx.attr._upb
-    if cc_provider == _UpbDefsWrappedCcInfo:
-        deps += ctx.attr._upb_reflection
+    files = _compile_upb_protos(ctx, generator, proto_info, proto_info.direct_sources)
+    deps = ctx.rule.attr.deps + getattr(ctx.attr, "_" + generator)
     dep_ccinfos = [dep[CcInfo] for dep in deps if CcInfo in dep]
     dep_ccinfos += [dep[_UpbWrappedCcInfo].cc_info for dep in deps if _UpbWrappedCcInfo in dep]
     dep_ccinfos += [dep[_UpbDefsWrappedCcInfo].cc_info for dep in deps if _UpbDefsWrappedCcInfo in dep]
-    if cc_provider == _UpbDefsWrappedCcInfo:
+    if generator == "upbdefs":
         if _UpbWrappedCcInfo not in target:
             fail("Target should have _UpbDefsWrappedCcInfo provider")
         dep_ccinfos += [target[_UpbWrappedCcInfo].cc_info]
     cc_info = _cc_library_func(
         ctx = ctx,
-        name = ctx.rule.attr.name + ctx.attr._ext,
+        name = ctx.rule.attr.name + "." + generator,
         hdrs = files.hdrs,
         srcs = files.srcs,
         copts = ctx.attr._copts[_UpbProtoLibraryCopts].copts,
@@ -237,10 +238,10 @@ def _upb_proto_aspect_impl(target, ctx, cc_provider, file_provider):
     return [cc_provider(cc_info = cc_info), file_provider(srcs = files)]
 
 def _upb_proto_library_aspect_impl(target, ctx):
-    return _upb_proto_aspect_impl(target, ctx, _UpbWrappedCcInfo, _WrappedGeneratedSrcsInfo)
+    return _upb_proto_aspect_impl(target, ctx, "upb", _UpbWrappedCcInfo, _WrappedGeneratedSrcsInfo)
 
 def _upb_proto_reflection_library_aspect_impl(target, ctx):
-    return _upb_proto_aspect_impl(target, ctx, _UpbDefsWrappedCcInfo, _WrappedDefsGeneratedSrcsInfo)
+    return _upb_proto_aspect_impl(target, ctx, "upbdefs", _UpbDefsWrappedCcInfo, _WrappedDefsGeneratedSrcsInfo)
 
 def _maybe_add(d):
     if not _is_bazel:
@@ -258,7 +259,7 @@ _upb_proto_library_aspect = aspect(
         "_copts": attr.label(
             default = "//:upb_proto_library_copts__for_generated_code_only_do_not_use",
         ),
-        "_upbc": attr.label(
+        "_gen_upb": attr.label(
             executable = True,
             cfg = "host",
             default = "//upbc:protoc-gen-upb",
@@ -275,7 +276,6 @@ _upb_proto_library_aspect = aspect(
             "//:generated_code_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
             "//:upb",
         ]),
-        "_ext": attr.string(default = ".upb"),
         "_fasttable_enabled": attr.label(default = "//:fasttable_enabled"),
     }),
     implementation = _upb_proto_library_aspect_impl,
@@ -308,10 +308,10 @@ _upb_proto_reflection_library_aspect = aspect(
         "_copts": attr.label(
             default = "//:upb_proto_library_copts__for_generated_code_only_do_not_use",
         ),
-        "_upbc": attr.label(
+        "_gen_upbdefs": attr.label(
             executable = True,
             cfg = "host",
-            default = "//upbc:protoc-gen-upb",
+            default = "//upbc:protoc-gen-upbdefs",
         ),
         "_protoc": attr.label(
             executable = True,
@@ -321,22 +321,12 @@ _upb_proto_reflection_library_aspect = aspect(
         "_cc_toolchain": attr.label(
             default = "@bazel_tools//tools/cpp:current_cc_toolchain",
         ),
-        # For unknown reasons, this gets overwritten.
-        "_upb": attr.label_list(
+        "_upbdefs": attr.label_list(
             default = [
-                "//:generated_code_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
                 "//:upb",
                 "//:reflection",
             ],
         ),
-        "_upb_reflection": attr.label_list(
-            default = [
-                "//:upb",
-                "//:reflection",
-            ],
-        ),
-        "_ext": attr.string(default = ".upbdefs"),
-        "_fasttable_enabled": attr.label(default = "//:fasttable_enabled"),
     }),
     implementation = _upb_proto_reflection_library_aspect_impl,
     provides = [

upbc/BUILD

@@ -6,15 +6,27 @@ load(
 licenses(["notice"])
 
 cc_library(
-    name = "upbc_generator",
+    name = "common",
+    hdrs = ["common.h"],
+    srcs = ["common.cc"],
+    copts = UPB_DEFAULT_CPPOPTS,
+    deps = [
+        "@com_google_protobuf//:protobuf",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_binary(
+    name = "protoc-gen-upb",
     srcs = [
-        "generator.cc",
+        "protoc-gen-upb.cc",
         "message_layout.cc",
         "message_layout.h",
     ],
-    hdrs = ["generator.h"],
     copts = UPB_DEFAULT_CPPOPTS,
+    visibility = ["//visibility:public"],
     deps = [
+        ":common",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
@@ -24,12 +36,18 @@ cc_library(
 )
 
 cc_binary(
-    name = "protoc-gen-upb",
-    srcs = ["main.cc"],
+    name = "protoc-gen-upbdefs",
+    srcs = [
+        "protoc-gen-upbdefs.cc",
+    ],
     copts = UPB_DEFAULT_CPPOPTS,
     visibility = ["//visibility:public"],
     deps = [
-        ":upbc_generator",
+        ":common",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings",
+        "@com_google_protobuf//:protobuf",
         "@com_google_protobuf//:protoc_lib",
     ],
 )

upbc/common.cc

@@ -0,0 +1,65 @@
+
+#include "absl/strings/str_replace.h"
+#include "upbc/common.h"
+
+namespace upbc {
+namespace {
+
+namespace protobuf = ::google::protobuf;
+
+void AddMessages(const protobuf::Descriptor* message,
+                 std::vector<const protobuf::Descriptor*>* messages) {
+  messages->push_back(message);
+  for (int i = 0; i < message->nested_type_count(); i++) {
+    AddMessages(message->nested_type(i), messages);
+  }
+}
+
+}  // namespace
+
+std::string StripExtension(absl::string_view fname) {
+  size_t lastdot = fname.find_last_of(".");
+  if (lastdot == std::string::npos) {
+    return std::string(fname);
+  }
+  return std::string(fname.substr(0, lastdot));
+}
+
+std::string ToCIdent(absl::string_view str) {
+  return absl::StrReplaceAll(str, {{".", "_"}, {"/", "_"}});
+}
+
+std::string ToPreproc(absl::string_view str) {
+  return absl::AsciiStrToUpper(ToCIdent(str));
+}
+
+void EmitFileWarning(const protobuf::FileDescriptor* file, Output& output) {
+  output(
+      "/* This file was generated by upbc (the upb compiler) from the input\n"
+      " * file:\n"
+      " *\n"
+      " *     $0\n"
+      " *\n"
+      " * Do not edit -- your changes will be discarded when the file is\n"
+      " * regenerated. */\n\n",
+      file->name());
+}
+
+std::vector<const protobuf::Descriptor*> SortedMessages(
+    const protobuf::FileDescriptor* file) {
+  std::vector<const protobuf::Descriptor*> messages;
+  for (int i = 0; i < file->message_type_count(); i++) {
+    AddMessages(file->message_type(i), &messages);
+  }
+  return messages;
+}
+
+std::string MessageName(const protobuf::Descriptor* descriptor) {
+  return ToCIdent(descriptor->full_name());
+}
+
+std::string MessageInit(const protobuf::Descriptor* descriptor) {
+  return MessageName(descriptor) + "_msginit";
+}
+
+}  // namespace upbc

upbc/common.h

@@ -0,0 +1,66 @@
+
+#ifndef UPBC_COMMON_H
+#define UPBC_COMMON_H
+
+#include <vector>
+
+#include "absl/strings/substitute.h"
+#include "google/protobuf/descriptor.h"
+#include "google/protobuf/io/zero_copy_stream.h"
+
+namespace upbc {
+
+class Output {
+ public:
+  Output(google::protobuf::io::ZeroCopyOutputStream* stream)
+      : stream_(stream) {}
+  ~Output() { stream_->BackUp((int)size_); }
+
+  template <class... Arg>
+  void operator()(absl::string_view format, const Arg&... arg) {
+    Write(absl::Substitute(format, arg...));
+  }
+
+ private:
+  void Write(absl::string_view data) {
+    while (!data.empty()) {
+      RefreshOutput();
+      size_t to_write = std::min(data.size(), size_);
+      memcpy(ptr_, data.data(), to_write);
+      data.remove_prefix(to_write);
+      ptr_ += to_write;
+      size_ -= to_write;
+    }
+  }
+
+  void RefreshOutput() {
+    while (size_ == 0) {
+      void *ptr;
+      int size;
+      if (!stream_->Next(&ptr, &size)) {
+        fprintf(stderr, "upbc: Failed to write to to output\n");
+        abort();
+      }
+      ptr_ = static_cast<char*>(ptr);
+      size_ = size;
+    }
+  }
+
+  google::protobuf::io::ZeroCopyOutputStream* stream_;
+  char *ptr_ = nullptr;
+  size_t size_ = 0;
+};
+
+std::string StripExtension(absl::string_view fname);
+std::string ToCIdent(absl::string_view str);
+std::string ToPreproc(absl::string_view str);
+void EmitFileWarning(const google::protobuf::FileDescriptor* file,
+                     Output& output);
+std::vector<const google::protobuf::Descriptor*> SortedMessages(
+    const google::protobuf::FileDescriptor* file);
+std::string MessageInit(const google::protobuf::Descriptor* descriptor);
+std::string MessageName(const google::protobuf::Descriptor* descriptor);
+
+}  // namespace upbc
+
+# endif  // UPBC_COMMON_H