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;
