Add a reserve method on ProxiedInRepeated PiperOrigin-RevId: 633179853
diff --git a/rust/cpp.rs b/rust/cpp.rs index 1bd1526..87b4825 100644 --- a/rust/cpp.rs +++ b/rust/cpp.rs
@@ -374,7 +374,8 @@ $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 @@ 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 @@ 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 @@ [< __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 @@ } } +/// 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,
diff --git a/rust/cpp_kernel/cpp_api.cc b/rust/cpp_kernel/cpp_api.cc index f37a1c7..753331c 100644 --- a/rust/cpp_kernel/cpp_api.cc +++ b/rust/cpp_kernel/cpp_api.cc
@@ -42,6 +42,10 @@ 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 @@ 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);
diff --git a/rust/repeated.rs b/rust/repeated.rs index 3183e42..c7c9762 100644 --- a/rust/repeated.rs +++ b/rust/repeated.rs
@@ -296,6 +296,10 @@ /// 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 @@ 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()); }
diff --git a/rust/upb.rs b/rust/upb.rs index 895f47b..858e7cf 100644 --- a/rust/upb.rs +++ b/rust/upb.rs
@@ -226,6 +226,16 @@ ) } } + 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 @@ } } +/// 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.
diff --git a/rust/upb/array.rs b/rust/upb/array.rs index 8c0abf8..a8b2fb1 100644 --- a/rust/upb/array.rs +++ b/rust/upb/array.rs
@@ -12,6 +12,7 @@ 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;
diff --git a/rust/upb/lib.rs b/rust/upb/lib.rs index 299044f..643018a 100644 --- a/rust/upb/lib.rs +++ b/rust/upb/lib.rs
@@ -4,8 +4,8 @@ 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;
diff --git a/src/google/protobuf/compiler/rust/enum.cc b/src/google/protobuf/compiler/rust/enum.cc index 857d283..0ed50e9 100644 --- a/src/google/protobuf/compiler/rust/enum.cc +++ b/src/google/protobuf/compiler/rust/enum.cc
@@ -450,6 +450,15 @@ $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
diff --git a/src/google/protobuf/compiler/rust/message.cc b/src/google/protobuf/compiler/rust/message.cc index 2b2babd..f6e4728 100644 --- a/src/google/protobuf/compiler/rust/message.cc +++ b/src/google/protobuf/compiler/rust/message.cc
@@ -180,6 +180,8 @@ {"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 @@ 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 @@ {"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 @@ $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 @@ $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 @@ {"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 @@ 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$