avoid allocating memory if all extension are cleared (#9345)

* Update extension_set.cc

avoid allocating memory if all extension are cleared

* Avoid allocating too much memory

* Accurately count flat memory size

* Update extension_set.cc

use the 4x growth.

* add unittest for ExtensionSet

src/google/protobuf/extension_set.cc

@@ -972,6 +972,26 @@ size_t SizeOfUnion(ItX it_xs, ItX end_xs, ItY it_ys, ItY end_ys) {
 
 void ExtensionSet::MergeFrom(const MessageLite* extendee,
                              const ExtensionSet& other) {
+  // check if the other is cleared
+  bool is_other_all_cleared = true;
+  if (PROTOBUF_PREDICT_TRUE(!other.is_large())){
+    for (const KeyValue* it = other.flat_begin(); it != other.flat_end(); it++){
+      is_other_all_cleared = is_other_all_cleared && it->second.is_cleared;
+      if (!is_other_all_cleared){
+        break;
+      }
+    }
+  } else {
+    for (LargeMap::iterator it = other.map_.large->begin(); it != other.map_.large->end(); it++){
+      is_other_all_cleared = is_other_all_cleared && it->second.is_cleared;
+      if (!is_other_all_cleared){
+        break;
+      }
+    }
+  }
+  if (is_other_all_cleared){
+    return;
+  }
   if (PROTOBUF_PREDICT_TRUE(!is_large())) {
     if (PROTOBUF_PREDICT_TRUE(!other.is_large())) {
       GrowCapacity(SizeOfUnion(flat_begin(), flat_end(), other.flat_begin(),

src/google/protobuf/extension_set_heavy.cc

@@ -359,7 +359,7 @@ int ExtensionSet::SpaceUsedExcludingSelf() const {
 }
 
 size_t ExtensionSet::SpaceUsedExcludingSelfLong() const {
-  size_t total_size = Size() * sizeof(KeyValue);
+  size_t total_size = (is_large() ? map_.large->size() : flat_capacity_) * sizeof(KeyValue);
   ForEach([&total_size](int /* number */, const Extension& ext) {
     total_size += ext.SpaceUsedExcludingSelfLong();
   });

src/google/protobuf/extension_set_unittest.cc

@@ -1338,6 +1338,18 @@ TEST(ExtensionSetTest, ConstInit) {
   EXPECT_EQ(set.NumExtensions(), 0);
 }
 
+TEST(ExtensionSetTest, ExtensionSetSpaceUsed) {
+  unittest::TestAllExtensions msg;
+  long l = msg.SpaceUsedLong();
+  msg.SetExtension(unittest::optional_int32_extension, 100);
+  unittest::TestAllExtensions msg2(msg);
+  long l2 = msg2.SpaceUsedLong();
+  msg.ClearExtension(unittest::optional_int32_extension);
+  unittest::TestAllExtensions msg3(msg);
+  long l3 = msg3.SpaceUsedLong();
+  EXPECT_TRUE((l2 - l) > (l3 - l));
+}
+
 }  // namespace
 }  // namespace internal
 }  // namespace protobuf