Allow aliasing submessages when owning arenas are fused

PiperOrigin-RevId: 691272284
protocolbuffers · Oct 30, 2024 · 1f5df18 · 1f5df18
1 parent bd38761
commit 1f5df18
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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);