Implement IntoProxied for messages

PiperOrigin-RevId: 628992357

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<T: Proxied> {
+    fn into(self, _private: Private) -> T;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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,

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

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$>) {
+                pub fn set_$raw_field_name$(&mut self,
-                  //~ TODO: Optimize this to not go through the
+                  val: impl $pb$::IntoProxied<$msg_type$>) {
-                  //~ FieldEntry.
+
-                  self.$raw_field_name$_entry().set(val);
+                  $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");
 }
 

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$,
                     )
                   )

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> {
+        impl<'msg> $pb$::IntoProxied<$Msg$> for $Msg$View<'msg> {
-          fn set_on<'dst>(
+          fn into(self, _private: $pbi$::Private) -> $Msg$ {
-            self, _private: $pbi$::Private, mutator: $pb$::Mut<'dst, $Msg$>)
+            let dst = $Msg$::new();
-            where $Msg$: 'dst {
+            unsafe { $copy_from_thunk$(dst.inner.msg, self.msg) };
-            unsafe { $copy_from_thunk$(mutator.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> {
+        impl<'msg> $pb$::IntoProxied<$Msg$> for $Msg$View<'msg> {
-          fn set_on<'dst>(
+          fn into(self, _private: $pbi$::Private) -> $Msg$ {
-            self, _private: $pbi$::Private, mutator: $pb$::Mut<'dst, $Msg$>)
+            let dst = $Msg$::new();
-            where $Msg$: 'dst {
             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",
+       {"into_proxied_impl", [&] { IntoProxiedForMessage(ctx, msg); }},
-        [&] { MessageSettableValueForView(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$
+        $into_proxied_impl$
-
-        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);
-          }
-        }
 
         $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)
           }