Add a reserve method on ProxiedInRepeated

PiperOrigin-RevId: 633179853

rust/cpp.rs

@@ -374,7 +374,8 @@ macro_rules! impl_repeated_primitives {
         $get_thunk:ident,
         $set_thunk:ident,
         $clear_thunk:ident,
-        $copy_from_thunk:ident $(,)?
+        $copy_from_thunk:ident,
+        $reserve_thunk:ident $(,)?
     ]),* $(,)?) => {
         $(
             extern "C" {
@@ -391,6 +392,9 @@ macro_rules! impl_repeated_primitives {
                     v: <$t as CppTypeConversions>::ElemType);
                 fn $clear_thunk(f: RawRepeatedField);
                 fn $copy_from_thunk(src: RawRepeatedField, dst: RawRepeatedField);
+                fn $reserve_thunk(
+                    f: RawRepeatedField,
+                    additional: usize);
             }
 
             unsafe impl ProxiedInRepeated for $t {
@@ -423,6 +427,9 @@ macro_rules! impl_repeated_primitives {
                 fn repeated_copy_from(src: View<Repeated<$t>>, mut dest: Mut<Repeated<$t>>) {
                     unsafe { $copy_from_thunk(src.as_raw(Private), dest.as_raw(Private)) }
                 }
+                fn repeated_reserve(mut f: Mut<Repeated<$t>>, additional: usize) {
+                    unsafe { $reserve_thunk(f.as_raw(Private), additional) }
+                }
             }
         )*
     };
@@ -438,6 +445,7 @@ macro_rules! impl_repeated_primitives {
                     [< __pb_rust_RepeatedField_ $t _set >],
                     [< __pb_rust_RepeatedField_ $t _clear >],
                     [< __pb_rust_RepeatedField_ $t _copy_from >],
+                    [< __pb_rust_RepeatedField_ $t _reserve >],
                 ],
             )*);
         }
@@ -474,6 +482,17 @@ pub fn cast_enum_repeated_mut<E: Enum + ProxiedInRepeated>(
     }
 }
 
+/// Cast a `RepeatedMut<SomeEnum>` to `RepeatedMut<c_int>` and call
+/// repeated_reserve.
+pub fn reserve_enum_repeated_mut<E: Enum + ProxiedInRepeated>(
+    private: Private,
+    repeated: RepeatedMut<E>,
+    additional: usize,
+) {
+    let int_repeated = cast_enum_repeated_mut(private, repeated);
+    ProxiedInRepeated::repeated_reserve(int_repeated, additional);
+}
+
 #[derive(Debug)]
 pub struct InnerMap {
     pub(crate) raw: RawMap,

rust/cpp_kernel/cpp_api.cc

@@ -42,6 +42,10 @@ extern "C" {
   void __pb_rust_RepeatedField_##rust_ty##_clear(                             \
       google::protobuf::RepeatedField<ty>* r) {                                         \
     r->Clear();                                                               \
+  }                                                                           \
+  void __pb_rust_RepeatedField_##rust_ty##_reserve(                           \
+      google::protobuf::RepeatedField<ty>* r, size_t additional) {                      \
+    r->Reserve(r->size() + additional);                                       \
   }
 
 expose_repeated_field_methods(int32_t, i32);
@@ -89,6 +93,10 @@ expose_repeated_field_methods(int64_t, i64);
   void __pb_rust_RepeatedField_##ty##_clear(                           \
       google::protobuf::RepeatedPtrField<std::string>* r) {                      \
     r->Clear();                                                        \
+  }                                                                    \
+  void __pb_rust_RepeatedField_##ty##_reserve(                         \
+      google::protobuf::RepeatedPtrField<std::string>* r, size_t additional) {   \
+    r->Reserve(r->size() + additional);                                \
   }
 
 expose_repeated_ptr_field_methods(ProtoStr);

rust/repeated.rs

@@ -296,6 +296,10 @@ pub unsafe trait ProxiedInRepeated: Proxied {
 
     /// Copies the values in the `src` repeated field into `dest`.
     fn repeated_copy_from(src: View<Repeated<Self>>, dest: Mut<Repeated<Self>>);
+
+    /// Ensures that the repeated field has enough space allocated to insert at
+    /// least `additional` values without an allocation.
+    fn repeated_reserve(repeated: Mut<Repeated<Self>>, additional: usize);
 }
 
 /// An iterator over the values inside of a [`View<Repeated<T>>`](RepeatedView).
