Introduce an upb::AssociatedMiniTable trait in Rust.

PiperOrigin-RevId: 657263802

rust/upb/BUILD

@@ -19,6 +19,7 @@ rust_library(
     srcs = [
         "arena.rs",
         "array.rs",
+        "associated_mini_table.rs",
         "ctype.rs",
         "extension_registry.rs",
         "lib.rs",

rust/upb/associated_mini_table.rs

@@ -0,0 +1,30 @@
+// 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 super::upb_MiniTable;
+
+/// A trait for types which have a constant 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
+/// operations on the upb_Message* will require the same upb_MiniTable* to be
+/// passed in as a parameter, which is referred to as 'the associated MiniTable
+/// for the upb_Message instance' in safety comments.
+///
+/// This trait is a way to statically associate a MiniTable with Rust types
+/// which hold upb_Message* to simplify ensuring the upb C invariants
+/// are maintained.
+///
+/// SAFETY:
+/// - The MiniTable pointer must be from Protobuf code generation and follow the
+///   corresponding invariants associated with upb's C API (the pointer should
+///   always have the same non-null value, the underlying pointee should never
+///   be modified and should have 'static lifetime).
+pub unsafe trait AssociatedMiniTable {
+    const MINI_TABLE: *const upb_MiniTable;
+}

rust/upb/lib.rs

@@ -18,6 +18,9 @@ pub use array::{
     upb_Array_Size, RawArray,
 };
 
+mod associated_mini_table;
+pub use associated_mini_table::AssociatedMiniTable;
+
 mod ctype;
 pub use ctype::CType;
 

src/google/protobuf/compiler/rust/accessors/repeated_field.cc

@@ -114,7 +114,7 @@ void RepeatedField::InMsgImpl(Context& ctx, const FieldDescriptor& field,
                     pub fn set_$raw_field_name$(&mut self, src: impl $pb$::IntoProxied<$pb$::Repeated<$RsType$>>) {
                       let minitable_field = unsafe {
                         $pbr$::upb_MiniTable_FindFieldByNumber(
-                          Self::raw_minitable($pbi$::Private),
+                          <Self as $pbr$::AssociatedMiniTable>::MINI_TABLE,
                           $field_number$
                         )
                       };

src/google/protobuf/compiler/rust/generator.cc

@@ -177,6 +177,16 @@ bool RustGenerator::Generate(const FileDescriptor* file,
       {"Phantom", "::__std::marker::PhantomData"},
   });
 
+  // On upb we need to enable this feature to be able to have const pointers to
+  // upb_MiniTables while using the upb C minitable codegen. This
+  // feature appears that it'll imminently be considered stabilized
+  // (https://github.com/rust-lang/rust/issues/128183), and if we emit the
+  // MiniTables as const directly in .rs rather than using the upb minitable .h
+  // this could be avoided regardless.
+  if (ctx.is_upb() && file == &rust_generator_context.primary_file()) {
+    ctx.Emit("#![feature(const_refs_to_static)]\n");
+  }
+
   ctx.Emit({{"kernel", KernelRsName(ctx.opts().kernel)}}, R"rs(
     extern crate protobuf_$kernel$ as __pb;
     extern crate std as __std;

src/google/protobuf/compiler/rust/message.cc

@@ -76,11 +76,10 @@ void MessageSerialize(Context& ctx, const Descriptor& msg) {
     case Kernel::kUpb:
       ctx.Emit({{"minitable", UpbMinitableName(msg)}},
                R"rs(
-        // SAFETY: $minitable$ is a static of a const object.
-        let mini_table = unsafe { $std$::ptr::addr_of!($minitable$) };
-        // SAFETY: $minitable$ is the one associated with raw_msg().
+        // SAFETY: `MINI_TABLE` is the one associated with `self.raw_msg()`.
         let encoded = unsafe {
-          $pbr$::wire::encode(self.raw_msg(), mini_table)
+          $pbr$::wire::encode(self.raw_msg(),
+              <Self as $pbr$::AssociatedMiniTable>::MINI_TABLE)
         };
         //~ TODO: This discards the info we have about the reason
         //~ of the failure, we should try to keep it instead.
@@ -139,8 +138,6 @@ void MessageClearAndParse(Context& ctx, const Descriptor& msg) {
       ctx.Emit({{"minitable", UpbMinitableName(msg)}},
                R"rs(
         let mut msg = Self::new();
-        // SAFETY: $minitable$ is a static of a const object.
-        let mini_table = unsafe { $std$::ptr::addr_of!($minitable$) };
 
         // SAFETY:
         // - `data.as_ptr()` is valid to read for `data.len()`
@@ -148,8 +145,10 @@ void MessageClearAndParse(Context& ctx, const Descriptor& msg) {
         // - `msg.arena().raw()` is held for the same lifetime as `msg`.
         let status = unsafe {
           $pbr$::wire::decode(
-              data, msg.raw_msg(),
-              mini_table, msg.arena())
+              data,
+              msg.raw_msg(),
+              <Self as $pbr$::AssociatedMiniTable>::MINI_TABLE,
+              msg.arena())
         };
         match status {
           Ok(_) => {
@@ -178,13 +177,12 @@ void MessageDebug(Context& ctx, const Descriptor& msg) {
       return;
 
     case Kernel::kUpb:
-      ctx.Emit({{"minitable", UpbMinitableName(msg)}},
-               R"rs(
-        let mini_table = unsafe { $std$::ptr::addr_of!($minitable$) };
+      ctx.Emit(
+          R"rs(
         let string = unsafe {
           $pbr$::debug_string(
             self.raw_msg(),
-            mini_table,
+            <Self as $pbr$::AssociatedMiniTable>::MINI_TABLE
           )
         };
         write!(f, "{}", string)
@@ -298,7 +296,7 @@ void IntoProxiedForMessage(Context& ctx, const Descriptor& msg) {
             unsafe { $pbr$::upb_Message_DeepCopy(
               dst.inner.msg,
               self.msg,
-              $std$::ptr::addr_of!($minitable$),
+              <Self as $pbr$::AssociatedMiniTable>::MINI_TABLE,
               dst.inner.arena.raw(),
             ) };
             dst
@@ -317,11 +315,19 @@ void IntoProxiedForMessage(Context& ctx, const Descriptor& msg) {
   ABSL_LOG(FATAL) << "unreachable";
 }
 
-void MessageGetMinitable(Context& ctx, const Descriptor& msg) {
+void UpbGeneratedMessageTraitImpls(Context& ctx, const Descriptor& msg) {
   if (ctx.opts().kernel == Kernel::kUpb) {
     ctx.Emit({{"minitable", UpbMinitableName(msg)}}, R"rs(
-      pub fn raw_minitable(_private: $pbi$::Private) -> *const $pbr$::upb_MiniTable {
-        unsafe { $std$::ptr::addr_of!($minitable$) }
+      unsafe impl $pbr$::AssociatedMiniTable for $Msg$ {
+        const MINI_TABLE: *const $pbr$::upb_MiniTable = unsafe { $std$::ptr::addr_of!($minitable$) };
+      }
+
+      unsafe impl $pbr$::AssociatedMiniTable for $Msg$View<'_> {
+        const MINI_TABLE: *const $pbr$::upb_MiniTable = unsafe { $std$::ptr::addr_of!($minitable$) };
+      }
+
+      unsafe impl $pbr$::AssociatedMiniTable for $Msg$Mut<'_> {
+        const MINI_TABLE: *const $pbr$::upb_MiniTable = unsafe { $std$::ptr::addr_of!($minitable$) };
       }
     )rs");
   }
@@ -345,17 +351,14 @@ void MessageMergeFrom(Context& ctx, const Descriptor& msg) {
       return;
     case Kernel::kUpb:
       ctx.Emit(
-          {
-              {"minitable", UpbMinitableName(msg)},
-          },
           R"rs(
           pub fn merge_from<'src>(&mut self, src: impl $pb$::Proxy<'src, Proxied = $Msg$>) {
             // SAFETY: self and src are both valid `$Msg$`s.
             unsafe {
               assert!(
-                $pbr$::upb_Message_MergeFrom(self.raw_msg(), 
+                $pbr$::upb_Message_MergeFrom(self.raw_msg(),
                   src.as_view().raw_msg(),
-                  $std$::ptr::addr_of!($minitable$),
+                  <Self as $pbr$::AssociatedMiniTable>::MINI_TABLE,
                   // Use a nullptr for the ExtensionRegistry.
                   $std$::ptr::null(),
                   self.arena().raw())
@@ -481,7 +484,6 @@ void MessageProxiedInRepeated(Context& ctx, const Descriptor& msg) {
     case Kernel::kUpb:
       ctx.Emit(
           {
-              {"minitable", UpbMinitableName(msg)},
               {"new_thunk", ThunkName(ctx, msg, "new")},
           },
           R"rs(
@@ -559,9 +561,9 @@ void MessageProxiedInRepeated(Context& ctx, const Descriptor& msg) {
             dest: $pb$::Mut<$pb$::Repeated<Self>>,
           ) {
               // SAFETY:
-              // - Elements of `src` and `dest` have message minitable `$minitable$`.
+              // - Elements of `src` and `dest` have message minitable `MINI_TABLE`.
               unsafe {
-                $pbr$::repeated_message_copy_from(src, dest, $std$::ptr::addr_of!($minitable$));
+                $pbr$::repeated_message_copy_from(src, dest, <Self as $pbr$::AssociatedMiniTable>::MINI_TABLE);
               }
           }
 
@@ -698,7 +700,6 @@ void MessageProxiedInMapValue(Context& ctx, const Descriptor& msg) {
     case Kernel::kUpb:
       ctx.Emit(
           {
-              {"minitable", UpbMinitableName(msg)},
               {"new_thunk", ThunkName(ctx, msg, "new")},
           },
           R"rs(
@@ -940,7 +941,8 @@ void GenerateRs(Context& ctx, const Descriptor& msg) {
           }
         }},
        {"into_proxied_impl", [&] { IntoProxiedForMessage(ctx, msg); }},
-       {"get_upb_minitable", [&] { MessageGetMinitable(ctx, msg); }},
+       {"upb_generated_message_trait_impls",
+        [&] { UpbGeneratedMessageTraitImpls(ctx, msg); }},
        {"msg_merge_from", [&] { MessageMergeFrom(ctx, msg); }},
        {"repeated_impl", [&] { MessageProxiedInRepeated(ctx, msg); }},
        {"map_value_impl", [&] { MessageProxiedInMapValue(ctx, msg); }},
@@ -1155,8 +1157,6 @@ void GenerateRs(Context& ctx, const Descriptor& msg) {
 
           $msg_merge_from$
 
-          $get_upb_minitable$
-
           $raw_arena_getter_for_msgmut$
 
           $accessor_fns_for_muts$
@@ -1237,8 +1237,6 @@ void GenerateRs(Context& ctx, const Descriptor& msg) {
             self.as_mut().merge_from(src);
           }
 
-          $get_upb_minitable$
-
           $accessor_fns$
         }  // impl $Msg$
 
@@ -1270,6 +1268,8 @@ void GenerateRs(Context& ctx, const Descriptor& msg) {
           }
         }
 
+        $upb_generated_message_trait_impls$
+
         extern "C" {
           $Msg_externs$