Apply the two phase extension registration approach to all extensions, not just

message extensions. Otherwise we might still miss some descriptor extensions
when parsing descriptors in phase two.

PiperOrigin-RevId: 618924435

src/google/protobuf/compiler/cpp/extension.cc

@@ -180,24 +180,61 @@ void ExtensionGenerator::GenerateDefinition(io::Printer* p) {
 }
 
 bool ExtensionGenerator::WillGenerateRegistration(InitPriority priority) {
-  // We only use priority 101 for weak descriptors.
+  // When not using weak descriptors we initialize everything on priority 102.
-  return priority != kInitPriority101 ||
+  if (!UsingImplicitWeakDescriptor(descriptor_->file(), options_)) {
-         (descriptor_->cpp_type() == descriptor_->CPPTYPE_MESSAGE &&
+    return priority == kInitPriority102;
-          UsingImplicitWeakDescriptor(descriptor_->file(), options_));
+  }
+  return true;
 }
 
 void ExtensionGenerator::GenerateRegistration(io::Printer* p,
                                               InitPriority priority) {
   ABSL_CHECK(WillGenerateRegistration(priority));
+  const bool using_implicit_weak_descriptors =
+      UsingImplicitWeakDescriptor(descriptor_->file(), options_);
+  const auto find_index = [](auto* desc) {
+    const std::vector<const Descriptor*> msgs =
+        FlattenMessagesInFile(desc->file());
+    return absl::c_find(msgs, desc) - msgs.begin();
+  };
   auto vars = p->WithVars(variables_);
+  auto vars2 = p->WithVars({{
+      {"extendee_table",
+       DescriptorTableName(descriptor_->containing_type()->file(), options_)},
+      {"extendee_index", find_index(descriptor_->containing_type())},
+      {"preregister", priority == kInitPriority101},
+  }});
   switch (descriptor_->cpp_type()) {
     case FieldDescriptor::CPPTYPE_ENUM:
-      p->Emit({{"enum_name", ClassName(descriptor_->enum_type(), true)}},
+      if (using_implicit_weak_descriptors) {
-              R"cc(
+        p->Emit({{"enum_name", ClassName(descriptor_->enum_type(), true)}},
-                ::_pbi::ExtensionSet::RegisterEnumExtension(
+                R"cc(
-                    &$extendee$::default_instance(), $number$, $field_type$,
+#if defined(PROTOBUF_INTERNAL_TEMPORARY_WEAK_EXTENSION_OPT_IN)
-                    $repeated$, $packed$, $enum_name$_IsValid),
+                  (::_pbi::ExtensionSet::ShouldRegisterAtThisTime(
-              )cc");
+                       {{&$extendee_table$, $extendee_index$}}, $preregister$)
+                       ? ::_pbi::ExtensionSet::RegisterEnumExtension(
+                             ::_pbi::GetPrototypeForWeakDescriptor(
+                                 &$extendee_table$, $extendee_index$, true),
+                             $number$, $field_type$, $repeated$, $packed$,
+                             $enum_name$_IsValid)
+                       : (void)0),
+#else
+                )cc");
+      }
+      if (priority == kInitPriority102) {
+        p->Emit({{"enum_name", ClassName(descriptor_->enum_type(), true)}},
+                R"cc(
+                  ::_pbi::ExtensionSet::RegisterEnumExtension(
+                      &$extendee$::default_instance(), $number$, $field_type$,
+                      $repeated$, $packed$, $enum_name$_IsValid),
+                )cc");
+      }
+      if (using_implicit_weak_descriptors) {
+        p->Emit(R"cc(
+#endif
+        )cc");
+      }
+
       break;
     case FieldDescriptor::CPPTYPE_MESSAGE: {
       const bool should_verify =
@@ -215,64 +252,76 @@ void ExtensionGenerator::GenerateRegistration(io::Printer* p,
            {"lazy", descriptor_->options().has_lazy()
                         ? descriptor_->options().lazy() ? "kLazy" : "kEager"
                         : "kUndefined"}});
-      const auto register_message = [&] {
+      if (using_implicit_weak_descriptors) {
-        p->Emit(R"cc(
-          ::_pbi::ExtensionSet::RegisterMessageExtension(
-              &$extendee$::default_instance(), $number$, $field_type$,
-              $repeated$, $packed$, &$message_type$::default_instance(),
-              $verify$, ::_pbi::LazyAnnotation::$lazy$),
-        )cc");
-      };
-      if (UsingImplicitWeakDescriptor(descriptor_->file(), options_)) {
-        const auto find_index = [](auto* desc) {
-          const std::vector<const Descriptor*> msgs =
-              FlattenMessagesInFile(desc->file());
-          return absl::c_find(msgs, desc) - msgs.begin();
-        };
         p->Emit(
             {
-                {"extendee_table",
-                 DescriptorTableName(descriptor_->containing_type()->file(),
-                                     options_)},
-                {"extendee_index", find_index(descriptor_->containing_type())},
                 {"extension_table",
                  DescriptorTableName(descriptor_->message_type()->file(),
                                      options_)},
                 {"extension_index", find_index(descriptor_->message_type())},
-                {"preregister", priority == kInitPriority101},
-                {"fallback_for_opt_out",
-                 // For now we have a fallback to use normal registration,
-                 // which should only happen at priority 102.
-                 // Once we turn this on by default we can remove the opt-in
-                 // and the fallback.
-                 [&] {
-                   if (priority != kInitPriority102) return;
-                   register_message();
-                 }},
             },
             R"cc(
-              //
 #if defined(PROTOBUF_INTERNAL_TEMPORARY_WEAK_EXTENSION_OPT_IN)
-              ::_pbi::ExtensionSet::RegisterWeakMessageExtension(
+              (::_pbi::ExtensionSet::ShouldRegisterAtThisTime(
-                  {&$extendee_table$, $extendee_index$}, $number$, $field_type$,
+                   {{&$extendee_table$, $extendee_index$},
-                  $repeated$, {&$extension_table$, $extension_index$}, $verify$,
+                    {&$extension_table$, $extension_index$}},
-                  ::_pbi::LazyAnnotation::$lazy$, $preregister$),
+                   $preregister$)
+                   ? ::_pbi::ExtensionSet::RegisterMessageExtension(
+                         ::_pbi::GetPrototypeForWeakDescriptor(
+                             &$extendee_table$, $extendee_index$, true),
+                         $number$, $field_type$, $repeated$, $packed$,
+                         ::_pbi::GetPrototypeForWeakDescriptor(
+                             &$extension_table$, $extension_index$, true),
+                         $verify$, ::_pbi::LazyAnnotation::$lazy$)
+                   : (void)0),
 #else
-              $fallback_for_opt_out$,
-#endif
             )cc");
-      } else {
+      }
-        register_message();
+      if (priority == kInitPriority102) {
+        p->Emit(R"cc(
+          ::_pbi::ExtensionSet::RegisterMessageExtension(
+              &$extendee$::default_instance(), $number$, $field_type$,
+              $repeated$, $packed$, &$message_type$::default_instance(),
+              $verify$, ::_pbi::LazyAnnotation::$lazy$),
+        )cc");
+      }
+      if (using_implicit_weak_descriptors) {
+        p->Emit(R"cc(
+#endif
+        )cc");
       }
       break;
     }
+
     default:
-      p->Emit(
+      if (using_implicit_weak_descriptors) {
-          R"cc(
+        p->Emit(R"cc(
-            ::_pbi::ExtensionSet::RegisterExtension(
+#if defined(PROTOBUF_INTERNAL_TEMPORARY_WEAK_EXTENSION_OPT_IN)
-                &$extendee$::default_instance(), $number$, $field_type$,
+          (::_pbi::ExtensionSet::ShouldRegisterAtThisTime(
-                $repeated$, $packed$),
+               {{&$extendee_table$, $extendee_index$}}, $preregister$)
-          )cc");
+               ? ::_pbi::ExtensionSet::RegisterExtension(
+                     ::_pbi::GetPrototypeForWeakDescriptor(&$extendee_table$,
+                                                           $extendee_index$,
+                                                           true),
+                     $number$, $field_type$, $repeated$, $packed$)
+               : (void)0),
+#else
+        )cc");
+      }
+      if (priority == kInitPriority102) {
+        p->Emit(
+            R"cc(
+              ::_pbi::ExtensionSet::RegisterExtension(
+                  &$extendee$::default_instance(), $number$, $field_type$,
+                  $repeated$, $packed$),
+            )cc");
+      }
+      if (using_implicit_weak_descriptors) {
+        p->Emit(R"cc(
+#endif
+        )cc");
+      }
+
       break;
   }
 }

