From e35cf1025b0e845ccb809737ffc725354f59e642 Mon Sep 17 00:00:00 2001
From: Hong Shin <hongshin@google.com>
Date: Mon, 11 Dec 2023 14:11:14 -0800
Subject: [PATCH] Add submessage support for string_mut

We augment the TYPE_STRING block inside message.cc:GetterForViewOrMut so that we properly return a Mut for strings when it comes to submessages.

PiperOrigin-RevId: 589944498
---
 rust/test/shared/simple_nested_test.rs       |  5 +-
 src/google/protobuf/compiler/rust/message.cc | 49 +++++++++++++++-----
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/rust/test/shared/simple_nested_test.rs b/rust/test/shared/simple_nested_test.rs
index a13d4b53f2..1fe175a8b4 100644
--- a/rust/test/shared/simple_nested_test.rs
+++ b/rust/test/shared/simple_nested_test.rs
@@ -80,6 +80,7 @@ fn test_nested_muts() {
             bool(): eq(false)
         })
     );
+    assert_that!(inner_msg.string_mut().get(), eq(""));
 
     set_and_test_mut!(inner_msg =>
         double_mut, 543.21;
@@ -93,9 +94,9 @@ fn test_nested_muts() {
         fixed32_mut, 77;
         fixed64_mut, 999;
         bool_mut, true;
-
+        string_mut, "hi";
     );
-    // TODO: add mutation tests for strings and bytes
+    // TODO: add mutation test for bytes
 }
 
 #[test]
diff --git a/src/google/protobuf/compiler/rust/message.cc b/src/google/protobuf/compiler/rust/message.cc
index c98b6f409e..7a73e93101 100644
--- a/src/google/protobuf/compiler/rust/message.cc
+++ b/src/google/protobuf/compiler/rust/message.cc
@@ -172,6 +172,7 @@ void GetterForViewOrMut(Context<FieldDescriptor> field, bool is_mut) {
   // If we're dealing with a Mut, the getter must be supplied
   // self.inner.msg() whereas a View has to be supplied self.msg
   auto self = is_mut ? "self.inner.msg()" : "self.msg";
+
   if (fieldType == FieldDescriptor::TYPE_MESSAGE) {
     Context<Descriptor> d = field.WithDesc(field.desc().message_type());
     // TODO: support messages which are defined in other crates.
@@ -217,18 +218,44 @@ void GetterForViewOrMut(Context<FieldDescriptor> field, bool is_mut) {
 
   auto rsType = PrimitiveRsTypeName(field.desc());
   if (fieldType == FieldDescriptor::TYPE_STRING) {
-    field.Emit(
-        {
-            {"field", fieldName},
-            {"self", self},
-            {"getter_thunk", getter_thunk},
-            {"RsType", rsType},
-        },
-        R"rs(
-              pub fn r#$field$(&self) -> &$RsType$ {
+    field.Emit({{"field", fieldName},
+                {"self", self},
+                {"getter_thunk", getter_thunk},
+                {"setter_thunk", setter_thunk},
+                {"RsType", rsType},
+                {"maybe_mutator",
+                 [&] {
+                   if (is_mut) {
+                     field.Emit({}, R"rs(
+                    pub fn r#$field$_mut(&self) -> $pb$::Mut<'_, $RsType$> {
+                       static VTABLE: $pbi$::BytesMutVTable = 
+                        $pbi$::BytesMutVTable::new(
+                          $pbi$::Private,
+                          $getter_thunk$,
+                          $setter_thunk$,
+                        );
+
+                       unsafe {
+                        <$pb$::Mut<$RsType$>>::from_inner(
+                          $pbi$::Private,
+                          $pbi$::RawVTableMutator::new(
+                            $pbi$::Private,
+                            self.inner,
+                            &VTABLE,
+                           )
+                        )
+                      }
+                    }
+                    )rs");
+                   }
+                 }}},
+               R"rs(
+              pub fn r#$field$(&self) -> $pb$::View<'_, $RsType$> {
                 let s = unsafe { $getter_thunk$($self$).as_ref() };
                 unsafe { __pb::ProtoStr::from_utf8_unchecked(s) }
               }
+
+              $maybe_mutator$
             )rs");
   } else if (fieldType == FieldDescriptor::TYPE_BYTES) {
     field.Emit(
@@ -239,7 +266,7 @@ void GetterForViewOrMut(Context<FieldDescriptor> field, bool is_mut) {
             {"RsType", rsType},
         },
         R"rs(
-              pub fn r#$field$(&self) -> &$RsType$ {
+              pub fn r#$field$(&self) -> $pb$::View<'_, $RsType$> {
                 unsafe { $getter_thunk$($self$).as_ref() }
               }
             )rs");
@@ -275,7 +302,7 @@ void GetterForViewOrMut(Context<FieldDescriptor> field, bool is_mut) {
                    }
                  }}},
                R"rs(
-            pub fn r#$field$(&self) -> $RsType$ {
+            pub fn r#$field$(&self) -> $pb$::View<'_, $RsType$> {
               unsafe { $getter_thunk$($self$) }
             }