Create a reflection helper for ExtensionIdentifier.

This will allow easy access to the FieldDescriptor of a generated extension.

PiperOrigin-RevId: 561480962

Author: mkruskal-google

src/google/protobuf/compiler/code_generator_unittest.cc

@@ -95,8 +95,7 @@ class TestGenerator : public CodeGenerator {
   absl::string_view minimum_edition_ = PROTOBUF_MINIMUM_EDITION;
   absl::string_view maximum_edition_ = PROTOBUF_MAXIMUM_EDITION;
   std::vector<const FieldDescriptor*> feature_extensions_ = {
-      DescriptorPool::generated_pool()->FindExtensionByNumber(
-          FeatureSet::descriptor(), pb::test.number())};
+      GetExtensionReflection(pb::test)};
 };
 
 class SimpleErrorCollector : public io::ErrorCollector {
@@ -190,9 +189,7 @@ TEST_F(CodeGeneratorTest, GetUnresolvedSourceFeaturesInherited) {
 
 TEST_F(CodeGeneratorTest, GetResolvedSourceFeaturesRoot) {
   TestGenerator generator;
-  generator.set_feature_extensions(
-      {DescriptorPool::generated_pool()->FindExtensionByNumber(
-          FeatureSet::descriptor(), pb::test.number())});
+  generator.set_feature_extensions({GetExtensionReflection(pb::test)});
   pool_.SetFeatureSetDefaults(*generator.BuildFeatureSetDefaults());
 
   ASSERT_THAT(BuildFile(DescriptorProto::descriptor()->file()), NotNull());
@@ -225,9 +222,7 @@ TEST_F(CodeGeneratorTest, GetResolvedSourceFeaturesRoot) {
 
 TEST_F(CodeGeneratorTest, GetResolvedSourceFeaturesInherited) {
   TestGenerator generator;
-  generator.set_feature_extensions(
-      {DescriptorPool::generated_pool()->FindExtensionByNumber(
-          FeatureSet::descriptor(), pb::test.number())});
+  generator.set_feature_extensions({GetExtensionReflection(pb::test)});
   pool_.SetFeatureSetDefaults(*generator.BuildFeatureSetDefaults());
 
   ASSERT_THAT(BuildFile(DescriptorProto::descriptor()->file()), NotNull());

src/google/protobuf/compiler/command_line_interface_unittest.cc

@@ -1482,8 +1482,7 @@ TEST_F(CommandLineInterfaceTest, FeatureExtensionError) {
     })schema");
 
   mock_generator_->set_feature_extensions(
-      {DescriptorPool::generated_pool()->FindExtensionByNumber(
-          FeatureSet::descriptor(), pb::test_invalid.number())});
+      {GetExtensionReflection(pb::test_invalid)});
 
   Run("protocol_compiler --proto_path=$tmpdir --test_out=$tmpdir "
       "--experimental_editions foo.proto");

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

@@ -96,8 +96,7 @@ class PROTOC_EXPORT CppGenerator : public CodeGenerator {
   }
 
   std::vector<const FieldDescriptor*> GetFeatureExtensions() const override {
-    return {DescriptorPool::generated_pool()->FindExtensionByNumber(
-        FeatureSet::descriptor(), pb::cpp.number())};
+    return {GetExtensionReflection(pb::cpp)};
   }
 
  private:

src/google/protobuf/compiler/mock_code_generator.h

@@ -149,8 +149,7 @@ class MockCodeGenerator : public CodeGenerator {
   absl::string_view minimum_edition_ = PROTOBUF_MINIMUM_EDITION;
   absl::string_view maximum_edition_ = PROTOBUF_MAXIMUM_EDITION;
   std::vector<const FieldDescriptor*> feature_extensions_ = {
-      DescriptorPool::generated_pool()->FindExtensionByNumber(
-          FeatureSet::descriptor(), pb::test.number())};
+      GetExtensionReflection(pb::test)};
 
   static std::string GetOutputFileContent(absl::string_view generator_name,
                                           absl::string_view parameter,

src/google/protobuf/descriptor_unittest.cc

@@ -7246,10 +7246,9 @@ class FeaturesTest : public FeaturesBaseTest {
 
     auto default_spec = FeatureResolver::CompileDefaults(
         FeatureSet::descriptor(),
-        {pb::CppFeatures::descriptor()->file()->extension(0),
-         pb::TestFeatures::descriptor()->file()->extension(0),
-         pb::TestMessage::descriptor()->extension(0),
-         pb::TestMessage::Nested::descriptor()->extension(0)},
+        {GetExtensionReflection(pb::cpp), GetExtensionReflection(pb::test),
+         GetExtensionReflection(pb::TestMessage::test_message),
+         GetExtensionReflection(pb::TestMessage::Nested::test_nested)},
         "2023", "2025");
     ASSERT_OK(default_spec);
     pool_.SetFeatureSetDefaults(std::move(default_spec).value());

src/google/protobuf/extension_set.h

@@ -1554,6 +1554,9 @@ template <>
 class ExtensionIdentifier<FeatureSet, MessageTypeTraits<::pb::CppFeatures>, 11,
                           false> {
  public:
+  using TypeTraits = MessageTypeTraits<::pb::CppFeatures>;
+  using Extendee = FeatureSet;
+
   explicit constexpr ExtensionIdentifier(int number) : number_(number) {}
 
   int number() const { return number_; }
@@ -1612,6 +1615,21 @@ void LinkExtensionReflection(
   internal::StrongReference(extension);
 }
 
+// Returns the field descriptor for a generated extension identifier.  This is
+// useful when doing reflection over generated extensions.
+template <typename ExtendeeType, typename TypeTraitsType,
+          internal::FieldType field_type, bool is_packed,
+          typename PoolType = DescriptorPool>
+const FieldDescriptor* GetExtensionReflection(
+    const google::protobuf::internal::ExtensionIdentifier<
+        ExtendeeType, TypeTraitsType, field_type, is_packed>& extension) {
+  return PoolType::generated_pool()->FindExtensionByNumber(
+      google::protobuf::internal::ExtensionIdentifier<ExtendeeType, TypeTraitsType,
+                                            field_type,
+                                            is_packed>::Extendee::descriptor(),
+      extension.number());
+}
+
 }  // namespace protobuf
 }  // namespace google
 

src/google/protobuf/extension_set_unittest.cc

@@ -41,6 +41,7 @@
 #include "absl/strings/cord.h"
 #include "absl/strings/match.h"
 #include "google/protobuf/arena.h"
+#include "google/protobuf/cpp_features.pb.h"
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/dynamic_message.h"
 #include "google/protobuf/io/coded_stream.h"
@@ -1430,6 +1431,18 @@ TEST(ExtensionSetTest, ExtensionSetSpaceUsed) {
   EXPECT_TRUE((l2 - l) > (l3 - l));
 }
 
+TEST(ExtensionSetTest, Descriptor) {
+  EXPECT_EQ(
+      GetExtensionReflection(unittest::optional_int32_extension),
+      unittest::TestAllExtensions::descriptor()->file()->FindExtensionByName(
+          "optional_int32_extension"));
+  EXPECT_NE(GetExtensionReflection(unittest::optional_int32_extension),
+            nullptr);
+  EXPECT_EQ(GetExtensionReflection(pb::cpp),
+            pb::CppFeatures::descriptor()->file()->FindExtensionByName("cpp"));
+  EXPECT_NE(GetExtensionReflection(pb::cpp), nullptr);
+}
+
 }  // namespace
 }  // namespace internal
 }  // namespace protobuf