Reserve/InternalExtend use capacity instead of size as input.

This allows to remove branch from `InternalExtend` which is beneficial for small fields.

Also refactored `UnsafeArenaAddAllocated` - it looks like it had an outdated comment. The function no longer deletes already allocated elements and hence does not need to be a template.

PiperOrigin-RevId: 562530131

src/google/protobuf/generated_message_reflection.cc

@@ -2484,7 +2484,7 @@ Message* Reflection::AddMessage(Message* message, const FieldDescriptor* field,
       // 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(result);
     }
 
     return result;
@@ -2526,7 +2526,7 @@ void Reflection::UnsafeArenaAddAllocatedMessage(Message* message,
     } else {
       repeated = MutableRaw<RepeatedPtrFieldBase>(message, field);
     }
-    repeated->UnsafeArenaAddAllocated<GenericTypeHandler<Message>>(new_entry);
+    repeated->UnsafeArenaAddAllocated(new_entry);
   }
 }
 

src/google/protobuf/repeated_field_unittest.cc

@@ -1694,8 +1694,8 @@ TEST(RepeatedPtrField, AddAllocated) {
   std::string* moo = new std::string("moo");
   field.AddAllocated(moo);
   EXPECT_EQ(index + 1, field.size());
-  // We should have discarded the cleared object.
-  EXPECT_EQ(0, field.ClearedCount());
+  // We should have retained the cleared object.
+  EXPECT_EQ(1, field.ClearedCount());
   EXPECT_EQ(moo, &field.Get(index));
 }
 

src/google/protobuf/repeated_ptr_field.cc

