From 91dfd8698a230d5d2d33b04f46fa20af0686d4f1 Mon Sep 17 00:00:00 2001
From: Derek Benson <derekbenson@google.com>
Date: Mon, 1 Jul 2024 13:15:30 -0700
Subject: [PATCH] add .. operator to proto! for merge_from. This mimics Rust's
 struct update syntax.
 https://doc.rust-lang.org/stable/book/ch05-01-defining-structs.html#creating-instances-from-other-instances-with-struct-update-syntax

PiperOrigin-RevId: 648464577
---
 rust/proto_macro.rs                  | 18 +++++++++++++++++
 rust/test/shared/proto_macro_test.rs | 29 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/rust/proto_macro.rs b/rust/proto_macro.rs
index 95b14232bf0c..a70afffa6226 100644
--- a/rust/proto_macro.rs
+++ b/rust/proto_macro.rs
@@ -40,6 +40,17 @@ macro_rules! proto {
 #[macro_export(local_inner_macros)]
 #[doc(hidden)]
 macro_rules! proto_internal {
+    (@merge $msg:ident $ident:ident : $expr:expr, $($rest:tt)*) => {
+        proto_internal!(@merge $msg $($rest)*);
+    };
+
+    (@merge $msg:ident $ident:ident : $expr:expr) => {
+    };
+
+    (@merge $msg:ident ..$expr:expr) => {
+        $msg.merge_from($expr);
+    };
+
     // nested message,
     (@msg $msg:ident $submsg:ident : $($msgtype:ident)::+ { $field:ident : $($value:tt)* }, $($rest:tt)*) => {
         proto_internal!(@msg $msg $submsg : $($msgtype)::+ { $field : $($value)* });
@@ -54,12 +65,14 @@ macro_rules! proto_internal {
     (@msg $msg:ident $submsg:ident : $($msgtype:ident)::+ { $field:ident : $($value:tt)* }) => {
         {
             let mut $msg: <$($msgtype)::+ as $crate::MutProxied>::Mut<'_> = $crate::__internal::paste!($msg.[<$submsg _mut>]());
+            proto_internal!(@merge $msg $field : $($value)*);
             proto_internal!(@msg $msg $field : $($value)*);
         }
     };
     (@msg $msg:ident $submsg:ident : ::$($msgtype:ident)::+ { $field:ident : $($value:tt)* }) => {
         {
             let mut $msg: <::$($msgtype)::+ as $crate::MutProxied>::Mut<'_> = $crate::__internal::paste!($msg.[<$submsg _mut>]());
+            proto_internal!(@merge $msg $field : $($value)*);
             proto_internal!(@msg $msg $field : $($value)*);
         }
     };
@@ -98,12 +111,17 @@ macro_rules! proto_internal {
         $crate::__internal::paste!($msg.[<set_ $ident>]($expr));
     };
 
+    (@msg $msg:ident ..$expr:expr) => {
+    };
+
     (@msg $msg:ident) => {};
+    (@merge $msg:ident) => {};
 
     // entry point
     ($msgtype:ty { $($tt:tt)* }) => {
         {
             let mut message = <$msgtype>::new();
+            proto_internal!(@merge message $($tt)*);
             proto_internal!(@msg message $($tt)*);
             message
         }
diff --git a/rust/test/shared/proto_macro_test.rs b/rust/test/shared/proto_macro_test.rs
index c7865cb1adc9..07c6f6bb80b5 100644
--- a/rust/test/shared/proto_macro_test.rs
+++ b/rust/test/shared/proto_macro_test.rs
@@ -122,3 +122,32 @@ fn test_recursive_msg() {
     assert_that!(msg.child().child().payload().optional_int32(), eq(42));
     assert_that!(msg.child().child().child().payload().optional_int32(), eq(43));
 }
+
+#[test]
+fn test_spread_msg() {
+    let msg = proto!(TestAllTypes { optional_nested_message: NestedMessage { bb: 42 } });
+    let msg2 = proto!(TestAllTypes { ..msg.as_view() });
+    assert_that!(msg2.optional_nested_message().bb(), eq(42));
+    let msg3 = proto!(TestAllTypes { optional_int32: 1, ..msg.as_view() });
+    assert_that!(msg3.optional_nested_message().bb(), eq(42));
+    assert_that!(msg3.optional_int32(), eq(1));
+}
+
+#[test]
+fn test_spread_nested_msg() {
+    let msg = proto!(NestedTestAllTypes {
+        child: NestedTestAllTypes {
+            payload: TestAllTypes { optional_int32: 41 },
+            child: NestedTestAllTypes {
+                child: NestedTestAllTypes { payload: TestAllTypes { optional_int32: 43 } },
+                payload: TestAllTypes { optional_int32: 42 }
+            }
+        }
+    });
+    let msg2 = proto!(NestedTestAllTypes {
+        child: NestedTestAllTypes { payload: TestAllTypes { optional_int32: 100 }, ..msg.child() }
+    });
+    assert_that!(msg2.child().payload().optional_int32(), eq(100));
+    assert_that!(msg2.child().child().payload().optional_int32(), eq(42));
+    assert_that!(msg2.child().child().child().payload().optional_int32(), eq(43));
+}