From 03ff7af37ec2289480e62b82484ac8e47f910e1d Mon Sep 17 00:00:00 2001
From: Joshua Haberman <haberman@google.com>
Date: Mon, 22 Jan 2024 15:30:15 -0800
Subject: [PATCH] Refactored a few protos:: methods to no longer access
 `upb_Extension` directly.

PiperOrigin-RevId: 600588767
---
 protos/protos.cc | 32 +++++++++++++-------------------
 1 file changed, 13 insertions(+), 19 deletions(-)

diff --git a/protos/protos.cc b/protos/protos.cc
index 6522405c12..80d72155e1 100644
--- a/protos/protos.cc
+++ b/protos/protos.cc
@@ -163,22 +163,15 @@ upb_Message* DeepClone(const upb_Message* source,
 absl::Status MoveExtension(upb_Message* message, upb_Arena* message_arena,
                            const upb_MiniTableExtension* ext,
                            upb_Message* extension, upb_Arena* extension_arena) {
-  upb_Extension* msg_ext =
-      _upb_Message_GetOrCreateExtension(message, ext, message_arena);
-  if (!msg_ext) {
-    return MessageAllocationError();
-  }
-  if (message_arena != extension_arena) {
-    // Try fuse, if fusing is not allowed or fails, create copy of extension.
-    if (!upb_Arena_Fuse(message_arena, extension_arena)) {
-      msg_ext->data.ptr = DeepClone(
-          extension, upb_MiniTableExtension_GetSubMessage(msg_ext->ext),
-          message_arena);
-      return absl::OkStatus();
-    }
+  if (message_arena != extension_arena &&
+      // Try fuse, if fusing is not allowed or fails, create copy of extension.
+      !upb_Arena_Fuse(message_arena, extension_arena)) {
+    extension = DeepClone(extension, upb_MiniTableExtension_GetSubMessage(ext),
+                          message_arena);
   }
-  msg_ext->data.ptr = extension;
-  return absl::OkStatus();
+  return _upb_Message_SetExtensionField(message, ext, &extension, message_arena)
+             ? absl::OkStatus()
+             : MessageAllocationError();
 }
 
 absl::Status SetExtension(upb_Message* message, upb_Arena* message_arena,
@@ -190,10 +183,11 @@ absl::Status SetExtension(upb_Message* message, upb_Arena* message_arena,
     return MessageAllocationError();
   }
   // Clone extension into target message arena.
-  msg_ext->data.ptr =
-      DeepClone(extension, upb_MiniTableExtension_GetSubMessage(msg_ext->ext),
-                message_arena);
-  return absl::OkStatus();
+  extension = DeepClone(extension, upb_MiniTableExtension_GetSubMessage(ext),
+                        message_arena);
+  return _upb_Message_SetExtensionField(message, ext, &extension, message_arena)
+             ? absl::OkStatus()
+             : MessageAllocationError();
 }
 
 }  // namespace internal