From d76fdc56bb2e9eb6005d09e5adb5518e2d67f2d7 Mon Sep 17 00:00:00 2001 From: Jakob Buchgraber Date: Mon, 29 Apr 2024 01:31:58 -0700 Subject: [PATCH] Implement IntoProxied for messages PiperOrigin-RevId: 628992357 --- rust/proxied.rs | 15 +++- rust/shared.rs | 3 +- rust/test/shared/accessors_test.rs | 19 ----- .../rust/accessors/singular_message.cc | 49 +++++++++++-- .../rust/accessors/singular_string.cc | 8 +-- src/google/protobuf/compiler/rust/message.cc | 69 +++++++++++-------- 6 files changed, 105 insertions(+), 58 deletions(-) diff --git a/rust/proxied.rs b/rust/proxied.rs index 41ddb6a743..beef53c15d 100644 --- a/rust/proxied.rs +++ b/rust/proxied.rs @@ -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` @@ -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 { + fn into(self, _private: Private) -> T; +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/shared.rs b/rust/shared.rs index 26d606a4ca..4eacf7c767 100644 --- a/rust/shared.rs +++ b/rust/shared.rs @@ -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, diff --git a/rust/test/shared/accessors_test.rs b/rust/test/shared/accessors_test.rs index fa11b86d07..36caf8b00c 100644 --- a/rust/test/shared/accessors_test.rs +++ b/rust/test/shared/accessors_test.rs @@ -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)); } @@ -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(); diff --git a/src/google/protobuf/compiler/rust/accessors/singular_message.cc b/src/google/protobuf/compiler/rust/accessors/singular_message.cc index 73acd60164..347f1d5dcd 100644 --- a/src/google/protobuf/compiler/rust/accessors/singular_message.cc +++ b/src/google/protobuf/compiler/rust/accessors/singular_message.cc @@ -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", [&] { @@ -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) } @@ -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"); }}, @@ -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()) { @@ -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")}, @@ -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"); } diff --git a/src/google/protobuf/compiler/rust/accessors/singular_string.cc b/src/google/protobuf/compiler/rust/accessors/singular_string.cc index 1de8f5b7ae..ee45d71efb 100644 --- a/src/google/protobuf/compiler/rust/accessors/singular_string.cc +++ b/src/google/protobuf/compiler/rust/accessors/singular_string.cc @@ -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 ); } @@ -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, ) @@ -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$, ) ) diff --git a/src/google/protobuf/compiler/rust/message.cc b/src/google/protobuf/compiler/rust/message.cc index 5ccf2541ac..67c569f4f5 100644 --- a/src/google/protobuf/compiler/rust/message.cc +++ b/src/google/protobuf/compiler/rust/message.cc @@ -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"); @@ -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", @@ -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$ @@ -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 } @@ -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) }