Skip to content

Commit

Permalink
Implement IntoProxied for messages
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 627426777
  • Loading branch information
@buchgr
buchgr authored and Copybara-Service committed Apr 29, 2024
1 parent 1d0028d commit 1d1cbca
Show file tree
Hide file tree
Showing 6 changed files with 105 additions and 58 deletions.
15 changes: 14 additions & 1 deletion rust/proxied.rs
View file
Expand Up @@ -134,7 +134,7 @@ pub trait ViewProxy<'msg>: 'msg + Sync + Unpin + Sized + Debug {
/// y: View<'b, T>,
/// ) -> [View<'b, T>; 2]
/// where
/// T: Proxied,
/// T: MutProxied,
/// 'a: 'b,
/// {
/// // `[x, y]` fails to compile because `'a` is not the same as `'b` and the `View`
Expand Down Expand Up @@ -295,6 +295,19 @@ where
}
}

/// A value to `Proxied`-value conversion that consumes the input value.
///
/// All setter functions accept types that implement `IntoProxied`. The purpose
/// of `IntoProxied` is to allow setting arbitrary values on Protobuf fields
/// with the minimal number of copies.
///
/// This trait must not be implemented on types outside the Protobuf codegen and
/// runtime. We expect it to change in backwards incompatible ways in the
/// future.
pub trait IntoProxied<T: Proxied> {
fn into(self, _private: Private) -> T;
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down
3 changes: 2 additions & 1 deletion rust/shared.rs
View file
Expand Up @@ -28,7 +28,8 @@ pub mod __public {
pub use crate::primitive::PrimitiveMut;
pub use crate::proto;
pub use crate::proxied::{
Mut, MutProxied, MutProxy, Proxied, ProxiedWithPresence, SettableValue, View, ViewProxy,
IntoProxied, Mut, MutProxied, MutProxy, Proxied, ProxiedWithPresence, SettableValue, View,
ViewProxy,
};
pub use crate::repeated::{
ProxiedInRepeated, Repeated, RepeatedIter, RepeatedMut, RepeatedView,
Expand Down
19 changes: 0 additions & 19 deletions rust/test/shared/accessors_test.rs
View file
Expand Up @@ -507,16 +507,9 @@ fn test_singular_msg_field() {

assert_that!(msg.has_optional_nested_message(), eq(false));
let mut nested_msg_mut = msg.optional_nested_message_mut();

// test reading an int inside a mut
assert_that!(nested_msg_mut.bb(), eq(0));

// Test setting an owned NestedMessage onto another message.
let mut new_nested = NestedMessage::new();
new_nested.set_bb(7);
nested_msg_mut.set(new_nested);
assert_that!(nested_msg_mut.bb(), eq(7));

assert_that!(msg.has_optional_nested_message(), eq(true));
}

Expand Down Expand Up @@ -762,18 +755,6 @@ fn test_msg_oneof_default_accessors() {
// TODO: Add tests covering a message-type field in a oneof.
}

#[test]
fn test_set_message_from_view() {
use protobuf::MutProxy;

let mut m1 = TestAllTypes::new();
m1.set_optional_int32(1);
let mut m2 = TestAllTypes::new();
m2.as_mut().set(m1.as_view());

assert_that!(m2.optional_int32(), eq(1i32));
}

#[test]
fn test_group() {
let mut m = TestAllTypes::new();
Expand Down
49 changes: 44 additions & 5 deletions src/google/protobuf/compiler/rust/accessors/singular_message.cc
View file
Expand Up @@ -34,6 +34,7 @@ void SingularMessage::InMsgImpl(Context& ctx, const FieldDescriptor& field,
{"getter_mut_thunk", ThunkName(ctx, field, "get_mut")},
{"clearer_thunk", ThunkName(ctx, field, "clear")},
{"hazzer_thunk", ThunkName(ctx, field, "has")},
{"set_allocated_thunk", ThunkName(ctx, field, "set")},
{
"getter_body",
[&] {
Expand Down Expand Up @@ -95,7 +96,7 @@ void SingularMessage::InMsgImpl(Context& ctx, const FieldDescriptor& field,
unsafe {
let has = self.has_$raw_field_name$();
$pbi$::new_vtable_field_entry($pbi$::Private,
self.as_mutator_message_ref(),
self.as_mutator_message_ref($pbi$::Private),
&VTABLE,
has)
}
Expand All @@ -112,14 +113,44 @@ void SingularMessage::InMsgImpl(Context& ctx, const FieldDescriptor& field,
}
)rs");
}},
{"setter_body",
[&] {
if (accessor_case == AccessorCase::VIEW) return;
if (ctx.is_upb()) {
ctx.Emit({}, R"rs(
// The message and arena are dropped after the setter. The
// memory remains allocated as we fuse the arena with the
// parent message's arena.
let mut msg = val.into($pbi$::Private);
self.as_mutator_message_ref($pbi$::Private)
.arena($pbi$::Private)
.fuse(msg.as_mutator_message_ref($pbi$::Private).arena($pbi$::Private));
unsafe {
$set_allocated_thunk$(self.as_mutator_message_ref($pbi$::Private).msg(),
msg.as_mutator_message_ref($pbi$::Private).msg());
}
)rs");
} else {
ctx.Emit({}, R"rs(
// Prevent the memory from being deallocated. The setter
// transfers ownership of the memory to the parent message.
let mut msg = std::mem::ManuallyDrop::new(val.into($pbi$::Private));
unsafe {
$set_allocated_thunk$(self.as_mutator_message_ref($pbi$::Private).msg(),
msg.as_mutator_message_ref($pbi$::Private).msg());
}
)rs");
}
}},
{"setter",
[&] {
if (accessor_case == AccessorCase::VIEW) return;
ctx.Emit(R"rs(
pub fn set_$raw_field_name$(&mut self, val: impl $pb$::SettableValue<$msg_type$>) {
//~ TODO: Optimize this to not go through the
//~ FieldEntry.
self.$raw_field_name$_entry().set(val);
pub fn set_$raw_field_name$(&mut self,
val: impl $pb$::IntoProxied<$msg_type$>) {
$setter_body$
}
)rs");
}},
Expand Down Expand Up @@ -157,6 +188,7 @@ void SingularMessage::InExternC(Context& ctx,
{"getter_mut_thunk", ThunkName(ctx, field, "get_mut")},
{"clearer_thunk", ThunkName(ctx, field, "clear")},
{"hazzer_thunk", ThunkName(ctx, field, "has")},
{"set_allocated_thunk", ThunkName(ctx, field, "set")},
{"getter_mut",
[&] {
if (ctx.is_cpp()) {
Expand Down Expand Up @@ -188,12 +220,16 @@ void SingularMessage::InExternC(Context& ctx,
$getter_mut$
fn $clearer_thunk$(raw_msg: $pbr$::RawMessage);
fn $hazzer_thunk$(raw_msg: $pbr$::RawMessage) -> bool;
fn $set_allocated_thunk$(raw_msg: $pbr$::RawMessage,
field_msg: $pbr$::RawMessage);
)rs");
}

void SingularMessage::InThunkCc(Context& ctx,
const FieldDescriptor& field) const {
ctx.Emit({{"QualifiedMsg", cpp::QualifiedClassName(field.containing_type())},
{"FieldMsg", cpp::QualifiedClassName(field.message_type())},
{"set_allocated_thunk", ThunkName(ctx, field, "set")},
{"getter_thunk", ThunkName(ctx, field, "get")},
{"getter_mut_thunk", ThunkName(ctx, field, "get_mut")},
{"clearer_thunk", ThunkName(ctx, field, "clear")},
Expand All @@ -208,6 +244,9 @@ void SingularMessage::InThunkCc(Context& ctx,
}
void $clearer_thunk$($QualifiedMsg$* msg) { msg->clear_$field$(); }
bool $hazzer_thunk$($QualifiedMsg$* msg) { return msg->has_$field$(); }
void $set_allocated_thunk$($QualifiedMsg$* msg, $FieldMsg$* sub_msg) {
msg->set_allocated_$field$(sub_msg);
}
)cc");
}

Expand Down
8 changes: 4 additions & 4 deletions src/google/protobuf/compiler/rust/accessors/singular_string.cc
View file
Expand Up @@ -87,13 +87,13 @@ void SingularString::InMsgImpl(Context& ctx, const FieldDescriptor& field,
pub fn set_$raw_field_name$(&mut self, val: impl std::convert::AsRef<$proxied_type$>) {
let string_view: $pbr$::PtrAndLen =
$pbr$::copy_bytes_in_arena_if_needed_by_runtime(
self.as_mutator_message_ref(),
self.as_mutator_message_ref($pbi$::Private),
val.as_ref().into()
).into();
unsafe {
$setter_thunk$(
self.as_mutator_message_ref().msg(),
self.as_mutator_message_ref($pbi$::Private).msg(),
string_view
);
}
Expand Down Expand Up @@ -160,7 +160,7 @@ void SingularString::InMsgImpl(Context& ctx, const FieldDescriptor& field,
let has = $hazzer_thunk$(self.raw_msg());
$pbi$::new_vtable_field_entry(
$pbi$::Private,
self.as_mutator_message_ref(),
self.as_mutator_message_ref($pbi$::Private),
$Msg$::$vtable_name$,
has,
)
Expand All @@ -176,7 +176,7 @@ void SingularString::InMsgImpl(Context& ctx, const FieldDescriptor& field,
$pbi$::Private,
$pbi$::RawVTableMutator::new(
$pbi$::Private,
self.as_mutator_message_ref(),
self.as_mutator_message_ref($pbi$::Private),
$Msg$::$vtable_name$,
)
)
Expand Down
69 changes: 41 additions & 28 deletions src/google/protobuf/compiler/rust/message.cc
View file
Expand Up @@ -226,33 +226,56 @@ void MessageDrop(Context& ctx, const Descriptor& msg) {
)rs");
}

void MessageSettableValueForView(Context& ctx, const Descriptor& msg) {
void IntoProxiedForMessage(Context& ctx, const Descriptor& msg) {
switch (ctx.opts().kernel) {
case Kernel::kCpp:
ctx.Emit({{"copy_from_thunk", ThunkName(ctx, msg, "copy_from")}}, R"rs(
impl<'msg> $pb$::SettableValue<$Msg$> for $Msg$View<'msg> {
fn set_on<'dst>(
self, _private: $pbi$::Private, mutator: $pb$::Mut<'dst, $Msg$>)
where $Msg$: 'dst {
unsafe { $copy_from_thunk$(mutator.inner.msg(), self.msg) };
impl<'msg> $pb$::IntoProxied<$Msg$> for $Msg$View<'msg> {
fn into(self, _private: $pbi$::Private) -> $Msg$ {
let dst = $Msg$::new();
unsafe { $copy_from_thunk$(dst.inner.msg, self.msg) };
dst
}
}
impl<'msg> $pb$::IntoProxied<$Msg$> for $Msg$Mut<'msg> {
fn into(self, _private: $pbi$::Private) -> $Msg$ {
$pb$::IntoProxied::into($pb$::ViewProxy::into_view(self), _private)
}
}
impl $pb$::IntoProxied<$Msg$> for $Msg$ {
fn into(self, _private: $pbi$::Private) -> $Msg$ {
self
}
}
)rs");
return;

case Kernel::kUpb:
// TODO: Add owned SettableValue impl for upb messages.
ctx.Emit({{"minitable", UpbMinitableName(msg)}}, R"rs(
impl<'msg> $pb$::SettableValue<$Msg$> for $Msg$View<'msg> {
fn set_on<'dst>(
self, _private: $pbi$::Private, mutator: $pb$::Mut<'dst, $Msg$>)
where $Msg$: 'dst {
impl<'msg> $pb$::IntoProxied<$Msg$> for $Msg$View<'msg> {
fn into(self, _private: $pbi$::Private) -> $Msg$ {
let dst = $Msg$::new();
unsafe { $pbr$::upb_Message_DeepCopy(
mutator.inner.msg(),
dst.inner.msg,
self.msg,
$std$::ptr::addr_of!($minitable$),
mutator.inner.arena($pbi$::Private).raw(),
dst.inner.arena.raw(),
) };
dst
}
}
impl<'msg> $pb$::IntoProxied<$Msg$> for $Msg$Mut<'msg> {
fn into(self, _private: $pbi$::Private) -> $Msg$ {
$pb$::IntoProxied::into($pb$::ViewProxy::into_view(self), _private)
}
}
impl $pb$::IntoProxied<$Msg$> for $Msg$ {
fn into(self, _private: $pbi$::Private) -> $Msg$ {
self
}
}
)rs");
Expand Down Expand Up @@ -797,8 +820,7 @@ void GenerateRs(Context& ctx, const Descriptor& msg) {
AccessorCase::MUT);
}
}},
{"settable_impl_for_view",
[&] { MessageSettableValueForView(ctx, msg); }},
{"into_proxied_impl", [&] { IntoProxiedForMessage(ctx, msg); }},
{"repeated_impl", [&] { MessageProxiedInRepeated(ctx, msg); }},
{"map_value_impl", [&] { MessageProxiedInMapValue(ctx, msg); }},
{"unwrap_upb",
Expand Down Expand Up @@ -960,17 +982,7 @@ void GenerateRs(Context& ctx, const Descriptor& msg) {
}
}
$settable_impl_for_view$
impl $pb$::SettableValue<$Msg$> for $Msg$ {
fn set_on<'dst>(
self, _private: $pbi$::Private, mutator: $pb$::Mut<'dst, $Msg$>)
where $Msg$: 'dst {
//~ TODO: b/320701507 - This current will copy the message and then
//~ drop it, this copy would be avoided on upb kernel.
self.as_view().set_on($pbi$::Private, mutator);
}
}
$into_proxied_impl$
$repeated_impl$
$map_value_impl$
Expand Down Expand Up @@ -1013,7 +1025,8 @@ void GenerateRs(Context& ctx, const Descriptor& msg) {
self.inner.msg()
}
fn as_mutator_message_ref(&mut self) -> $pbr$::MutatorMessageRef<'msg> {
pub fn as_mutator_message_ref(&mut self, _private: $pbi$::Private)
-> $pbr$::MutatorMessageRef<'msg> {
self.inner
}
Expand Down Expand Up @@ -1059,7 +1072,7 @@ void GenerateRs(Context& ctx, const Descriptor& msg) {
self.inner.msg
}
fn as_mutator_message_ref(&mut self) -> $pbr$::MutatorMessageRef {
pub fn as_mutator_message_ref(&mut self, _private: $pbi$::Private) -> $pbr$::MutatorMessageRef {
$pbr$::MutatorMessageRef::new($pbi$::Private, &mut self.inner)
}
Expand Down

0 comments on commit 1d1cbca

Please sign in to comment.