Add a reserve method on ProxiedInRepeated

PiperOrigin-RevId: 620291225
Derek Benson 2024-03-29 11:10:30 -07:00 committed by Copybara-Service
parent d5b7d7a87f
commit 1bebb17655
8 changed files with 77 additions and 3 deletions

View File

@ -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 >],
],
)*);
}

View File

@ -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 capacity) { \
r->Reserve(r->size() + capacity); \
}
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 capacity) { \
r->Reserve(r->size() + capacity); \
}
expose_repeated_ptr_field_methods(ProtoStr);

View File

@ -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());
}

View File

@ -226,6 +226,13 @@ macro_rules! impl_repeated_base {
)
}
}
fn repeated_reserve(mut f: Mut<Repeated<$t>>, additional: usize) {
unsafe {
let arena = f.raw_arena(Private);
let size = upb_Array_Size(f.as_raw(Private));
upb_Array_Reserve(f.as_raw(Private), size + additional, arena);
}
}
};
}

View File

@ -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;

View File

@ -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;

View File

@ -450,6 +450,19 @@ 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 a valid `upb_Array*`.
unsafe {
let mut f = $pbr$::cast_enum_repeated_mut($pbi$::Private, r);
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));
}
}
}
// SAFETY: this is an enum type

View File

@ -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,14 @@ 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 +478,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 +1177,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 +1243,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$