From 0c9389e916eb69a4c21ebacfe39227076622ebc6 Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Tue, 5 Mar 2024 17:04:21 -0800
Subject: [PATCH] Adds mutable extension accessors to the codegen for upb C for
 getting mutable extension messages.

PiperOrigin-RevId: 613017999
---
 upb/message/test.cc             |  7 +++++++
 upb_generator/protoc-gen-upb.cc | 25 +++++++++++++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/upb/message/test.cc b/upb/message/test.cc
index 236b973184a14..433632db14a85 100644
--- a/upb/message/test.cc
+++ b/upb/message/test.cc
@@ -107,6 +107,13 @@ TEST(MessageTest, Extensions) {
                              defpool.ptr(), 0, arena.ptr(), status.ptr()))
       << status.error_message();
   VerifyMessage(ext_msg3);
+
+  // Test setters and mutable accessors
+  upb_test_TestExtensions* ext_msg4 = upb_test_TestExtensions_new(arena.ptr());
+  upb_test_TestExtensions_set_optional_int32_ext(ext_msg4, 123, arena.ptr());
+  protobuf_test_messages_proto3_TestAllTypesProto3_set_optional_int32(
+      upb_test_mutable_optional_msg_ext(ext_msg4, arena.ptr()), 456);
+  VerifyMessage(ext_msg4);
 }
 
 void VerifyMessageSet(const upb_test_TestMessageSet* mset_msg) {
diff --git a/upb_generator/protoc-gen-upb.cc b/upb_generator/protoc-gen-upb.cc
index 871819dd78ee0..2b0e9c30cbb17 100644
--- a/upb_generator/protoc-gen-upb.cc
+++ b/upb_generator/protoc-gen-upb.cc
@@ -259,7 +259,7 @@ std::string GetFieldRep(const DefPoolPair& pools, upb::FieldDefPtr field) {
 }
 
 void GenerateExtensionInHeader(const DefPoolPair& pools, upb::FieldDefPtr ext,
-                               Output& output) {
+                               const Options& options, Output& output) {
   output(
       R"cc(
         UPB_INLINE bool $0_has_$1(const struct $2* msg) {
@@ -311,6 +311,27 @@ void GenerateExtensionInHeader(const DefPoolPair& pools, upb::FieldDefPtr ext,
         CTypeConst(ext), ExtensionIdentBase(ext), ext.name(),
         MessageName(ext.containing_type()), ExtensionLayout(ext),
         GetFieldRep(pools, ext));
+
+    // Message extensions also have a Msg_mutable_foo() accessor that will
+    // create the sub-message if it doesn't already exist.
+    if (ext.ctype() == kUpb_CType_Message &&
+        !UPB_DESC(MessageOptions_map_entry)(ext.containing_type().options())) {
+      output(
+          R"cc(
+            UPB_INLINE struct $0* $1_mutable_$2(struct $3* msg,
+                                                upb_Arena* arena) {
+              struct $0* sub = (struct $0*)$1_$2(msg);
+              if (sub == NULL) {
+                sub = (struct $0*)_upb_Message_New($4, arena);
+                if (sub) $1_set_$2(msg, sub, arena);
+              }
+              return sub;
+            }
+          )cc",
+          MessageName(ext.message_type()), ExtensionIdentBase(ext), ext.name(),
+          MessageName(ext.containing_type()),
+          MessageMiniTableRef(ext.message_type(), options));
+    }
   }
 }
 
@@ -930,7 +951,7 @@ void WriteHeader(const DefPoolPair& pools, upb::FileDefPtr file,
   }
 
   for (auto ext : this_file_exts) {
-    GenerateExtensionInHeader(pools, ext, output);
+    GenerateExtensionInHeader(pools, ext, options, output);
   }
 
   if (absl::string_view(file.name()) == "google/protobuf/descriptor.proto" ||