Move to proto_common for all upb aspects to fix numerous tricky edge cases and simplify the code

PiperOrigin-RevId: 527904449

.github/workflows/bazel_tests.yml

@@ -30,8 +30,8 @@ jobs:
           - { NAME: "macOS", BAZEL: bazel, CC: clang, os: macos-11 }
           - { NAME: "Windows", BAZEL: bazel, os: windows-2019, startup-flags: "--output_user_root=C:/tmp", flags: "--config=cpp17_msvc", targets: "upb/... upbc/... python/... protos/... protos_generator/..." }
           # We support two Bazel versions back per https://opensource.google/documentation/policies/cplusplus-support
-          - { NAME: "Bazel 4.1.0", BAZEL: bazel-4.1.0-linux-x86_64, CC: clang, os: ubuntu-20-large }
           - { NAME: "Bazel 5.3.0", BAZEL: bazel-5.3.0-linux-x86_64, CC: clang, os: ubuntu-20-large }
+          - { NAME: "Bazel 6.1.1", BAZEL: bazel-6.1.1-linux-x86_64, CC: clang, os: ubuntu-20-large }
 
     name: ${{ matrix.NAME }}
 
@@ -52,10 +52,10 @@ jobs:
              wget -O $FILENAME https://github.com/bazelbuild/bazel/releases/download/$VERSION/${{ matrix.BAZEL }}
              chmod a+x $FILENAME
         if: ${{ matrix.BAZEL != 'bazel' }}
-      - name: Check compiler versions
+      - name: Check compiler version
         if: matrix.CC
         run: ${{ matrix.CC }} --version
-      - name: Check Bazel versions
+      - name: Check Bazel version
         run: ${{ matrix.BAZEL }} --version
       - id: bazel-cache
         name: Set up Bazel caching

.github/workflows/python_tests.yml

@@ -29,6 +29,9 @@ jobs:
           export_environment_variables: true
       - name: Use gcloud CLI
         run: gcloud info
+      - name: Check compiler version
+        if: matrix.CC
+        run: ${{ matrix.CC }} --version
       - name: Configure Docker
         run: gcloud auth configure-docker -q us-docker.pkg.dev
       - name: Pull Docker Image

WORKSPACE

@@ -56,13 +56,15 @@ http_archive(
 )
 
 load("@com_google_googletest//:googletest_deps.bzl", "googletest_deps")
-
 googletest_deps()
 
 load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
-
 rules_pkg_dependencies()
 