@@ -52,46 +52,52 @@ namespace protobuf {
 
 namespace internal {
 
-void** RepeatedPtrFieldBase::InternalExtend(int extend_amount) {
-  int new_size = current_size_ + extend_amount;
-  if (total_size_ >= new_size) {
-    // N.B.: rep_ is non-nullptr because extend_amount is always > 0, hence
-    // total_size must be non-zero since it is lower-bounded by new_size.
-    return elements() + current_size_;
+void RepeatedPtrFieldBase::UnsafeArenaAddAllocated(void* value) {
+  MaybeExtend();
+  void** elems = elements();
+  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.
+    elems[allocated_size_] = elems[current_size_];
   }
+  ++allocated_size_;
+  elems[ExchangeCurrentSize(current_size_ + 1)] = value;
+}
+
+void RepeatedPtrFieldBase::InternalExtend(int extend_amount) {
+  ABSL_DCHECK(extend_amount > 0);
   constexpr size_t ptr_size = sizeof(rep()->elements[0]);
+  int new_capacity = Capacity() + extend_amount;
   Arena* arena = GetOwningArena();
-  new_size = internal::CalculateReserveSize<void*, kRepHeaderSize>(total_size_,
-                                                                   new_size);
+  new_capacity = internal::CalculateReserveSize<void*, kRepHeaderSize>(
+      Capacity(), new_capacity);
   ABSL_CHECK_LE(
-      static_cast<int64_t>(new_size),
+      static_cast<int64_t>(new_capacity),
       static_cast<int64_t>(
           (std::numeric_limits<size_t>::max() - kRepHeaderSize) / ptr_size))
       << "Requested size is too large to fit into size_t.";
-  size_t bytes = kRepHeaderSize + ptr_size * new_size;
+  size_t bytes = kRepHeaderSize + ptr_size * new_capacity;
   Rep* new_rep;
   void* old_tagged_ptr = tagged_rep_or_elem_;
   if (arena == nullptr) {
     internal::SizedPtr res = internal::AllocateAtLeast(bytes);
-    new_size = static_cast<int>((res.n - kRepHeaderSize) / ptr_size);
     new_rep = reinterpret_cast<Rep*>(res.p);
+    new_rep->capacity = static_cast<int>((res.n - kRepHeaderSize) / ptr_size);
   } else {
     new_rep = reinterpret_cast<Rep*>(Arena::CreateArray<char>(arena, bytes));
+    new_rep->capacity = new_capacity;
   }
 
   if (using_sso()) {
-    new_rep->allocated_size = old_tagged_ptr != nullptr ? 1 : 0;
     new_rep->elements[0] = old_tagged_ptr;
   } else {
     Rep* old_rep =
         reinterpret_cast<Rep*>(reinterpret_cast<uintptr_t>(old_tagged_ptr) - 1);
-    if (old_rep->allocated_size > 0) {
-      memcpy(new_rep->elements, old_rep->elements,
-             old_rep->allocated_size * ptr_size);
+    if (allocated_size_ > 0) {
+      memcpy(new_rep->elements, old_rep->elements, allocated_size_ * ptr_size);
     }
-    new_rep->allocated_size = old_rep->allocated_size;
 
-    size_t old_size = total_size_ * ptr_size + kRepHeaderSize;
+    size_t old_size = Capacity() * ptr_size + kRepHeaderSize;
     if (arena == nullptr) {
       internal::SizedDelete(old_rep, old_size);
     } else {
@@ -101,88 +107,48 @@ void** RepeatedPtrFieldBase::InternalExtend(int extend_amount) {
 
   tagged_rep_or_elem_ =
       reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(new_rep) + 1);
-  total_size_ = new_size;
-  return &new_rep->elements[current_size_];
 }
 
-void RepeatedPtrFieldBase::Reserve(int new_size) {
-  if (new_size > current_size_) {
-    InternalExtend(new_size - current_size_);
+void RepeatedPtrFieldBase::Reserve(int capacity) {
+  int n = Capacity();
+  if (capacity > n) {
+    InternalExtend(capacity - n);
   }
 }
 
 void RepeatedPtrFieldBase::DestroyProtos() {
-  ABSL_DCHECK(tagged_rep_or_elem_);
   ABSL_DCHECK(arena_ == nullptr);
-  if (using_sso()) {
-    delete static_cast<MessageLite*>(tagged_rep_or_elem_);
-
-  } else {
+  void** elems = elements();
+  for (int i = 0; i < allocated_size_; i++) {
+    delete static_cast<MessageLite*>(elems[i]);
+  }
+  if (!using_sso()) {
     Rep* r = rep();
-    int n = r->allocated_size;
-    void* const* elements = r->elements;
-    for (int i = 0; i < n; i++) {
-      delete static_cast<MessageLite*>(elements[i]);
-    }
-    const size_t size = total_size_ * sizeof(elements[0]) + kRepHeaderSize;
-    internal::SizedDelete(r, size);
-    tagged_rep_or_elem_ = nullptr;
+    internal::SizedDelete(r, r->capacity * sizeof(elems[0]) + kRepHeaderSize);
   }
 }
 
 void* RepeatedPtrFieldBase::AddOutOfLineHelper(void* obj) {
-  if (tagged_rep_or_elem_ == nullptr) {
-    ABSL_DCHECK_EQ(current_size_, 0);
-    ABSL_DCHECK(using_sso());
-    ABSL_DCHECK_EQ(allocated_size(), 0);
-    ExchangeCurrentSize(1);
-    tagged_rep_or_elem_ = obj;
-    return obj;
-  }
-  if (using_sso() || rep()->allocated_size == total_size_) {
-    InternalExtend(1);  // Equivalent to "Reserve(total_size_ + 1)"
-  }
-  Rep* r = rep();
-  ++r->allocated_size;
-  r->elements[ExchangeCurrentSize(current_size_ + 1)] = obj;
-  return obj;
+  return AddAllocatedForParse(obj);
 }
 
 void RepeatedPtrFieldBase::CloseGap(int start, int num) {
-  if (using_sso()) {
-    if (start == 0 && num == 1) {
-      tagged_rep_or_elem_ = nullptr;
-    }
-  } else {
     // Close up a gap of "num" elements starting at offset "start".
-    Rep* r = rep();
-    for (int i = start + num; i < r->allocated_size; ++i)
-      r->elements[i - num] = r->elements[i];
-    r->allocated_size -= num;
-  }
+  void** elems = elements();
+  for (int i = start + num; i < allocated_size_; ++i) elems[i - num] = elems[i];
+  allocated_size_ -= num;
   ExchangeCurrentSize(current_size_ - num);
 }
 
 MessageLite* RepeatedPtrFieldBase::AddWeak(const MessageLite* prototype) {
-  if (current_size_ < allocated_size()) {
-    return reinterpret_cast<MessageLite*>(
+  if (current_size_ < allocated_size_) {
+    return static_cast<MessageLite*>(
         element_at(ExchangeCurrentSize(current_size_ + 1)));
   }
-  if (allocated_size() == total_size_) {
-    Reserve(total_size_ + 1);
-  }
   MessageLite* result = prototype
                             ? prototype->New(arena_)
                             : Arena::CreateMessage<ImplicitWeakMessage>(arena_);
-  if (using_sso()) {
-    ExchangeCurrentSize(current_size_ + 1);
-    tagged_rep_or_elem_ = result;
-  } else {
-    Rep* r = rep();
-    ++r->allocated_size;
-    r->elements[ExchangeCurrentSize(current_size_ + 1)] = result;
-  }
-  return result;
+  return static_cast<MessageLite*>(AddAllocatedForParse(result));
 }
 
 void InternalOutOfLineDeleteMessageLite(MessageLite* message) {