Fix a bug with custom features outside of the pb package.

This was always intended to be possible, but was accidentally broken by a recent improvement in C++.  This expands the generated defaults to all features, adding an implicit requirement that the message type of each feature extension is unique.  The only violation of this today was in our test features.

PiperOrigin-RevId: 836350099

Author: mkruskal-google

src/google/protobuf/BUILD.bazel

@@ -1004,6 +1004,7 @@ filegroup(
         "map_proto2_unittest.proto",
         "map_unittest.proto",
         "unittest.proto",
+        "unittest_custom_features.proto",
         "unittest_custom_options.proto",
         "unittest_embed_optimize_for.proto",
         "unittest_empty.proto",
@@ -1089,6 +1090,7 @@ protobuf_test_proto_library(
         "map_unittest.proto",
         "unittest.proto",
         "unittest_arena.proto",
+        "unittest_custom_features.proto",
         "unittest_custom_options.proto",
         "unittest_delimited.proto",
         "unittest_delimited_import.proto",

src/google/protobuf/compiler/code_generator_unittest.cc

@@ -17,12 +17,12 @@
 #include "absl/log/absl_log.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_format.h"
-#include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
 #include "google/protobuf/compiler/parser.h"
 #include "google/protobuf/io/tokenizer.h"
 #include "google/protobuf/io/zero_copy_stream_impl_lite.h"
 #include "google/protobuf/test_textproto.h"
+#include "google/protobuf/unittest_custom_features.pb.h"
 #include "google/protobuf/unittest_features.pb.h"
 
 // Must be included last.
@@ -281,6 +281,45 @@ TEST_F(CodeGeneratorTest, GetResolvedSourceFeatureExtension) {
   EXPECT_EQ(ext2.source_feature(), ext1.source_feature());
 }
 
+TEST_F(CodeGeneratorTest, GetResolvedSourceFeatureExtensionCustom) {
+  TestGenerator generator;
+  generator.set_feature_extensions(
+      {GetExtensionReflection(custom_features::test)});
+  ASSERT_OK(pool_.SetFeatureSetDefaults(*generator.BuildFeatureSetDefaults()));
+
+  ASSERT_THAT(BuildFile(DescriptorProto::descriptor()->file()), NotNull());
+  ASSERT_THAT(
+      BuildFile(custom_features::TestCustomFeatures::descriptor()->file()),
+      NotNull());
+  auto file = BuildFile(R"schema(
+    edition = "2023";
+    package proto2_unittest;
+
+    import "google/protobuf/unittest_custom_features.proto";
+
+    option features.(custom_features.test).file_feature = VALUE6;
+    option features.(custom_features.test).source_feature = VALUE5;
+  )schema");
+  ASSERT_THAT(file, NotNull());
+  const custom_features::TestCustomFeatures& ext1 =
+      TestGenerator::GetResolvedSourceFeatureExtension(*file,
+                                                       custom_features::test);
+  const custom_features::TestCustomFeatures& ext2 =
+      TestGenerator::GetResolvedSourceFeatures(*file).GetExtension(
+          custom_features::test);
+
+  // Since the pool provides the feature set defaults, there should be no
+  // difference between the two results.
+  EXPECT_EQ(ext1.enum_feature(), custom_features::EnumFeature::VALUE1);
+  EXPECT_EQ(ext1.field_feature(), custom_features::EnumFeature::VALUE1);
+  EXPECT_EQ(ext1.file_feature(), custom_features::EnumFeature::VALUE6);
+  EXPECT_EQ(ext1.source_feature(), custom_features::EnumFeature::VALUE5);
+  EXPECT_EQ(ext2.enum_feature(), ext1.enum_feature());
+  EXPECT_EQ(ext2.field_feature(), ext1.field_feature());
+  EXPECT_EQ(ext2.file_feature(), ext1.file_feature());
+  EXPECT_EQ(ext2.source_feature(), ext1.source_feature());
+}
+
 TEST_F(CodeGeneratorTest, GetResolvedSourceFeatureExtensionEditedDefaults) {
   FeatureSetDefaults defaults = ParseTextOrDie(R"pb(
     minimum_edition: EDITION_PROTO2

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

@@ -101,67 +101,59 @@ bool ExtensionGenerator::IsScoped() const {
   return descriptor_->extension_scope() != nullptr;
 }
 
-namespace {
-bool ShouldGenerateFeatureSetDefaultData(absl::string_view extension) {
-  return extension == "pb.java" || extension == "pb.java_mutable" ||
-         extension == "pb.test" || extension == "pb.proto1";
+void ExtensionGenerator::GenerateFeatureDefaults(io::Printer* p) const {
+  auto var = p->WithVars(variables_);
+  auto annotate = p->WithAnnotations({{"name", descriptor_}});
+  if (descriptor_->message_type() == nullptr) return;
+  absl::string_view extendee = descriptor_->containing_type()->full_name();
+  if (extendee != "google.protobuf.FeatureSet") return;
+
+  absl::StatusOr<FeatureSetDefaults> defaults =
+      FeatureResolver::CompileDefaults(descriptor_->containing_type(),
+                                       {descriptor_}, ProtocMinimumEdition(),
+                                       MaximumKnownEdition());
+  if (!defaults.ok()) {
+    return;
+  }
+  p->Emit({{"defaults", absl::Base64Escape(defaults->SerializeAsString())},
+           {"extension_type", ClassName(descriptor_->message_type(), true)},
+           {"function_name", "GetFeatureSetDefaultsData"}},
+          R"cc(
+            namespace internal {
+            template <>
+            inline ::absl::string_view $function_name$<$extension_type$>() {
+              static constexpr char kDefaults[] = "$defaults$";
+              return kDefaults;
+            }
+            }  // namespace internal
+          )cc");
 }
-}  // namespace
 
 void ExtensionGenerator::GenerateDeclaration(io::Printer* p) const {
   auto var = p->WithVars(variables_);
   auto annotate = p->WithAnnotations({{"name", descriptor_}});
   p->Emit(
-      {{"constant_qualifier",
-        // If this is a class member, it needs to be declared
-        //   `static constexpr`.
-        // Otherwise, it will be
-        //   `inline constexpr`.
-        IsScoped() ? "static" : ""},
-       {"id_qualifier",
-        // If this is a class member, it needs to be declared "static".
-        // Otherwise, it needs to be "extern".  In the latter case, it
-        // also needs the DLL export/import specifier.
-        IsScoped() ? "static"
-        : options_.dllexport_decl.empty()
-            ? "extern"
-            : absl::StrCat(options_.dllexport_decl, " extern")},
-       {"feature_set_defaults",
-        [&] {
-          if (!ShouldGenerateFeatureSetDefaultData(descriptor_->full_name())) {
-            return;
-          }
-          if (descriptor_->message_type() == nullptr) return;
-          absl::string_view extendee =
-              descriptor_->containing_type()->full_name();
-          if (extendee != "google.protobuf.FeatureSet") return;
-
-          std::vector<const FieldDescriptor*> extensions = {descriptor_};
-          absl::StatusOr<FeatureSetDefaults> defaults =
-              FeatureResolver::CompileDefaults(
-                  descriptor_->containing_type(), extensions,
-                  ProtocMinimumEdition(), MaximumKnownEdition());
-          ABSL_CHECK_OK(defaults);
-          p->Emit(
-              {{"defaults", absl::Base64Escape(defaults->SerializeAsString())},
-               {"extension_type", ClassName(descriptor_->message_type(), true)},
-               {"function_name", "GetFeatureSetDefaultsData"}},
-              R"cc(
-                namespace internal {
-                template <>
-                inline ::absl::string_view $function_name$<$extension_type$>() {
-                  static constexpr char kDefaults[] = "$defaults$";
-                  return kDefaults;
-                }
-                }  // namespace internal
-              )cc");
-        }}},
+      {
+          {"constant_qualifier",
+           // If this is a class member, it needs to be declared
+           //   `static constexpr`.
+           // Otherwise, it will be
+           //   `inline constexpr`.
+           IsScoped() ? "static" : ""},
+          {"id_qualifier",
+           // If this is a class member, it needs to be declared "static".
+           // Otherwise, it needs to be "extern".  In the latter case, it
+           // also needs the DLL export/import specifier.
+           IsScoped() ? "static"
+           : options_.dllexport_decl.empty()
+               ? "extern"
+               : absl::StrCat(options_.dllexport_decl, " extern")},
+      },
       R"cc(
         inline $constant_qualifier $constexpr int $constant_name$ = $number$;
         $id_qualifier$ $pbi$::ExtensionIdentifier<
             $extendee$, $pbi$::$type_traits$, $field_type$, $packed$>
             $name$;
-        $feature_set_defaults$;
       )cc");
 }
 

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

@@ -58,6 +58,9 @@ class PROTOC_EXPORT ExtensionGenerator {
   // Source file stuff.
   void GenerateDefinition(io::Printer* p);
 
+  // Feature-specific defaults.
+  void GenerateFeatureDefaults(io::Printer* p) const;
+
   // Extension registration can happen at different priority levels depending on
   // the features used.
   //

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

@@ -1588,6 +1588,10 @@ void FileGenerator::GenerateForwardDeclarations(io::Printer* p) {
     decl.second.PrintTopLevelDecl(p, options_);
   }
 
+  for (auto& extension_generator : extension_generators_) {
+    extension_generator->GenerateFeatureDefaults(p);
+  }
+
   if (IsFileDescriptorProto(file_, options_)) {
     ns.ChangeTo(absl::StrCat(ProtobufNamespace(options_), "::internal"));
     p->Emit(R"cc(

src/google/protobuf/compiler/java/java_features.pb.h

@@ -80,13 +80,6 @@ void InitializeLazyExtensionSet();
 }  // namespace google
 namespace pb {
 class CppFeatures;
-namespace internal {
-// Forward-declares the function for FeatureSet extensions to make it visible
-// to the internal feature helper. It should hold and return serialized
-// FeatureSetDefaults data.
-template <class T>
-inline ::absl::string_view GetFeatureSetDefaultsData();
-}  // namespace internal
 }  // namespace pb
 
 namespace google {
@@ -99,6 +92,12 @@ class FindExtensionTest;
 // Forward-declared from message.h.
 PROTOBUF_EXPORT bool IsDescendant(const Message& root, const Message& message);
 
+// Forward-declares the function for FeatureSet extensions to make it visible
+// to the internal feature helper. It should hold and return serialized
+// FeatureSetDefaults data.
+template <class T>
+inline ::absl::string_view GetFeatureSetDefaultsData();
+
 namespace v2 {
 class TableDrivenMessage;
 }  // namespace v2

src/google/protobuf/cpp_features.pb.h

@@ -154,18 +154,9 @@ TEST(FeatureResolverTest, DefaultsTestMessageExtension) {
   EXPECT_EQ(merged->message_encoding(), FeatureSet::LENGTH_PREFIXED);
   EXPECT_FALSE(merged->HasExtension(pb::test));
 
-  const pb::TestFeatures& ext =
+  const pb::TestMessageFeatures& ext =
       merged->GetExtension(pb::TestMessage::test_message);
   EXPECT_EQ(ext.file_feature(), pb::VALUE3);
-  EXPECT_EQ(ext.extension_range_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.message_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.field_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.oneof_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.enum_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.enum_entry_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.service_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.method_feature(), pb::VALUE1);
-  EXPECT_FALSE(ext.bool_field_feature());
 }
 
 TEST(FeatureResolverTest, DefaultsTestNestedExtension) {
@@ -179,18 +170,9 @@ TEST(FeatureResolverTest, DefaultsTestNestedExtension) {
   EXPECT_EQ(merged->message_encoding(), FeatureSet::LENGTH_PREFIXED);
   EXPECT_FALSE(merged->HasExtension(pb::test));
 
-  const pb::TestFeatures& ext =
+  const pb::TestNestedFeatures& ext =
       merged->GetExtension(pb::TestMessage::Nested::test_nested);
   EXPECT_EQ(ext.file_feature(), pb::VALUE3);
-  EXPECT_EQ(ext.extension_range_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.message_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.field_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.oneof_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.enum_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.enum_entry_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.service_feature(), pb::VALUE1);
-  EXPECT_EQ(ext.method_feature(), pb::VALUE1);
-  EXPECT_FALSE(ext.bool_field_feature());
 }
 
 TEST(FeatureResolverTest, DefaultsGeneratedPoolCustom) {

src/google/protobuf/extension_set.h

@@ -114,7 +114,7 @@ class PROTOBUF_EXPORT InternalFeatureHelper {
 
     auto lang_features_ret =
         ParseAndGetEditionResolvedFeatureSet(
-            ::pb::internal::GetFeatureSetDefaultsData<ExtType>(),
+            ::google::protobuf::internal::GetFeatureSetDefaultsData<ExtType>(),
             GetEdition(descriptor))
             .GetExtension(extension);
     lang_features_ret.MergeFrom(lang_features);