Update Rust generator to output minidescriptors instead of C minitables

With this change, we no longer emit any C generated code for the Rust upb
implementation. Instead, we output minidescriptors in the Rust generated code,
and use these to construct the minitables at run time. Each message's
`AssociatedMiniTable::mini_table()` method will lazily build the minitable on
first use.

We need a place to put the minidescriptor and lazily-built minitable for enums,
so I added a new `AssociatedMiniTableEnum` trait for this purpose. This is
only necessary for closed enums.

There are a couple of things that were somewhat difficult about building the
minitables at runtime:
 - Each map entry is represented as its own implicit message type and needs a
   minitable. To facilitate this, I updated the generator to output a
   visibility-controlled struct to represent each map entry. This way we can
   implement `AssociatedMiniTable` on it and treat it like any other message
   for the purpose of building minitables.
 - The naive way of recursively building minitables is vulnerable to deadlock
   in the case of cycles in the message graph. Such cycles are allowed within a
   single .proto file, so we need to be able to handle this. My change handles
   this problem by treating each strongly-connected component (SCC) as its own
   unit, and then the graph of SCCs has no cycles. Each SCC has an arbitrariliy
   chosen representative responsible for linking all messages in the SCC.

PiperOrigin-RevId: 761654774

Author: acozzette

rust/bazel/aspects.bzl

@@ -7,7 +7,6 @@ load("@rules_rust//rust/private:providers.bzl", "CrateInfo", "DepInfo", "DepVari
 
 # buildifier: disable=bzl-visibility
 load("@rules_rust//rust/private:rustc.bzl", "rustc_compile_action")
-load("//bazel:upb_minitable_proto_library.bzl", "UpbMinitableCcInfo", "upb_minitable_proto_library_aspect")
 load("//bazel/common:proto_common.bzl", "proto_common")
 load("//bazel/common:proto_info.bzl", "ProtoInfo")
 load("//bazel/private:cc_proto_aspect.bzl", "cc_proto_aspect")
@@ -347,9 +346,8 @@ def _rust_proto_aspect_common(target, ctx, is_upb):
         is_upb,
     )
 
-    if is_upb:
-        thunks_cc_info = target[UpbMinitableCcInfo].cc_info
-    else:
+    dep_variant_info_for_native_gencode = []
+    if not is_upb:
         dep_cc_infos = []
         for dep in proto_deps:
             dep_cc_infos.append(dep[CcInfo])
@@ -363,6 +361,8 @@ def _rust_proto_aspect_common(target, ctx, is_upb):
             cc_infos = [target[CcInfo]] + [dep[CcInfo] for dep in ctx.attr._cpp_thunks_deps] + dep_cc_infos,
         ) for thunk in cc_thunks_gencode])
 
+        dep_variant_info_for_native_gencode = [DepVariantInfo(cc_info = thunks_cc_info)]
+
     runtime = proto_lang_toolchain.runtime
     dep_variant_info_for_runtime = DepVariantInfo(
         crate_info = runtime[CrateInfo] if CrateInfo in runtime else None,
@@ -371,14 +371,12 @@ def _rust_proto_aspect_common(target, ctx, is_upb):
         build_info = None,
     )
 
