Extend the weak descriptor message feature to include extensions.

An extension declaration now allows both messages to be weak and will create the prototype on the fly. PiperOrigin-RevId: 588811002
protocolbuffers · Dec 7, 2023 · 8edf0bd · 8edf0bd
1 parent 9760292
commit 8edf0bd
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 31 deletions.
diff --git a/src/google/protobuf/compiler/cpp/extension.cc b/src/google/protobuf/compiler/cpp/extension.cc
@@ -12,6 +12,7 @@
 #include "google/protobuf/compiler/cpp/extension.h"
 
 #include <string>
+#include <vector>
 
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_replace.h"
@@ -188,35 +189,58 @@ void ExtensionGenerator::GenerateRegistration(io::Printer* p) {
                     $repeated$, $packed$, $enum_name$_IsValid),
               )cc");
       break;
-    case FieldDescriptor::CPPTYPE_MESSAGE:
-      p->Emit({{"verify",
-                [&] {
-                  const bool should_verify =
-                      // Only verify msgs.
-                      descriptor_->cpp_type() ==
-                          FieldDescriptor::CPPTYPE_MESSAGE &&
-                      // Options say to verify.
-                      ShouldVerify(descriptor_->message_type(), options_,
-                                   scc_analyzer_) &&
-                      ShouldVerify(descriptor_->containing_type(), options_,
-                                   scc_analyzer_);
-                  if (should_verify) {
-                    p->Emit("&$message_type$::InternalVerify,");
-                  } else {
-                    p->Emit("nullptr,");
-                  }
-                }},
-               {"message_type", FieldMessageTypeName(descriptor_, options_)},
-               {"lazy", descriptor_->options().has_lazy()
-                            ? descriptor_->options().lazy() ? "kLazy" : "kEager"
-                            : "kUndefined"}},
-              R"cc(
-                ::_pbi::ExtensionSet::RegisterMessageExtension(
-                    &$extendee$::default_instance(), $number$, $field_type$,
-                    $repeated$, $packed$, &$message_type$::default_instance(),
-                    $verify$, ::_pbi::LazyAnnotation::$lazy$),
-              )cc");
+    case FieldDescriptor::CPPTYPE_MESSAGE: {
+      const bool should_verify =
+          // Only verify msgs.
+          descriptor_->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE &&
+          // Options say to verify.
+          ShouldVerify(descriptor_->message_type(), options_, scc_analyzer_) &&
+          ShouldVerify(descriptor_->containing_type(), options_, scc_analyzer_);
+      const auto message_type = FieldMessageTypeName(descriptor_, options_);
+      auto v = p->WithVars(
+          {{"verify", should_verify
+                          ? absl::StrCat("&", message_type, "::InternalVerify")
+                          : "nullptr"},
+           {"message_type", message_type},
+           {"lazy", descriptor_->options().has_lazy()
+                        ? descriptor_->options().lazy() ? "kLazy" : "kEager"
+                        : "kUndefined"}});
+      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())},
+            },
+            R"cc(
+              ::_pbi::ExtensionSet::RegisterMessageExtension(
+                  ::_pbi::GetPrototypeForWeakDescriptor(&$extendee_table$,
+                                                        $extendee_index$),
+                  $number$, $field_type$, $repeated$, $packed$,
+                  ::_pbi::GetPrototypeForWeakDescriptor(&$extension_table$,
+                                                        $extension_index$),
+                  $verify$, ::_pbi::LazyAnnotation::$lazy$),
+            )cc");
+      } else {
+        p->Emit(R"cc(
+          ::_pbi::ExtensionSet::RegisterMessageExtension(
+              &$extendee$::default_instance(), $number$, $field_type$,
+              $repeated$, $packed$, &$message_type$::default_instance(),
+              $verify$, ::_pbi::LazyAnnotation::$lazy$),
+        )cc");
+      }
       break;
+    }
     default:
       p->Emit(
           R"cc(

diff --git a/src/google/protobuf/compiler/cpp/helpers.h b/src/google/protobuf/compiler/cpp/helpers.h
@@ -772,9 +772,13 @@ void ListAllTypesForServices(const FileDescriptor* fd,
 // their codegen.
 // LITE messages do not participate at all in this feature.
 //
-// For extensions, the identifiers currently pin both the extended and extendee
-// messages. This is the status quo, but not the desired end state which should
-// change in a future update to the feature.
+// For extensions, the identifiers currently pin the extendee. The extended is
+// assumed to by pinned elsewhere since we already have an instance of it when
+// we call `.GetExtension` et al. The extension identifier itself is not
+// automatically pinned, so it has to be used to participate in the graph.
+// Registration of the extensions do not pin the extended or the extendee. At
+// registration time we will eagerly create a prototype object if one is
+// missing to insert in the extension table in ExtensionSet.
 //
 // For services, the TU unconditionally pins the request/response objects.
 // This is the status quo for simplicitly to avoid modifying the RPC layer. It

diff --git a/src/google/protobuf/generated_message_reflection.cc b/src/google/protobuf/generated_message_reflection.cc
@@ -39,6 +39,7 @@
 #include "google/protobuf/inlined_string_field.h"
 #include "google/protobuf/map_field.h"
 #include "google/protobuf/map_field_inl.h"
+#include "google/protobuf/message.h"
 #include "google/protobuf/raw_ptr.h"
 #include "google/protobuf/repeated_field.h"
 #include "google/protobuf/unknown_field_set.h"
@@ -3819,6 +3820,24 @@ bool SplitFieldHasExtraIndirection(const FieldDescriptor* field) {
   return field->is_repeated();
 }
 
+const Message* GetPrototypeForWeakDescriptor(const DescriptorTable* table,
+                                             int index) {
+  // First, make sure we inject the surviving default instances.
+  InitProtobufDefaults();
+
+  // Now check if the table has it. If so, return it.
+  if (const auto* msg = table->default_instances[index]) {
+    return msg;
+  }
+
+  // Fallback to dynamic messages.
+  // Register the dep and generate the prototype via the generated pool.
+  AssignDescriptors(table);
+  ABSL_CHECK(table->file_level_metadata[index].descriptor != nullptr);
+  return MessageFactory::generated_factory()->GetPrototype(
+      table->file_level_metadata[index].descriptor);
+}
+
 }  // namespace internal
 }  // namespace protobuf
 }  // namespace google

diff --git a/src/google/protobuf/generated_message_reflection.h b/src/google/protobuf/generated_message_reflection.h
@@ -344,6 +344,11 @@ struct PROTOBUF_EXPORT AddDescriptorsRunner {
   explicit AddDescriptorsRunner(const DescriptorTable* table);
 };
 
+// Retrieves the existing prototype out of a descriptor table.
+// If it doesn't exist, asks the generated message factory for one.
+const Message* GetPrototypeForWeakDescriptor(const DescriptorTable* table,
+                                             int index);
+
 struct DenseEnumCacheInfo {
   std::atomic<const std::string**> cache;
   int min_val;