From 1f5df18680a5084fc9b4e4150b22f0b388d678dd Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Tue, 29 Oct 2024 21:18:17 -0700
Subject: [PATCH] Allow aliasing submessages when owning arenas are fused

PiperOrigin-RevId: 691272284
---
 hpb_generator/gen_accessors.cc        |  5 +++--
 hpb_generator/gen_repeated_fields.cc  |  2 +-
 hpb_generator/tests/test_generated.cc | 20 ++++++++++++++++++++
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/hpb_generator/gen_accessors.cc b/hpb_generator/gen_accessors.cc
index 9a70f6aef0..d038374963 100644
--- a/hpb_generator/gen_accessors.cc
+++ b/hpb_generator/gen_accessors.cc
@@ -115,7 +115,8 @@ void WriteFieldAccessorsInHeader(const protobuf::Descriptor* desc,
                           * Re-points submessage to the given target.
                           *
                           * REQUIRES:
-                          * - both messages must be in the same arena.
+                          * - both messages must be in the same arena, or in two
+                          * fused arenas.
                           */
                          void set_alias_$2($0 target);
                        )cc",
@@ -282,7 +283,7 @@ void WriteAccessorsInSource(const protobuf::Descriptor* desc, Context& ctx) {
                     (upb_Message*)($3_mutable_$5(msg_, $6)), $6);
               }
               void $0::set_alias_$2($1 target) {
-                ABSL_CHECK_EQ(arena_, hpb::interop::upb::GetArena(target));
+                ABSL_CHECK(upb_Arena_IsFused(arena_, hpb::interop::upb::GetArena(target)));
                 upb_Message_SetBaseFieldMessage(
                     UPB_UPCAST(msg_),
                     upb_MiniTable_FindFieldByNumber($7::minitable(), $8),
diff --git a/hpb_generator/gen_repeated_fields.cc b/hpb_generator/gen_repeated_fields.cc
index 0b2650d758..bbc6fa4537 100644
--- a/hpb_generator/gen_repeated_fields.cc
+++ b/hpb_generator/gen_repeated_fields.cc
@@ -161,7 +161,7 @@ void WriteRepeatedMessageAccessor(const protobuf::Descriptor* desc,
         }
 
         bool $0::add_alias_$2($1 target) {
-          ABSL_CHECK_EQ(arena_, hpb::interop::upb::GetArena(target));
+          ABSL_CHECK(upb_Arena_IsFused(arena_, hpb::interop::upb::GetArena(target)));
           size_t size = 0;
           $3_$2(msg_, &size);
           auto elements = $3_resize_$2(msg_, size + 1, arena_);
diff --git a/hpb_generator/tests/test_generated.cc b/hpb_generator/tests/test_generated.cc
index ec2cb102ae..f9fb8132f7 100644
--- a/hpb_generator/tests/test_generated.cc
+++ b/hpb_generator/tests/test_generated.cc
@@ -1310,6 +1310,26 @@ TEST(CppGeneratedCode, SetAliasFailsForDifferentArena) {
   EXPECT_DEATH(parent.set_alias_child(child), "hpb::interop::upb::GetArena");
 }
 
+TEST(CppGeneratedCode, SetAliasSucceedsForDifferentArenaFused) {
+  hpb::Arena arena;
+  auto parent1 = hpb::CreateMessage<Parent>(arena);
+  auto child = parent1.mutable_child();
+  child->set_peeps(12);
+
+  hpb::Arena other_arena;
+  auto parent2 = hpb::CreateMessage<Parent>(other_arena);
+  arena.Fuse(other_arena);
+
+  parent2.set_alias_child(child);
+
+  ASSERT_EQ(parent1.child()->peeps(), parent2.child()->peeps());
+  ASSERT_EQ(hpb::interop::upb::GetMessage(parent1.child()),
+            hpb::interop::upb::GetMessage(parent2.child()));
+  auto childPtr = hpb::Ptr<Child>(child);
+  ASSERT_EQ(hpb::interop::upb::GetMessage(childPtr),
+            hpb::interop::upb::GetMessage(parent1.child()));
+}
+
 TEST(CppGeneratedCode, SetAliasRepeated) {
   hpb::Arena arena;
   auto child = hpb::CreateMessage<Child>(arena);