Adds mutable extension accessors to the codegen for upb C and adds a `::protos::MutableExtension` method to upb C++ for getting a mutable extension.

PiperOrigin-RevId: 597088960

protos/protos.h

@@ -18,6 +18,7 @@
 #include "upb/message/copy.h"
 #include "upb/message/internal/accessors.h"
 #include "upb/message/internal/extension.h"
+#include "upb/message/promote.h"
 #include "upb/mini_table/extension.h"
 #include "upb/wire/decode.h"
 #include "upb/wire/encode.h"
@@ -429,6 +430,25 @@ absl::StatusOr<Ptr<const Extension>> GetExtension(
   return GetExtension(protos::Ptr(message), id);
 }
 
+template <int&..., typename T, typename Extendee, typename Extension,
+          typename = EnableIfProtosClass<T>, typename = EnableIfMutableProto<T>>
+absl::StatusOr<Ptr<Extension>> MutableExtension(
+    Ptr<T> message,
+    const ::protos::internal::ExtensionIdentifier<Extendee, Extension>& id) {
+  upb_Arena* arena = ::protos::internal::GetArena(message);
+  upb_Extension* ext;
+  upb_GetExtension_Status status = upb_MiniTable_GetOrPromoteOrCreateExtension(
+      internal::GetInternalMsg(message), id.mini_table_ext(), 0, arena, &ext);
+  if (status == kUpb_GetExtension_OutOfMemory) {
+    return MessageAllocationError();
+  } else if (status != kUpb_GetExtension_Ok) {
+    return ExtensionNotFoundError(
+        upb_MiniTableExtension_Number(id.mini_table_ext()));
+  }
+  return Ptr<Extension>(::protos::internal::CreateMessageProxy<Extension>(
+      static_cast<upb_Message*>(ext->data.ptr), arena));
+}
+
 template <typename T>
 ABSL_MUST_USE_RESULT bool Parse(Ptr<T> message, absl::string_view bytes) {
   static_assert(!std::is_const_v<T>);

protos_generator/tests/test_generated.cc

@@ -824,6 +824,47 @@ TEST(CppGeneratedCode, GetExtensionOnImmutableChild) {
             ::protos::GetExtension(recursive_child, theme).value()->ext_name());
 }
 
+TEST(CppGeneratedCode, MutableExtension) {
+  ::protos::Arena arena;
+  ::protos::Ptr<TestModel> model = ::protos::CreateMessage<TestModel>(arena);
+  EXPECT_EQ(false, ::protos::HasExtension(model, theme));
+  absl::StatusOr<::protos::Ptr<ThemeExtension>> extension1 =
+      ::protos::MutableExtension(model, theme);
+  EXPECT_EQ(true, extension1.ok());
+  (*extension1)->set_ext_name("Hello World");
+  EXPECT_EQ("Hello World",
+            ::protos::GetExtension(model, theme).value()->ext_name());
+}
+
+TEST(CppGeneratedCode, MutableExtensionOnMutableChild) {
+  TestModel model;
+  ::protos::Ptr<TestModel> mutable_recursive_child =
+      model.mutable_recursive_child();
+  EXPECT_EQ(false, ::protos::HasExtension(mutable_recursive_child, theme));
+  absl::StatusOr<::protos::Ptr<ThemeExtension>> extension1 =
+      ::protos::MutableExtension(mutable_recursive_child, theme);
+  EXPECT_EQ(true, extension1.ok());
+  (*extension1)->set_ext_name("Hello World");
+  EXPECT_EQ("Hello World",
+            ::protos::GetExtension(mutable_recursive_child, theme)
+                .value()
+                ->ext_name());
+}
+
+TEST(CppGeneratedCode, MutableExtensionOnImmutableChild) {
+  TestModel model;
+  ::protos::Ptr<TestModel> mutable_recursive_child =
+      model.mutable_recursive_child();
+  EXPECT_EQ(false, ::protos::HasExtension(mutable_recursive_child, theme));
+  absl::StatusOr<::protos::Ptr<ThemeExtension>> extension1 =
+      ::protos::MutableExtension(mutable_recursive_child, theme);
+  EXPECT_EQ(true, extension1.ok());
+  (*extension1)->set_ext_name("Hello World");
+  ::protos::Ptr<const TestModel> recursive_child = model.recursive_child();
+  EXPECT_EQ("Hello World",
+            ::protos::GetExtension(recursive_child, theme).value()->ext_name());
+}
+
 TEST(CppGeneratedCode, SerializeUsingArena) {
   TestModel model;
   model.set_str1("Hello World");

upb/message/internal/accessors.h

@@ -259,6 +259,32 @@ UPB_INLINE void _upb_Message_GetExtensionField(
   }
 }
 