-    dep_variant_info_for_native_gencode = DepVariantInfo(cc_info = thunks_cc_info)
-
     dep_variant_info = _compile_rust(
         ctx = ctx,
         attr = ctx.rule.attr,
         src = entry_point_rs_output,
         extra_srcs = rs_gencode,
-        deps = [dep_variant_info_for_runtime, dep_variant_info_for_native_gencode] + dep_variant_infos,
+        deps = [dep_variant_info_for_runtime] + dep_variant_info_for_native_gencode + dep_variant_infos,
         runtime = runtime,
     )
     return [RustProtoInfo(
@@ -396,7 +394,7 @@ def _make_proto_library_aspect(is_upb):
     return aspect(
         implementation = (_rust_upb_proto_aspect_impl if is_upb else _rust_cc_proto_aspect_impl),
         attr_aspects = ["deps"],
-        requires = ([upb_minitable_proto_library_aspect] if is_upb else [cc_proto_aspect]),
+        requires = ([] if is_upb else [cc_proto_aspect]),
         attrs = {
             "_collect_cc_coverage": attr.label(
                 default = Label("@rules_rust//util:collect_coverage"),

rust/test/rust_proto_library_unit_test/rust_proto_library_unit_test.bzl

@@ -116,10 +116,6 @@ def _rust_upb_aspect_test_impl(ctx):
     # The action needs to produce a .rlib artifact (sometimes .rmeta as well, not tested here).
     asserts.true(env, rustc_action.outputs.to_list()[0].path.endswith(".rlib"))
 
-    # The aspect needs to provide CcInfo that passes UPB gencode to the eventual linking.
-    _find_linker_input(target_under_test[RustProtoInfo], "child.upb_minitable")
-    _find_linker_input(target_under_test[RustProtoInfo], "parent.upb_minitable")
-
     return analysistest.end(env)
 
 rust_upb_aspect_test = analysistest.make(_rust_upb_aspect_test_impl)

rust/test/shared/BUILD

@@ -320,6 +320,7 @@ rust_test(
     },
     deps = [
         "//rust:protobuf_cpp_export",
+        "//rust/test:map_unittest_cpp_rust_proto",
         "//rust/test:nested_cpp_rust_proto",
         "@crate_index//:googletest",
     ],
@@ -333,6 +334,7 @@ rust_test(
     },
     deps = [
         "//rust:protobuf_upb_export",
+        "//rust/test:map_unittest_upb_rust_proto",
         "//rust/test:nested_upb_rust_proto",
         "@crate_index//:googletest",
     ],

rust/test/shared/simple_nested_test.rs

@@ -6,9 +6,11 @@
 // https://developers.google.com/open-source/licenses/bsd
 
 use googletest::prelude::*;
+use map_unittest_rust_proto::TestRecursiveMapMessage;
 use nested_rust_proto::outer::inner::InnerEnum;
 use nested_rust_proto::outer::InnerView;
 use nested_rust_proto::*;
+use protobuf::prelude::*;
 
 #[gtest]
 fn test_deeply_nested_message() {
@@ -120,3 +122,12 @@ fn test_recursive_mut() {
     // let nested = rec.rec_mut().rec_mut().rec_mut();
     // assert_that!(nested.num(), eq(0));
 }
+
+#[gtest]
+fn test_recursive_map() {
+    let mut m1 = TestRecursiveMapMessage::new();
+    m1.a_mut().insert("k", TestRecursiveMapMessage::new());
+    let m2 = TestRecursiveMapMessage::parse(&m1.serialize().unwrap()).unwrap();
+    assert_that!(m2.a().len(), eq(1));
+    expect_that!(m2.a().get("k"), not(none()));
+}

rust/test/upb/BUILD

@@ -49,3 +49,16 @@ rust_test(
         "@crate_index//:googletest",
     ],
 )
+
+rust_test(
+    name = "thread_local_arena_test",
+    srcs = ["thread_local_arena_test.rs"],
+    aliases = {
+        "//rust:protobuf_upb_export": "protobuf",
+    },
+    deps = [
+        "//rust:protobuf_upb_export",
+        "//rust/test:unittest_upb_rust_proto",
+        "@crate_index//:googletest",
+    ],
+)

rust/test/upb/thread_local_arena_test.rs

@@ -0,0 +1,27 @@
+// Protocol Buffers - Google's data interchange format
+// Copyright 2024 Google LLC.  All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+use googletest::prelude::*;
+use protobuf::prelude::*;
+use unittest_rust_proto::TestAllTypes;
+
+#[gtest]
+fn test_thread_local_arena() {
+    // We use a thread-local arena to allocate minitables, so let's verify that a minitable remains
+    // valid even after that the thread that initialized it has exited.
+    let handle = std::thread::spawn(|| {
+        let mut m = TestAllTypes::new();
+        m.set_optional_int32(3);
+        let _ = m.serialize();
+    });
+    handle.join().unwrap();
+
+    let mut m1 = TestAllTypes::new();
+    m1.set_optional_int32(5);
+    let m2 = TestAllTypes::parse(&m1.serialize().unwrap()).unwrap();
+    assert_that!(m1.optional_int32(), eq(m2.optional_int32()));
+}

rust/upb.rs

@@ -33,6 +33,18 @@ pub type RawRepeatedField = upb::RawArray;
 pub type RawMap = upb::RawMap;
 pub type PtrAndLen = upb::StringView;
 
+// This struct represents a raw minitable pointer. We need it to be Send and Sync so that we can
+// store it in a static OnceLock for lazy initialization of minitables. It should not be used for
+// any other purpose.
+pub struct MiniTablePtr(pub *mut upb_MiniTable);
+unsafe impl Send for MiniTablePtr {}
+unsafe impl Sync for MiniTablePtr {}
+
+// Same as above, but for enum minitables.
+pub struct MiniTableEnumPtr(pub *const upb_MiniTableEnum);
+unsafe impl Send for MiniTableEnumPtr {}
+unsafe impl Sync for MiniTableEnumPtr {}
+
 impl From<&ProtoStr> for PtrAndLen {
     fn from(s: &ProtoStr) -> Self {
         let bytes = s.as_bytes();
@@ -61,6 +73,12 @@ impl ScratchSpace {
     }
 }
 
+thread_local! {
+    // We need to avoid dropping this Arena, because we use it to build mini tables that
+    // effectively have 'static lifetimes.
+    pub static THREAD_LOCAL_ARENA: ManuallyDrop<Arena> = ManuallyDrop::new(Arena::new());
+}
+
 #[doc(hidden)]
 pub type SerializedData = upb::OwnedArenaBox<[u8]>;
 

rust/upb/BUILD

@@ -52,6 +52,7 @@ cc_library(
         "//upb:message",
         "//upb:message_compare",
         "//upb:message_copy",
+        "//upb/mini_descriptor",
         "//upb/mini_table",
         "//upb/text:debug",
     ],

rust/upb/associated_mini_table.rs

@@ -6,9 +6,10 @@
 // https://developers.google.com/open-source/licenses/bsd
 
 use super::upb_MiniTable;
+use super::upb_MiniTableEnum;
 
-/// A trait for types which have a constant associated MiniTable (e.g.
-/// generated messages, and their mut and view proxy types).
+/// A trait for types which have an associated MiniTable (e.g. generated
+/// messages, and their mut and view proxy types).
 ///
 /// upb_Message in C is effectively a DST type, where instances are created with
 /// a MiniTable (and have a size dependent on the given MiniTable). Many upb
@@ -20,12 +21,6 @@ use super::upb_MiniTable;
 /// which hold upb_Message* to simplify ensuring the upb C invariants
 /// are maintained.
 ///
-/// Note that this would prefer to be a `const MINI_TABLE: *const upb_MiniTable`
-/// to statically associate a single MiniTable, but as long as the MiniTable is
-/// an extern "C" we cannot do that without the unstable `const_refs_to_static`.
-/// After that feature is stabilized (or if we move the MiniTable generation to
-/// .rs) this will be switched.
-///
 /// SAFETY:
 /// - The MiniTable pointer must be from Protobuf code generation and follow the
 ///   corresponding invariants associated with upb's C API (the pointer should
@@ -34,3 +29,8 @@ use super::upb_MiniTable;
 pub unsafe trait AssociatedMiniTable {
     fn mini_table() -> *const upb_MiniTable;
 }
+
+/// A trait for closed enums that have an associated MiniTable.
+pub unsafe trait AssociatedMiniTableEnum {
+    fn mini_table() -> *const upb_MiniTableEnum;
+}

rust/upb/lib.rs

@@ -19,7 +19,7 @@ pub use array::{
 };
 
 mod associated_mini_table;
-pub use associated_mini_table::AssociatedMiniTable;
+pub use associated_mini_table::{AssociatedMiniTable, AssociatedMiniTableEnum};
 
 mod ctype;
 pub use ctype::CType;
@@ -41,8 +41,9 @@ pub use message_value::{upb_MessageValue, upb_MutableMessageValue};
 
 mod mini_table;
 pub use mini_table::{
-    upb_MiniTable, upb_MiniTableField, upb_MiniTable_FindFieldByNumber,
-    upb_MiniTable_GetFieldByIndex, upb_MiniTable_SubMessage, RawMiniTable, RawMiniTableField,
+    upb_MiniTable, upb_MiniTableEnum, upb_MiniTableEnum_Build, upb_MiniTableField,
+    upb_MiniTable_Build, upb_MiniTable_FindFieldByNumber, upb_MiniTable_GetFieldByIndex,
+    upb_MiniTable_Link, upb_MiniTable_SubMessage, upb_Status, RawMiniTable, RawMiniTableField,
 };
 
 mod opaque_pointee;