@@ -489,6 +493,8 @@ where
     ViewT: Into<View<'view, T>>,
 {
     fn extend<I: IntoIterator<Item = ViewT>>(&mut self, iter: I) {
+        let iter = iter.into_iter();
+        T::repeated_reserve(self.as_mut(), iter.size_hint().0);
         for item in iter {
             self.push(item.into());
         }

rust/upb.rs

@@ -226,6 +226,16 @@ macro_rules! impl_repeated_base {
                 )
             }
         }
+        fn repeated_reserve(mut f: Mut<Repeated<$t>>, additional: usize) {
+            // SAFETY:
+            // - `upb_Array_Reserve` is unsafe but assumed to be sound when called on a
+            //   valid array.
+            unsafe {
+                let arena = f.raw_arena(Private);
+                let size = upb_Array_Size(f.as_raw(Private));
+                assert!(upb_Array_Reserve(f.as_raw(Private), size + additional, arena));
+            }
+        }
     };
 }
 
@@ -373,6 +383,17 @@ pub fn cast_enum_repeated_mut<E: Enum + ProxiedInRepeated>(
     }
 }
 
+/// Cast a `RepeatedMut<SomeEnum>` to `RepeatedMut<i32>` and call
+/// repeated_reserve.
+pub fn reserve_enum_repeated_mut<E: Enum + ProxiedInRepeated>(
+    private: Private,
+    repeated: RepeatedMut<E>,
+    additional: usize,
+) {
+    let int_repeated = cast_enum_repeated_mut(private, repeated);
+    ProxiedInRepeated::repeated_reserve(int_repeated, additional);
+}
+
 /// Returns a static empty RepeatedView.
 pub fn empty_array<T: ?Sized + ProxiedInRepeated>() -> RepeatedView<'static, T> {
     // TODO: Consider creating a static empty array in C.

rust/upb/array.rs

@@ -12,6 +12,7 @@ extern "C" {
     pub fn upb_Array_Get(arr: RawArray, i: usize) -> upb_MessageValue;
     pub fn upb_Array_Append(arr: RawArray, val: upb_MessageValue, arena: RawArena) -> bool;
     pub fn upb_Array_Resize(arr: RawArray, size: usize, arena: RawArena) -> bool;
+    pub fn upb_Array_Reserve(arr: RawArray, size: usize, arena: RawArena) -> bool;
     pub fn upb_Array_MutableDataPtr(arr: RawArray) -> *mut std::ffi::c_void;
     pub fn upb_Array_DataPtr(arr: RawArray) -> *const std::ffi::c_void;
     pub fn upb_Array_GetMutable(arr: RawArray, i: usize) -> upb_MutableMessageValue;

rust/upb/lib.rs

@@ -4,8 +4,8 @@ pub use arena::{upb_Arena, Arena, RawArena};
 mod array;
 pub use array::{
     upb_Array, upb_Array_Append, upb_Array_DataPtr, upb_Array_Get, upb_Array_GetMutable,
-    upb_Array_MutableDataPtr, upb_Array_New, upb_Array_Resize, upb_Array_Set, upb_Array_Size,
-    RawArray,
+    upb_Array_MutableDataPtr, upb_Array_New, upb_Array_Reserve, upb_Array_Resize, upb_Array_Set,
+    upb_Array_Size, RawArray,
 };
 
 mod ctype;

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

@@ -450,6 +450,15 @@ void GenerateEnumDefinition(Context& ctx, const EnumDescriptor& desc) {
           $pbr$::cast_enum_repeated_mut($pbi$::Private, dest)
             .copy_from($pbr$::cast_enum_repeated_view($pbi$::Private, src))
         }
+
+        fn repeated_reserve(
+            mut r: $pb$::Mut<$pb$::Repeated<Self>>,
+            additional: usize,
+        ) {
+            // SAFETY:
+            // - `f.as_raw()` is valid.
+            $pbr$::reserve_enum_repeated_mut($pbi$::Private, r, additional);
+        }
       }
 
       // SAFETY: this is an enum type

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