src/google/protobuf/compiler/cpp/file.cc

@@ -1077,11 +1077,16 @@ static void GatherAllCustomOptionTypes(
     reflection.ListFields(msg, &fields);
 
     for (auto* field : fields) {
-      const auto* field_msg = field->message_type();
-      if (field_msg == nullptr) continue;
       if (field->is_extension()) {
+        // Always add the extended.
         const Descriptor* desc = msg.GetDescriptor();
         out[desc->full_name()] = desc;
+      }
+
+      // Add and recurse of the extendee if it is a message.
+      const auto* field_msg = field->message_type();
+      if (field_msg == nullptr) continue;
+      if (field->is_extension()) {
         out[field_msg->full_name()] = field_msg;
       }
       if (field->is_repeated()) {

src/google/protobuf/extension_set.h

@@ -20,6 +20,7 @@
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
+#include <initializer_list>
 #include <string>
 #include <type_traits>
 #include <utility>
@@ -160,14 +161,6 @@ struct ExtensionInfo {
 };
 
 
-// Reference to a prototype via its DescriptorTable.
-// This way we can generate them on the fly if they are missing when Weak
-// Descriptor messages are enabled.
-struct WeakPrototypeRef {
-  const internal::DescriptorTable* table;
-  int index;
-};
-
 // An ExtensionFinder is an object which looks up extension definitions.  It
 // must implement this method:
 //
@@ -236,16 +229,18 @@ class PROTOBUF_EXPORT ExtensionSet {
                                        LazyEagerVerifyFnType verify_func,
                                        LazyAnnotation is_lazy);
 
-  // As RegisterMessageExtension, but for the weak descriptor message mode.
+  // In weak descriptor mode we register extensions in two phases.
-  // It will perform the registration in two phases to guarantee we can parse
+  // This function determines if it is the right time to register a particular
-  // descriptors properly.
+  // extension.
-  static void RegisterWeakMessageExtension(internal::WeakPrototypeRef extendee,
+  // During "preregistration" we only register extensions that have all their
-                                           int number, FieldType type,
+  // types linked in.
-                                           bool is_repeated,
+  struct WeakPrototypeRef {
-                                           internal::WeakPrototypeRef prototype,
+    const internal::DescriptorTable* table;
-                                           LazyEagerVerifyFnType verify_func,
+    int index;
-                                           LazyAnnotation is_lazy,
+  };
-                                           bool is_preregistration);
+  static bool ShouldRegisterAtThisTime(
+      std::initializer_list<WeakPrototypeRef> messages,
+      bool is_preregistration);
 
   // =================================================================
 

src/google/protobuf/extension_set_heavy.cc

@@ -14,6 +14,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <initializer_list>
 #include <vector>
 
 #include "absl/base/attributes.h"
@@ -426,39 +427,14 @@ uint8_t* ExtensionSet::SerializeMessageSetWithCachedSizesToArray(
 }
 
 #if defined(PROTOBUF_DESCRIPTOR_WEAK_MESSAGES_ALLOWED)
-// First, register all the extensions that have both messages linked in.
+bool ExtensionSet::ShouldRegisterAtThisTime(
-// This will include all messages used as extensions in .proto options.
+    std::initializer_list<WeakPrototypeRef> messages, bool is_preregistration) {
-// In the second phase, we generate the missing prototypes, but that requires
+  bool has_all = true;
-// parsing descriptors, which in turn require the extensions from the first
+  for (auto ref : messages) {
-// phase.
+    has_all = has_all && GetPrototypeForWeakDescriptor(ref.table, ref.index,
-void ExtensionSet::RegisterWeakMessageExtension(
+                                                       false) != nullptr;
-    internal::WeakPrototypeRef extendee, int number, FieldType type,
-    bool is_repeated, internal::WeakPrototypeRef prototype,
-    LazyEagerVerifyFnType verify_func, LazyAnnotation is_lazy,
-    bool is_preregistration) {
-  auto* extendee_msg =
-      GetPrototypeForWeakDescriptor(extendee.table, extendee.index, false);
-  auto* prototype_msg =
-      GetPrototypeForWeakDescriptor(prototype.table, prototype.index, false);
-
-  const bool have_both = extendee_msg != nullptr && prototype_msg != nullptr;
-  if (is_preregistration != have_both) {
-    // This is done on the other phase.
-    return;
   }
-
+  return has_all == is_preregistration;
-  if (extendee_msg == nullptr) {
-    extendee_msg =
-        GetPrototypeForWeakDescriptor(extendee.table, extendee.index, true);
-  }
-  if (prototype_msg == nullptr) {
-    prototype_msg =
-        GetPrototypeForWeakDescriptor(prototype.table, prototype.index, true);
-  }
-
-  ExtensionSet::RegisterMessageExtension(
-      extendee_msg, number, type, is_repeated,
-      /*is_packed=*/false, prototype_msg, verify_func, is_lazy);
 }
 #endif  // PROTOBUF_DESCRIPTOR_WEAK_MESSAGES_ALLOWED
 

src/google/protobuf/io/printer.cc

@@ -24,6 +24,7 @@
 #include "absl/log/absl_log.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/escaping.h"
+#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_split.h"
@@ -127,6 +128,12 @@ Printer::Format Printer::TokenizeFormat(absl::string_view format_string,
       format.is_raw_string = false;
       raw_string_indent = 0;
     }
+
+    // This means we have a preprocessor directive and we should not have eaten
+    // the newline.
+    if (!at_start_of_line_ && absl::StartsWith(format_string, "#")) {
+      format_string = orig;
+    }
   }
 
   // We now split the remaining format string into lines and discard:

src/google/protobuf/io/printer_unittest.cc

@@ -713,6 +713,44 @@ TEST_F(PrinterTest, EmitWithIndent) {
             "  };\n");
 }
 
+TEST_F(PrinterTest, EmitWithPreprocessor) {
+  {
+    Printer printer(output());
+    auto v = printer.WithIndent();
+    printer.Emit({{"value",
+                   [&] {
+                     printer.Emit(R"cc(
+#if FOO
+                       0,
+#else
+                       1,
+#endif
+                     )cc");
+                   }},
+                  {"on_new_line",
+                   [&] {
+                     printer.Emit(R"cc(
+#pragma foo
+                     )cc");
+                   }}},
+                 R"cc(
+                   int val = ($value$, 0);
+                   $on_new_line$;
+                 )cc");
+  }
+
+  EXPECT_EQ(written(),
+            R"(  int val = (
+  #if FOO
+                         0,
+  #else
+                         1,
+  #endif
+   0);
+  #pragma foo
+)");
+}
+
 
 TEST_F(PrinterTest, EmitSameNameAnnotation) {
   FakeAnnotationCollector collector;