Pass down `Arena*` to `UnsafeArenaAddAllocated`.
PiperOrigin-RevId: 807460420
diff --git a/src/google/protobuf/generated_message_reflection.cc b/src/google/protobuf/generated_message_reflection.cc
index 0186341..5b7749f 100644
--- a/src/google/protobuf/generated_message_reflection.cc
+++ b/src/google/protobuf/generated_message_reflection.cc
@@ -2719,7 +2719,8 @@
// We can guarantee here that repeated and result are either both heap
// allocated or arena owned. So it is safe to call the unsafe version
// of AddAllocated.
- repeated->UnsafeArenaAddAllocated<GenericTypeHandler<Message> >(result);
+ repeated->UnsafeArenaAddAllocated<GenericTypeHandler<Message>>(
+ message->GetArena(), result);
}
return result;
@@ -2763,7 +2764,8 @@
} else {
repeated = MutableRaw<RepeatedPtrFieldBase>(message, field);
}
- repeated->UnsafeArenaAddAllocated<GenericTypeHandler<Message>>(new_entry);
+ repeated->UnsafeArenaAddAllocated<GenericTypeHandler<Message>>(
+ message->GetArena(), new_entry);
SetHasBitForRepeated(message, field);
}
}
diff --git a/src/google/protobuf/repeated_ptr_field.h b/src/google/protobuf/repeated_ptr_field.h
index 2189456..6d7d082 100644
--- a/src/google/protobuf/repeated_ptr_field.h
+++ b/src/google/protobuf/repeated_ptr_field.h
@@ -476,12 +476,13 @@
}
template <typename TypeHandler>
- void UnsafeArenaAddAllocated(Value<TypeHandler>* value) {
+ void UnsafeArenaAddAllocated(Arena* arena, Value<TypeHandler>* value) {
+ ABSL_DCHECK_EQ(arena, GetArena());
ABSL_DCHECK_NE(value, nullptr);
// Make room for the new pointer.
if (SizeAtCapacity()) {
// The array is completely full with no cleared objects, so grow it.
- InternalExtend(1, GetArena());
+ InternalExtend(1, arena);
++rep()->allocated_size;
} else if (AllocatedSizeAtCapacity()) {
// There is no more space in the pointer array because it contains some
@@ -489,7 +490,7 @@
// this case because otherwise a loop calling AddAllocated() followed by
// Clear() would leak memory.
using H = CommonHandler<TypeHandler>;
- Delete<H>(element_at(current_size_), GetArena());
+ Delete<H>(element_at(current_size_), arena);
} else if (current_size_ < allocated_size()) {
// We have some cleared objects. We don't care about their order, so we
// can just move the first one to the end to make space.
@@ -548,6 +549,8 @@
// or load (mine).
Value<TypeHandler>* value, Arena* value_arena, Arena* my_arena) {
using H = CommonHandler<TypeHandler>;
+ ABSL_DCHECK_EQ(my_arena, GetArena());
+ ABSL_DCHECK_EQ(value_arena, TypeHandler::GetArena(value));
// Ensure that either the value is in the same arena, or if not, we do the
// appropriate thing: Own() it (if it's on heap and we're in an arena) or
// copy it to our arena/heap (otherwise).
@@ -558,7 +561,7 @@
value = cast<TypeHandler>(CloneSlow(my_arena, *value));
}
- UnsafeArenaAddAllocated<H>(value);
+ UnsafeArenaAddAllocated<H>(my_arena, value);
}
// TODO - Outline this function so a future change can use a
@@ -1680,7 +1683,7 @@
template <typename Element>
inline void RepeatedPtrField<Element>::UnsafeArenaAddAllocated(Element* value) {
- RepeatedPtrFieldBase::UnsafeArenaAddAllocated<TypeHandler>(value);
+ RepeatedPtrFieldBase::UnsafeArenaAddAllocated<TypeHandler>(GetArena(), value);
}
template <typename Element>
diff --git a/src/google/protobuf/repeated_ptr_field_unittest.cc b/src/google/protobuf/repeated_ptr_field_unittest.cc
index 62e07be..874141d 100644
--- a/src/google/protobuf/repeated_ptr_field_unittest.cc
+++ b/src/google/protobuf/repeated_ptr_field_unittest.cc
@@ -630,7 +630,7 @@
ElemT* concrete_new_elem = static_cast<ElemT*>(base_new_elem);
concrete_new_elem->set_bb(456);
base_field->UnsafeArenaAddAllocated<GenericTypeHandler<Message>>(
- base_new_elem);
+ &arena, base_new_elem);
Message* base_new_elem_roundtrip =
base_field->UnsafeArenaReleaseLast<GenericTypeHandler<Message>>();
ASSERT_NE(base_new_elem_roundtrip, nullptr);