+// Gets a extension message or creates a default message and sets the extension
+// if it doesn't already exist.
+UPB_INLINE bool _upb_Message_GetOrCreateExtensionSubmessage(
+    struct upb_Message* msg, const upb_MiniTableExtension* mt_ext,
+    struct upb_Message** val, upb_Arena* a) {
+  const upb_MiniTableField* f = &mt_ext->UPB_PRIVATE(field);
+  UPB_ASSUME(upb_MiniTableField_IsExtension(f));
+  const struct upb_Extension* const_ext = _upb_Message_Getext(msg, mt_ext);
+  if (const_ext) {
+    // Extension exists, get a mutable version of it.
+    struct upb_Extension* ext =
+        _upb_Message_GetOrCreateExtension(msg, mt_ext, a);
+    *val = (struct upb_Message*)ext->data.ptr;
+    return true;
+  }
+  // Extension doesn't exist, create a new message and set it.
+  struct upb_Message* ext_msg =
+      _upb_Message_New(upb_MiniTableExtension_GetSubMessage(mt_ext), a);
+  if (!ext_msg) return false;
+  struct upb_Extension* ext = _upb_Message_GetOrCreateExtension(msg, mt_ext, a);
+  if (!ext) return false;
+  ext->data.ptr = ext_msg;
+  *val = ext_msg;
+  return true;
+}
+
 UPB_INLINE void _upb_Message_SetNonExtensionField(
     struct upb_Message* msg, const upb_MiniTableField* field, const void* val) {
   UPB_ASSUME(!upb_MiniTableField_IsExtension(field));

upb/message/promote.c

@@ -64,17 +64,11 @@ static upb_UnknownToMessageRet upb_MiniTable_ParseUnknownMessage(
   return ret;
 }
 
-upb_GetExtension_Status upb_MiniTable_GetOrPromoteExtension(
+// Promotes unknown data to an extension by finding the field number in unknown
+// data and decoding it.
+static upb_GetExtension_Status upb_MiniTable_PromoteExtension(
     upb_Message* msg, const upb_MiniTableExtension* ext_table,
-    int decode_options, upb_Arena* arena, const upb_Extension** extension) {
-  UPB_ASSERT(upb_MiniTableField_CType(upb_MiniTableExtension_AsField(
-                 ext_table)) == kUpb_CType_Message);
-  *extension = _upb_Message_Getext(msg, ext_table);
-  if (*extension) {
-    return kUpb_GetExtension_Ok;
-  }
-
-  // Check unknown fields, if available promote.
+    int decode_options, upb_Arena* arena, upb_Extension** extension) {
   int field_number = upb_MiniTableExtension_Number(ext_table);
   upb_FindUnknownRet result = upb_Message_FindUnknown(msg, field_number, 0);
   if (result.status != kUpb_FindUnknown_Ok) {
@@ -98,19 +92,73 @@ upb_GetExtension_Status upb_MiniTable_GetOrPromoteExtension(
     case kUpb_UnknownToMessage_Ok:
       break;
   }
-  upb_Message* extension_msg = parse_result.message;
   // Add to extensions.
   upb_Extension* ext = _upb_Message_GetOrCreateExtension(msg, ext_table, arena);
   if (!ext) {
     return kUpb_GetExtension_OutOfMemory;
   }
-  memcpy(&ext->data, &extension_msg, sizeof(extension_msg));
+  ext->data.ptr = parse_result.message;
   *extension = ext;
   const char* delete_ptr = upb_Message_GetUnknown(msg, &len) + ofs;
   upb_Message_DeleteUnknown(msg, delete_ptr, result.len);
   return kUpb_GetExtension_Ok;
 }
 
+upb_GetExtension_Status upb_MiniTable_GetOrPromoteExtension(
+    upb_Message* msg, const upb_MiniTableExtension* ext_table,
+    int decode_options, upb_Arena* arena, const upb_Extension** extension) {
+  UPB_ASSERT(upb_MiniTableField_CType(upb_MiniTableExtension_AsField(
+                 ext_table)) == kUpb_CType_Message);
+  *extension = _upb_Message_Getext(msg, ext_table);
+  if (*extension) {
+    return kUpb_GetExtension_Ok;
+  }
+
+  // Check unknown fields, if available promote.
+  upb_Extension* ext;
+  upb_GetExtension_Status promote_result = upb_MiniTable_PromoteExtension(
+      msg, ext_table, decode_options, arena, &ext);
+  if (promote_result == kUpb_GetExtension_Ok) {
+    *extension = ext;
+  }
+  return promote_result;
+}
+
+upb_GetExtension_Status upb_MiniTable_GetOrPromoteOrCreateExtension(
+    upb_Message* msg, const upb_MiniTableExtension* ext_table,
+    int decode_options, upb_Arena* arena, upb_Extension** extension) {
+  UPB_ASSERT(upb_MiniTableField_CType(upb_MiniTableExtension_AsField(
+                 ext_table)) == kUpb_CType_Message);
+  const upb_Extension* const_extension = _upb_Message_Getext(msg, ext_table);
+  if (const_extension) {
+    // Extension exists on the message, get the mutable version of it.
+    *extension = _upb_Message_GetOrCreateExtension(msg, ext_table, arena);
+    return kUpb_GetExtension_Ok;
+  }
+
+  // Check unknown fields, if available promote.
+  upb_GetExtension_Status promote_result = upb_MiniTable_PromoteExtension(
+      msg, ext_table, decode_options, arena, extension);
+  if (promote_result != kUpb_GetExtension_NotPresent) {
+    return promote_result;
+  }
+
+  // Extension not found, create a new default message.
+  const upb_MiniTable* extension_table =
+      upb_MiniTableExtension_GetSubMessage(ext_table);
+  upb_Message* extension_msg = upb_Message_New(extension_table, arena);
+  if (!extension_msg) {
+    return kUpb_GetExtension_OutOfMemory;
+  }
+  upb_Extension* ext = _upb_Message_GetOrCreateExtension(msg, ext_table, arena);
+  if (!ext) {
+    return kUpb_GetExtension_OutOfMemory;
+  }
+  ext->data.ptr = extension_msg;
+  *extension = ext;
+  return kUpb_GetExtension_Ok;
+}
+
 static upb_FindUnknownRet upb_FindUnknownRet_ParseError(void) {
   return (upb_FindUnknownRet){.status = kUpb_FindUnknown_ParseError};
 }

upb/message/promote.h

@@ -42,6 +42,15 @@ upb_GetExtension_Status upb_MiniTable_GetOrPromoteExtension(
     upb_Message* msg, const upb_MiniTableExtension* ext_table,
     int decode_options, upb_Arena* arena, const upb_Extension** extension);
 
+// Returns a mutable message extension, promotes an unknown field to a mutable
+// extension, or creates a new extension.
+//
+// TODO Only supports extension fields that are messages,
+// expand support to include non-message types.
+upb_GetExtension_Status upb_MiniTable_GetOrPromoteOrCreateExtension(
+    upb_Message* msg, const upb_MiniTableExtension* ext_table,
+    int decode_options, upb_Arena* arena, upb_Extension** extension);
+
 typedef enum {
   kUpb_FindUnknown_Ok,
   kUpb_FindUnknown_NotPresent,

upb/message/test.cc

@@ -104,6 +104,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) {

upb_generator/protoc-gen-upb.cc

@@ -27,6 +27,7 @@
 #include "absl/strings/string_view.h"
 #include "absl/strings/substitute.h"
 #include "upb/base/descriptor_constants.h"
+#include "upb/base/status.hpp"
 #include "upb/base/string_view.h"
 #include "upb/reflection/def.hpp"
 #include "upb/wire/types.h"
@@ -256,7 +257,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) {
@@ -308,6 +309,31 @@ 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) {
+              const upb_MiniTableExtension* ext = &$4;
+              UPB_ASSUME(upb_MiniTableField_IsScalar(&ext->UPB_PRIVATE(field)));
+              UPB_ASSUME(upb_MiniTableField_IsSubMessage(&ext->UPB_PRIVATE(field)));
+              UPB_ASSUME(UPB_PRIVATE(_upb_MiniTableField_GetRep)(
+                             &ext->UPB_PRIVATE(field)) == $5);
+              upb_Message* ret;
+              bool ok = _upb_Message_GetOrCreateExtensionSubmessage(
+                  (upb_Message*)msg, ext, &ret, arena);
+              if (!ok) return NULL;
+              return ($0*)ret;
+            }
+          )cc",
+          MessageName(ext.message_type()), ExtensionIdentBase(ext), ext.name(),
+          MessageName(ext.containing_type()), ExtensionLayout(ext),
+          GetFieldRep(pools, ext));
+    }
   }
 }
 
@@ -927,7 +953,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" ||