From 61662a536b9008602ee465abddc13c24bf8f8b47 Mon Sep 17 00:00:00 2001
From: Vitaly Goldshteyn <goldvitaly@google.com>
Date: Thu, 21 Nov 2024 04:51:45 -0800
Subject: [PATCH] Avoid recursive call to `Insert` in the flat case.

PiperOrigin-RevId: 698735648
---
 src/google/protobuf/extension_set.cc | 52 +++++++++++++++++++---------
 src/google/protobuf/extension_set.h  |  2 ++
 2 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/src/google/protobuf/extension_set.cc b/src/google/protobuf/extension_set.cc
index 070bed0e1e68..383731f005cb 100644
--- a/src/google/protobuf/extension_set.cc
+++ b/src/google/protobuf/extension_set.cc
@@ -21,6 +21,7 @@
 #include <type_traits>
 #include <utility>
 
+#include "absl/base/attributes.h"
 #include "absl/base/optimization.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/hash/hash.h"
@@ -1744,25 +1745,44 @@ ExtensionSet::Extension* ExtensionSet::FindOrNullInLargeMap(int key) {
       const_this->FindOrNullInLargeMap(key));
 }
 
+ABSL_ATTRIBUTE_NOINLINE
+std::pair<ExtensionSet::Extension*, bool>
+ExtensionSet::InternalInsertIntoLargeMap(int key) {
+  ABSL_DCHECK(is_large());
+  auto maybe = map_.large->insert({key, Extension()});
+  return {&maybe.first->second, maybe.second};
+}
+
 std::pair<ExtensionSet::Extension*, bool> ExtensionSet::Insert(int key) {
   if (ABSL_PREDICT_FALSE(is_large())) {
-    auto maybe = map_.large->insert({key, Extension()});
-    return {&maybe.first->second, maybe.second};
+    return InternalInsertIntoLargeMap(key);
+  }
+  uint16_t i = flat_size_;
+  KeyValue* flat = map_.flat;
+  // Iterating from the back to benefit the case where the keys are inserted in
+  // increasing order.
+  for (; i > 0; --i) {
+    int map_key = flat[i - 1].first;
+    if (map_key == key) {
+      return {&flat[i - 1].second, false};
+    }
+    if (map_key < key) {
+      break;
+    }
   }
-  KeyValue* end = flat_end();
-  KeyValue* it = flat_begin();
-  for (; it != end && it->first <= key; ++it) {
-    if (it->first == key) return {&it->second, false};
-  }
-  if (flat_size_ < flat_capacity_) {
-    std::copy_backward(it, end, end + 1);
-    ++flat_size_;
-    it->first = key;
-    it->second = Extension();
-    return {&it->second, true};
-  }
-  GrowCapacity(flat_size_ + 1);
-  return Insert(key);
+  if (flat_size_ == flat_capacity_) {
+    GrowCapacity(flat_size_ + 1);
+    if (ABSL_PREDICT_FALSE(is_large())) {
+      return InternalInsertIntoLargeMap(key);
+    }
+    flat = map_.flat;  // Reload flat pointer after GrowCapacity.
+  }
+
+  std::copy_backward(flat + i, flat + flat_size_, flat + flat_size_ + 1);
+  ++flat_size_;
+  flat[i].first = key;
+  flat[i].second = Extension();
+  return {&flat[i].second, true};
 }
 
 void ExtensionSet::GrowCapacity(size_t minimum_new_capacity) {
diff --git a/src/google/protobuf/extension_set.h b/src/google/protobuf/extension_set.h
index 346f5e42c5de..1a3b606055f4 100644
--- a/src/google/protobuf/extension_set.h
+++ b/src/google/protobuf/extension_set.h
@@ -817,6 +817,8 @@ class PROTOBUF_EXPORT ExtensionSet {
   // finds the already-existing Extension for that key (returns false).
   // The Extension* will point to the new-or-found Extension.
   std::pair<Extension*, bool> Insert(int key);
+  // Same as insert for the large map.
+  std::pair<Extension*, bool> InternalInsertIntoLargeMap(int key);
 
   // Grows the flat_capacity_.
   // If flat_capacity_ > kMaximumFlatCapacity, converts to LargeMap.