@@ -180,6 +180,8 @@ void MessageExterns(Context& ctx, const Descriptor& msg) {
               {"repeated_clear_thunk", ThunkName(ctx, msg, "repeated_clear")},
               {"repeated_copy_from_thunk",
                ThunkName(ctx, msg, "repeated_copy_from")},
+              {"repeated_reserve_thunk",
+               ThunkName(ctx, msg, "repeated_reserve")},
           },
           R"rs(
           fn $new_thunk$() -> $pbr$::RawMessage;
@@ -193,6 +195,7 @@ void MessageExterns(Context& ctx, const Descriptor& msg) {
           fn $repeated_get_mut_thunk$(raw: $pbr$::RawRepeatedField, index: usize) -> $pbr$::RawMessage;
           fn $repeated_clear_thunk$(raw: $pbr$::RawRepeatedField);
           fn $repeated_copy_from_thunk$(dst: $pbr$::RawRepeatedField, src: $pbr$::RawRepeatedField);
+          fn $repeated_reserve_thunk$(raw: $pbr$::RawRepeatedField, additional: usize);
         )rs");
       return;
 
@@ -310,6 +313,8 @@ void MessageProxiedInRepeated(Context& ctx, const Descriptor& msg) {
               {"repeated_clear_thunk", ThunkName(ctx, msg, "repeated_clear")},
               {"repeated_copy_from_thunk",
                ThunkName(ctx, msg, "repeated_copy_from")},
+              {"repeated_reserve_thunk",
+               ThunkName(ctx, msg, "repeated_reserve")},
           },
           R"rs(
         unsafe impl $pb$::ProxiedInRepeated for $Msg$ {
@@ -372,6 +377,15 @@ void MessageProxiedInRepeated(Context& ctx, const Descriptor& msg) {
               $repeated_copy_from_thunk$(dest.as_raw($pbi$::Private), src.as_raw($pbi$::Private));
             }
           }
+
+          fn repeated_reserve(
+            mut f: $pb$::Mut<$pb$::Repeated<Self>>,
+            additional: usize,
+          ) {
+            // SAFETY:
+            // - `f.as_raw()` is a valid `RepeatedPtrField*`.
+            unsafe { $repeated_reserve_thunk$(f.as_raw($pbi$::Private), additional) }
+          }
         }
       )rs");
       return;
@@ -465,6 +479,18 @@ void MessageProxiedInRepeated(Context& ctx, const Descriptor& msg) {
                 $pbr$::repeated_message_copy_from(src, dest, $std$::ptr::addr_of!($minitable$));
               }
           }
+
+          fn repeated_reserve(
+            mut f: $pb$::Mut<$pb$::Repeated<Self>>,
+            additional: usize,
+          ) {
+            // SAFETY:
+            // - `f.as_raw()` is a valid `upb_Array*`.
+            unsafe {
+              let size = $pbr$::upb_Array_Size(f.as_raw($pbi$::Private));
+              $pbr$::upb_Array_Reserve(f.as_raw($pbi$::Private), size + additional, f.raw_arena($pbi$::Private));
+            }
+          }
         }
       )rs");
       return;
@@ -1152,6 +1178,7 @@ void GenerateThunksCc(Context& ctx, const Descriptor& msg) {
        {"repeated_add_thunk", ThunkName(ctx, msg, "repeated_add")},
        {"repeated_clear_thunk", ThunkName(ctx, msg, "repeated_clear")},
        {"repeated_copy_from_thunk", ThunkName(ctx, msg, "repeated_copy_from")},
+       {"repeated_reserve_thunk", ThunkName(ctx, msg, "repeated_reserve")},
        {"nested_msg_thunks",
         [&] {
           for (int i = 0; i < msg.nested_type_count(); ++i) {
@@ -1217,6 +1244,11 @@ void GenerateThunksCc(Context& ctx, const Descriptor& msg) {
           const google::protobuf::RepeatedPtrField<$QualifiedMsg$>& src) {
           dst = src;
         }
+        void $repeated_reserve_thunk$(
+          google::protobuf::RepeatedPtrField<$QualifiedMsg$>* field,
+          size_t additional) {
+          field->Reserve(field->size() + additional);
+        }
 
         $accessor_thunks$