+load("@rules_proto//proto:repositories.bzl", "rules_proto_dependencies", "rules_proto_toolchains")
+rules_proto_dependencies()
+rules_proto_toolchains()
+
 load("//bazel:system_python.bzl", "system_python")
 system_python(
     name = "system_python",

bazel/BUILD

@@ -46,13 +46,22 @@ bzl_library(
     srcs = ["py_proto_library.bzl"],
 )
 
+bzl_library(
+    name = "proto_common_wrapper_bzl",
+    srcs = ["proto_common_wrapper.bzl"],
+    visibility = ["//protos/bazel:__pkg__"],
+    deps = [
+        "@rules_proto//proto:defs",
+    ],
+)
+
 bzl_library(
     name = "upb_proto_library_bzl",
     srcs = ["upb_proto_library.bzl"],
     visibility = ["//visibility:public"],
     deps = [
-        "@bazel_skylib//lib:paths",
-        "@bazel_tools//tools/cpp:toolchain_utils.bzl",
+        ":proto_common_wrapper_bzl",
         "@rules_proto//proto:defs",
+        "@bazel_tools//tools/cpp:toolchain_utils.bzl",
     ],
 )

bazel/proto_common_wrapper.bzl

@@ -0,0 +1,68 @@
+# Copyright (c) 2023, Google LLC
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#     * Redistributions of source code must retain the above copyright
+#       notice, this list of conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above copyright
+#       notice, this list of conditions and the following disclaimer in the
+#       documentation and/or other materials provided with the distribution.
+#     * Neither the name of Google LLC nor the
+#       names of its contributors may be used to endorse or promote products
+#       derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL Google LLC BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""A wrapper around proto_common.compile() that fixes some surprising behaviors.
+
+More info in: https://github.com/bazelbuild/bazel/issues/18263
+"""
+
+load("@rules_proto//proto:defs.bzl", "proto_common")
+
+def output_dir(ctx, proto_info):
+    """Returns the output directory where generated proto files will be placed.
+
+    Args:
+      ctx: Rule context.
+      proto_info: ProtoInfo provider.
+
+    Returns:
+      A string specifying the output directory
+    """
+    proto_root = proto_info.proto_source_root
+    if proto_root.startswith(ctx.bin_dir.path):
+        path = proto_root
+    else:
+        path = ctx.bin_dir.path + "/" + proto_root
+
+    if proto_root == ".":
+        path = ctx.bin_dir.path
+    return path
+
+def proto_common_compile(ctx, proto_info, proto_lang_toolchain_info, generated_files):
+    """A wrapper around proto_common.compile that automatically calculates the output dir.
+
+    Args:
+      ctx: Rule context.
+      proto_info: ProtoInfo provider.
+      proto_lang_toolchain_info: ProtoLangToolchainInfo provider.
+      generated_files: The files we expect to be generated from this protoc invocation.
+    """
+    proto_common.compile(
+        actions = ctx.actions,
+        proto_info = proto_info,
+        proto_lang_toolchain_info = proto_lang_toolchain_info,
+        generated_files = generated_files,
+        plugin_output = output_dir(ctx, proto_info),
+    )

bazel/upb_proto_library.bzl

@@ -28,19 +28,9 @@
   - upb_proto_reflection_library()
 """
 
-load("@bazel_skylib//lib:paths.bzl", "paths")
-
-# begin:google_only
-# load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain", "use_cpp_toolchain")
-# end:google_only
-
-# begin:github_only
-# Compatibility code for Bazel 4.x. Remove this when we drop support for Bazel 4.x.
-load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain")
-
-def use_cpp_toolchain():
-    return ["@bazel_tools//tools/cpp:toolchain_type"]
-# end:github_only
+load(":proto_common_wrapper.bzl", "output_dir", "proto_common_compile")
+load("@rules_proto//proto:defs.bzl", "proto_common")
+load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain", "use_cpp_toolchain")
 
 # Generic support code #########################################################
 
@@ -52,61 +42,6 @@ _is_google3 = False
 # _is_google3 = True
 # end:google_only
 
-def _get_real_short_path(file):
-    # For some reason, files from other archives have short paths that look like:
-    #   ../com_google_protobuf/google/protobuf/descriptor.proto
-    short_path = file.short_path
-    if short_path.startswith("../"):
-        second_slash = short_path.index("/", 3)
-        short_path = short_path[second_slash + 1:]
-
-    # Sometimes it has another few prefixes like:
-    #   _virtual_imports/any_proto/google/protobuf/any.proto
-    #   benchmarks/_virtual_imports/100_msgs_proto/benchmarks/100_msgs.proto
-    # We want just google/protobuf/any.proto.
-    virtual_imports = "_virtual_imports/"
-    if virtual_imports in short_path:
-        short_path = short_path.split(virtual_imports)[1].split("/", 1)[1]
-    return short_path
-
-def _get_real_root(ctx, file):
-    real_short_path = _get_real_short_path(file)
-    root = file.path[:-len(real_short_path) - 1]
-
-    if not _is_google3 and ctx.rule.attr.strip_import_prefix:
-        root = paths.join(root, ctx.rule.attr.strip_import_prefix[1:])
-    return root
-
-def _generate_output_file(ctx, src, extension):
-    package = ctx.label.package
-    if not _is_google3:
-        strip_import_prefix = ctx.rule.attr.strip_import_prefix
-        if strip_import_prefix and strip_import_prefix != "/":
-            if not package.startswith(strip_import_prefix[1:]):
-                fail("%s does not begin with prefix %s" % (package, strip_import_prefix))
-            package = package[len(strip_import_prefix):]
-
-    real_short_path = _get_real_short_path(src)
-    real_short_path = paths.relativize(real_short_path, package)
-    output_filename = paths.replace_extension(real_short_path, extension)
-    ret = ctx.actions.declare_file(output_filename)
-    return ret
-
-def _generate_include_path(src, out, extension):
-    short_path = _get_real_short_path(src)
-    short_path = paths.replace_extension(short_path, extension)
-    if not out.path.endswith(short_path):
-        fail("%s does not end with %s" % (out.path, short_path))
-
-    return out.path[:-len(short_path)]
-
-def _filter_none(elems):
-    out = []
-    for elem in elems:
-        if elem:
-            out.append(elem)
-    return out
-
 def _cc_library_func(ctx, name, hdrs, srcs, copts, includes, dep_ccinfos):
     """Like cc_library(), but callable from rules.
 
@@ -205,57 +140,12 @@ _WrappedDefsGeneratedSrcsInfo = provider(
     fields = ["srcs"],
 )
 
-def _compile_upb_protos(ctx, generator, proto_info, proto_sources):
-    if len(proto_sources) == 0:
-        return GeneratedSrcsInfo(srcs = [], hdrs = [], thunks = [], includes = [])
-
-    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]
-    thunks = []
-    if generator == "upb":
-        thunks = [_generate_output_file(ctx, name, ext + ".thunks.c") for name in proto_sources]
-    transitive_sets = proto_info.transitive_descriptor_sets.to_list()
-
-    args = ctx.actions.args()
-    args.use_param_file(param_file_arg = "@%s")
-    args.set_param_file_format("multiline")
-
-    args.add("--" + generator + "_out=" + _get_real_root(ctx, srcs[0]))
-    args.add("--plugin=protoc-gen-" + generator + "=" + tool.path)
-    args.add("--descriptor_set_in=" + ctx.configuration.host_path_separator.join([f.path for f in transitive_sets]))
-    args.add_all(proto_sources, map_each = _get_real_short_path)
-
-    ctx.actions.run(
-        inputs = depset(
-            direct = [proto_info.direct_descriptor_set],
-            transitive = [proto_info.transitive_descriptor_sets],
-        ),
-        tools = [tool],
-        outputs = srcs + hdrs,
-        executable = ctx.executable._protoc,
-        arguments = [args],
-        progress_message = "Generating upb protos for :" + ctx.label.name,
-        mnemonic = "GenUpbProtos",
-    )
-    if generator == "upb":
-        ctx.actions.run_shell(
-            inputs = hdrs,
-            outputs = thunks,
-            command = " && ".join([
-                "sed 's/UPB_INLINE //' {} > {}".format(hdr.path, thunk.path)
-                for (hdr, thunk) in zip(hdrs, thunks)
-            ]),
-            progress_message = "Generating thunks for upb protos API for: " + ctx.label.name,
-            mnemonic = "GenUpbProtosThunks",
-        )
-    return GeneratedSrcsInfo(
-        srcs = srcs,
-        hdrs = hdrs,
-        thunks = thunks,
-        includes = [_generate_include_path(proto_sources[0], hdrs[0], ext + ".h")],
-    )
+def _filter_none(elems):
+    out = []
+    for elem in elems:
+        if elem:
+            out.append(elem)
+    return out
 
 def _upb_proto_rule_impl(ctx):
     if len(ctx.attr.deps) != 1:
@@ -290,10 +180,64 @@ def _upb_proto_rule_impl(ctx):
         cc_info,
     ]
 
+def _get_lang_toolchain(ctx, generator):
+    lang_toolchain_name = "_" + generator + "_toolchain"
+    return getattr(ctx.attr, lang_toolchain_name)[proto_common.ProtoLangToolchainInfo]
+
+def _compile_upb_protos(ctx, generator, proto_info):
+    proto_sources = proto_info.direct_sources
+    if len(proto_sources) == 0:
+        return [], []
+
+    srcs = proto_common.declare_generated_files(
+        ctx.actions,
+        extension = "." + generator + ".c",
+        proto_info = proto_info,
+    )
+    hdrs = proto_common.declare_generated_files(
+        ctx.actions,
+        extension = "." + generator + ".h",
+        proto_info = proto_info,
+    )
+
+    proto_common_compile(
+        ctx = ctx,
+        proto_info = proto_info,
+        proto_lang_toolchain_info = _get_lang_toolchain(ctx, generator),
+        generated_files = srcs + hdrs,
+    )
+
+    if generator == "upb":
+        thunks = proto_common.declare_generated_files(
+            ctx.actions,
+            extension = "." + generator + ".thunks.c",
+            proto_info = proto_info,
+        )
+        ctx.actions.run_shell(
+            inputs = hdrs,
+            outputs = thunks,
+            command = " && ".join([
+                "sed 's/UPB_INLINE //' {} > {}".format(hdr.path, thunk.path)
+                for (hdr, thunk) in zip(hdrs, thunks)
+            ]),
+            progress_message = "Generating thunks for upb protos API for: " + ctx.label.name,
+            mnemonic = "GenUpbProtosThunks",
+        )
+    else:
+        thunks = []
+
+    return srcs, hdrs, thunks
+
 def _upb_proto_aspect_impl(target, ctx, generator, cc_provider, file_provider):
     proto_info = target[ProtoInfo]
-    files = _compile_upb_protos(ctx, generator, proto_info, proto_info.direct_sources)
-    deps = ctx.rule.attr.deps + getattr(ctx.attr, "_" + generator)
+    srcs, hdrs, thunks = _compile_upb_protos(ctx, generator, proto_info)
+    files = GeneratedSrcsInfo(
+        srcs = srcs,
+        hdrs = hdrs,
+        thunks = thunks,
+    )
+    runtime = _get_lang_toolchain(ctx, generator).runtime
+    deps = ctx.rule.attr.deps + [runtime]
     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]
@@ -306,7 +250,7 @@ def _upb_proto_aspect_impl(target, ctx, generator, cc_provider, file_provider):
         name = ctx.rule.attr.name + "." + generator,
         hdrs = files.hdrs,
         srcs = files.srcs,
-        includes = files.includes,
+        includes = [output_dir(ctx, proto_info)],
         copts = ctx.attr._copts[UpbProtoLibraryCoptsInfo].copts,
         dep_ccinfos = dep_ccinfos,
     )
@@ -317,7 +261,7 @@ def _upb_proto_aspect_impl(target, ctx, generator, cc_provider, file_provider):
             name = ctx.rule.attr.name + "." + generator + ".thunks",
             hdrs = [],
             srcs = files.thunks,
-            includes = files.includes,
+            includes = [output_dir(ctx, proto_info)],
             copts = ctx.attr._copts[UpbProtoLibraryCoptsInfo].copts,
             dep_ccinfos = dep_ccinfos + [cc_info],
         )
@@ -347,6 +291,10 @@ def _maybe_add(d):
             cfg = "exec",
             default = "@bazel_tools//tools/cpp:grep-includes",
         )
+    else:
+        d["_cc_toolchain"] = attr.label(
+            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
+        )
     return d
 
 # upb_proto_library() ##########################################################
@@ -356,23 +304,10 @@ upb_proto_library_aspect = aspect(
         "_copts": attr.label(
             default = "//:upb_proto_library_copts__for_generated_code_only_do_not_use",
         ),
-        "_gen_upb": attr.label(
-            executable = True,
-            cfg = "exec",
-            default = "//upbc:protoc-gen-upb_stage1",
+        "_upb_toolchain": attr.label(
+            default = Label("//upbc:protoc-gen-upb_toolchain"),
+            cfg = getattr(proto_common, "proto_lang_toolchain_cfg", "target"),
         ),
-        "_protoc": attr.label(
-            executable = True,
-            cfg = "exec",
-            default = "@com_google_protobuf//:protoc",
-        ),
-        "_cc_toolchain": attr.label(
-            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
-        ),
-        "_upb": attr.label_list(default = [
-            "//:generated_code_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
-        ]),
-        "_fasttable_enabled": attr.label(default = "//:fasttable_enabled"),
     }),
     implementation = upb_proto_library_aspect_impl,
     provides = [
@@ -386,7 +321,6 @@ upb_proto_library_aspect = aspect(
 )
 
 upb_proto_library = rule(
-    output_to_genfiles = True,
     implementation = _upb_proto_rule_impl,
     attrs = {
         "deps": attr.label_list(
@@ -404,23 +338,9 @@ _upb_proto_reflection_library_aspect = aspect(
         "_copts": attr.label(
             default = "//:upb_proto_library_copts__for_generated_code_only_do_not_use",
         ),
-        "_gen_upbdefs": attr.label(
-            executable = True,
-            cfg = "exec",
-            default = "//upbc:protoc-gen-upbdefs",
-        ),
-        "_protoc": attr.label(
-            executable = True,
-            cfg = "exec",
-            default = "@com_google_protobuf//:protoc",
-        ),
-        "_cc_toolchain": attr.label(
-            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
-        ),
-        "_upbdefs": attr.label_list(
-            default = [
-                "//:generated_reflection_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
-            ],
+        "_upbdefs_toolchain": attr.label(
+            default = Label("//upbc:protoc-gen-upbdefs_toolchain"),
+            cfg = getattr(proto_common, "proto_lang_toolchain_cfg", "target"),
         ),
     }),
     implementation = _upb_proto_reflection_library_aspect_impl,
@@ -439,7 +359,6 @@ _upb_proto_reflection_library_aspect = aspect(
 )
 
 upb_proto_reflection_library = rule(
-    output_to_genfiles = True,
     implementation = _upb_proto_rule_impl,
     attrs = {
         "deps": attr.label_list(

bazel/workspace_deps.bzl

@@ -46,6 +46,16 @@ def upb_deps():
         sha256 = "8a298e832762eda1830597d64fe7db58178aa84cd5926d76d5b744d6558941c2",
     )
 
+    maybe(
+        http_archive,
+        name = "rules_proto",
+        sha256 = "dc3fb206a2cb3441b485eb1e423165b231235a1ea9b031b4433cf7bc1fa460dd",
+        strip_prefix = "rules_proto-5.3.0-21.7",
+        urls = [
+            "https://github.com/bazelbuild/rules_proto/archive/refs/tags/5.3.0-21.7.tar.gz",
+        ],
+    )
+
     maybe(
         _github_archive,
         name = "rules_python",

cmake/make_cmakelists.py

@@ -280,6 +280,12 @@ class WorkspaceFileFunctions(object):
   def googletest_deps(self):
     pass
 
+  def rules_proto_dependencies(self):
+    pass
+
+  def rules_proto_toolchains(self):
+    pass
+
 
 class Converter(object):
   def __init__(self):

protos/BUILD

@@ -68,6 +68,7 @@ cc_library(
 cc_library(
     name = "generated_protos_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
     hdrs = [
+        "protos.h",
         "protos_internal.h",
     ],
     copts = UPB_DEFAULT_CPPOPTS,
@@ -75,6 +76,10 @@ cc_library(
     deps = [
         ":protos",
         ":protos_internal",
+        "//:mini_table_internal",
+        "//:upb",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
     ],
 )
 

protos/bazel/BUILD

@@ -32,7 +32,7 @@ bzl_library(
     srcs = ["upb_cc_proto_library.bzl"],
     visibility = ["//visibility:public"],
     deps = [
-        "@bazel_skylib//lib:paths",
+        "//bazel:proto_common_wrapper_bzl",
         "//bazel:upb_proto_library_bzl",
         "@bazel_tools//tools/cpp:toolchain_utils.bzl",
     ],

protos/bazel/upb_cc_proto_library.bzl

@@ -27,20 +27,10 @@
   - upb_cc_proto_library()
 """
 
-load("@bazel_skylib//lib:paths.bzl", "paths")
+load("//bazel:proto_common_wrapper.bzl", "proto_common_compile")
 load("//bazel:upb_proto_library.bzl", "GeneratedSrcsInfo", "UpbWrappedCcInfo", "upb_proto_library_aspect")
-
-# begin:google_only
-# load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain", "use_cpp_toolchain")
-#
-# end:google_only
-# begin:github_only
-# Compatibility code for Bazel 4.x. Remove this when we drop support for Bazel 4.x.
-load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain")
-
-def use_cpp_toolchain():
-    return ["@bazel_tools//tools/cpp:toolchain_type"]
-# end:github_only
+load("@rules_proto//proto:defs.bzl", "proto_common")
+load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain", "use_cpp_toolchain")
 
 # Generic support code #########################################################
 
@@ -52,34 +42,6 @@ _is_google3 = False
 # _is_google3 = True
 # end:google_only
 
-def _get_real_short_path(file):
-    # For some reason, files from other archives have short paths that look like:
-    #   ../com_google_protobuf/google/protobuf/descriptor.proto
-    short_path = file.short_path
-    if short_path.startswith("../"):
-        second_slash = short_path.index("/", 3)
-        short_path = short_path[second_slash + 1:]
-
-    # Sometimes it has another few prefixes like:
-    #   _virtual_imports/any_proto/google/protobuf/any.proto
-    #   benchmarks/_virtual_imports/100_msgs_proto/benchmarks/100_msgs.proto
-    # We want just google/protobuf/any.proto.
-    virtual_imports = "_virtual_imports/"
-    if virtual_imports in short_path:
-        short_path = short_path.split(virtual_imports)[1].split("/", 1)[1]
-    return short_path
-
-def _get_real_root(file):
-    real_short_path = _get_real_short_path(file)
-    return file.path[:-len(real_short_path) - 1]
-
-def _generate_output_file(ctx, src, extension):
-    real_short_path = _get_real_short_path(src)
-    real_short_path = paths.relativize(real_short_path, ctx.label.package)
-    output_filename = paths.replace_extension(real_short_path, extension)
-    ret = ctx.actions.declare_file(output_filename)
-    return ret
-
 def _filter_none(elems):
     out = []
     for elem in elems:
@@ -165,36 +127,37 @@ upb_cc_proto_library_copts = rule(
 _UpbCcWrappedCcInfo = provider("Provider for cc_info for protos", fields = ["cc_info"])
 _WrappedCcGeneratedSrcsInfo = provider("Provider for generated sources", fields = ["srcs"])
 
+def _get_lang_toolchain(ctx, generator):
+    lang_toolchain_name = "_" + generator + "_toolchain"
+    return getattr(ctx.attr, lang_toolchain_name)[proto_common.ProtoLangToolchainInfo]
+
 def _compile_upb_cc_protos(ctx, generator, proto_info, proto_sources):
     if len(proto_sources) == 0:
         return GeneratedSrcsInfo(srcs = [], hdrs = [])
 
-    tool = getattr(ctx.executable, "_gen_" + generator)
-    srcs = [_generate_output_file(ctx, name, ".upb.proto.cc") for name in proto_sources]
-    hdrs = [_generate_output_file(ctx, name, ".upb.proto.h") for name in proto_sources]
-    hdrs += [_generate_output_file(ctx, name, ".upb.fwd.h") for name in proto_sources]
-    transitive_sets = proto_info.transitive_descriptor_sets.to_list()
-
-    args = ctx.actions.args()
-    args.use_param_file(param_file_arg = "@%s")
-    args.set_param_file_format("multiline")
-
-    args.add("--" + generator + "_out=" + _get_real_root(srcs[0]))
-    args.add("--plugin=protoc-gen-" + generator + "=" + tool.path)
-    args.add("--descriptor_set_in=" + ctx.configuration.host_path_separator.join([f.path for f in transitive_sets]))
-    args.add_all(proto_sources, map_each = _get_real_short_path)
-
-    ctx.actions.run(
-        inputs = depset(
-            direct = [proto_info.direct_descriptor_set],
-            transitive = [proto_info.transitive_descriptor_sets],
-        ),
-        tools = [tool],
-        outputs = srcs + hdrs,
-        executable = ctx.executable._protoc,
-        arguments = [args],
-        progress_message = "Generating upb cc protos for :" + ctx.label.name,
+    srcs = proto_common.declare_generated_files(
+        ctx.actions,
+        extension = ".upb.proto.cc",
+        proto_info = proto_info,
+    )
+    hdrs = proto_common.declare_generated_files(
+        ctx.actions,
+        extension = ".upb.proto.h",
+        proto_info = proto_info,
+    )
+    hdrs += proto_common.declare_generated_files(
+        ctx.actions,
+        extension = ".upb.fwd.h",
+        proto_info = proto_info,
     )
+
+    proto_common_compile(
+        ctx = ctx,
+        proto_info = proto_info,
+        proto_lang_toolchain_info = _get_lang_toolchain(ctx, generator),
+        generated_files = srcs + hdrs,
+    )
+
     return GeneratedSrcsInfo(srcs = srcs, hdrs = hdrs)
 
 def _upb_cc_proto_rule_impl(ctx):
@@ -231,7 +194,8 @@ def _upb_cc_proto_rule_impl(ctx):
 def _upb_cc_proto_aspect_impl(target, ctx, generator, cc_provider, file_provider):
     proto_info = target[ProtoInfo]
     files = _compile_upb_cc_protos(ctx, generator, proto_info, proto_info.direct_sources)
-    deps = ctx.rule.attr.deps + getattr(ctx.attr, "_" + generator)
+    runtime = _get_lang_toolchain(ctx, generator).runtime
+    deps = ctx.rule.attr.deps + [runtime] + ctx.attr._aspect_deps
     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[_UpbCcWrappedCcInfo].cc_info for dep in deps if _UpbCcWrappedCcInfo in dep]
@@ -258,6 +222,10 @@ def _maybe_add(d):
             cfg = "exec",
             default = "@bazel_tools//tools/cpp:grep-includes",
         )
+    else:
+        d["_cc_toolchain"] = attr.label(
+            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
+        )
     return d
 
 _upb_cc_proto_library_aspect = aspect(
@@ -265,22 +233,14 @@ _upb_cc_proto_library_aspect = aspect(
         "_ccopts": attr.label(
             default = "//protos:upb_cc_proto_library_copts__for_generated_code_only_do_not_use",
         ),
-        "_gen_upbprotos": attr.label(
-            executable = True,
-            cfg = "exec",
-            default = "//protos_generator:protoc-gen-upb-protos",
-        ),
-        "_protoc": attr.label(
-            executable = True,
-            cfg = "exec",
-            default = "@com_google_protobuf//:protoc",
-        ),
-        "_cc_toolchain": attr.label(
-            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
+        "_upbprotos_toolchain": attr.label(
+            default = Label("//protos_generator:protoc-gen-upbprotos_toolchain"),
+            cfg = getattr(proto_common, "proto_lang_toolchain_cfg", "target"),
         ),
-        "_upbprotos": attr.label_list(
+        # proto_lang_toolchain(runtime="//xyz") only allows a single "runtime", so we need to add
+        # extra deps here.
+        "_aspect_deps": attr.label_list(
             default = [
-                # TODO: Add dependencies for cc runtime (absl/string etc..)
                 "//:generated_cpp_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
                 "//protos:generated_protos_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
                 "@com_google_absl//absl/strings",

protos_generator/BUILD

@@ -81,3 +81,13 @@ cc_library(
         "@com_google_protobuf//src/google/protobuf/compiler:code_generator",
     ],
 )
+
+proto_lang_toolchain(
+    name = "protoc-gen-upbprotos_toolchain",
+    command_line = "--upbprotos_out=%s",
+    plugin = ":protoc-gen-upb-protos",
+    plugin_format_flag = "--plugin=protoc-gen-upbprotos=%s",
+    progress_message = "Generating upb C++ protos",
+    runtime = "//protos:generated_protos_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
+    visibility = ["//visibility:public"],
+)

protos_generator/protoc-gen-upb-protos.cc

@@ -129,11 +129,6 @@ void WriteHeader(const protobuf::FileDescriptor* file, Output& output) {
 
 #include "protos/protos.h"
 #include "protos/protos_internal.h"
-#include "upb/upb.hpp"
-
-#include "absl/strings/string_view.h"
-#include "absl/status/statusor.h"
-#include "upb/message/internal.h"
 #include "upb/message/copy.h"
       )cc",
       ToPreproc(file->name()));
@@ -149,8 +144,6 @@ void WriteHeader(const protobuf::FileDescriptor* file, Output& output) {
     }
   }
 
-  output("#include \"upb/port/def.inc\"\n");
-
   const std::vector<const protobuf::Descriptor*> this_file_messages =
       SortedMessages(file);
   const std::vector<const protobuf::FieldDescriptor*> this_file_exts =
@@ -182,7 +175,6 @@ void WriteHeader(const protobuf::FileDescriptor* file, Output& output) {
 
   WriteEndNamespace(file, output);
 
-  output("\n#include \"upb/port/undef.inc\"\n\n");
   // End of "C" section.
 
   output("#endif  /* $0_UPB_PROTO_H_ */\n", ToPreproc(file->name()));
@@ -196,8 +188,6 @@ void WriteSource(const protobuf::FileDescriptor* file, Output& output,
   output(
       R"cc(
 #include <stddef.h>
-#include "absl/strings/string_view.h"
-#include "upb/message/copy.h"
 #include "upb/message/internal.h"
 #include "protos/protos.h"
 #include "$0"
@@ -207,7 +197,6 @@ void WriteSource(const protobuf::FileDescriptor* file, Output& output,
   for (int i = 0; i < file->dependency_count(); i++) {
     output("#include \"$0\"\n", CppHeaderFilename(file->dependency(i)));
   }
-  output("#include \"upb/port/def.inc\"\n");
 
   WriteStartNamespace(file, output);
   WriteMessageImplementations(file, output);
@@ -215,8 +204,6 @@ void WriteSource(const protobuf::FileDescriptor* file, Output& output,
       SortedExtensions(file);
   WriteExtensionIdentifiers(this_file_exts, output);
   WriteEndNamespace(file, output);
-
-  output("#include \"upb/port/undef.inc\"\n\n");
 }
 
 void WriteMessageImplementations(const protobuf::FileDescriptor* file,

upbc/BUILD

@@ -255,6 +255,16 @@ bootstrap_cc_binary(
     ],
 )
 
+proto_lang_toolchain(
+    name = "protoc-gen-upb_toolchain",
+    command_line = "--upb_out=%s",
+    plugin = ":protoc-gen-upb_stage1",
+    plugin_format_flag = "--plugin=protoc-gen-upb=%s",
+    progress_message = "Generating upb protos",
+    runtime = "//:generated_code_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
+    visibility = ["//visibility:public"],
+)
+
 cc_binary(
     name = "protoc-gen-upbdefs",
     srcs = [
@@ -277,6 +287,16 @@ cc_binary(
     ],
 )
 
+proto_lang_toolchain(
+    name = "protoc-gen-upbdefs_toolchain",
+    command_line = "--upbdefs_out=%s",
+    plugin = ":protoc-gen-upbdefs",
+    plugin_format_flag = "--plugin=protoc-gen-upbdefs=%s",
+    progress_message = "Generating upb protos",
+    runtime = "//:generated_reflection_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
+    visibility = ["//visibility:public"],
+)
+
 cc_binary(
     name = "protoc-gen-upbdev",
     srcs = [