Improve MemoizeProjection by allowing any kind of descriptor.

Also, make it a static function to prevent bugs where the pool and the field
don't match. We can get the pool from the field anyway.

PiperOrigin-RevId: 741531141

Author: protobuf-github-bot

src/google/protobuf/descriptor.h

@@ -2527,28 +2527,30 @@ class PROTOBUF_EXPORT DescriptorPool {
     T value;
   };
 
-  // Memoize a projection of a field.  This is used to cache the results of
-  // calling a function on a field, used for expensive descriptor calculations.
-  template <typename Func>
-  const auto& MemoizeProjection(const FieldDescriptor* field, Func func) const {
-    using ResultT = std::decay_t<decltype(func(field))>;
-    ABSL_DCHECK(field->file()->pool() == this);
+  // Memoize a projection of a descriptor. This is used to cache the results of
+  // calling a function on a descriptor, used for expensive descriptor
+  // calculations.
+  template <typename Desc, typename Func>
+  static const auto& MemoizeProjection(const Desc* descriptor, Func func) {
+    using ResultT = std::decay_t<decltype(func(descriptor))>;
+    auto* pool = descriptor->file()->pool();
     static_assert(std::is_empty_v<Func>);
     // This static bool is unique per-Func, so its address can be used as a key.
     static bool type_key;
-    auto key = std::pair<const void*, const void*>(field, &type_key);
+    auto key = std::pair<const void*, const void*>(descriptor, &type_key);
     {
-      absl::ReaderMutexLock lock(&field_memo_table_mutex_);
-      auto it = field_memo_table_->find(key);
-      if (it != field_memo_table_->end()) {
+      absl::ReaderMutexLock lock(&pool->field_memo_table_mutex_);
+      auto it = pool->field_memo_table_->find(key);
+      if (it != pool->field_memo_table_->end()) {
         return internal::DownCast<const MemoData<ResultT>&>(*it->second).value;
       }
     }
     auto result = std::make_unique<MemoData<ResultT>>();
-    result->value = func(field);
+    result->value = func(descriptor);
     {
-      absl::MutexLock lock(&field_memo_table_mutex_);
-      auto insert_result = field_memo_table_->insert({key, std::move(result)});
+      absl::MutexLock lock(&pool->field_memo_table_mutex_);
+      auto insert_result =
+          pool->field_memo_table_->insert({key, std::move(result)});
       auto it = insert_result.first;
       return internal::DownCast<const MemoData<ResultT>&>(*it->second).value;
     }

src/google/protobuf/descriptor_unittest.cc

@@ -12230,10 +12230,9 @@ TEST_F(DescriptorPoolFeaturesTest, ResolvesFeaturesFor) {
 
 class DescriptorPoolMemoizationTest : public ::testing::Test {
  protected:
-  template <typename Func>
-  const auto& MemoizeProjection(const DescriptorPool* pool,
-                                const FieldDescriptor* field, Func func) {
-    return pool->MemoizeProjection(field, func);
+  template <typename Desc, typename Func>
+  const auto& MemoizeProjection(const Desc* descriptor, Func func) {
+    return DescriptorPool::MemoizeProjection(descriptor, func);
   };
 };
 
@@ -12243,13 +12242,12 @@ TEST_F(DescriptorPoolMemoizationTest, MemoizeProjectionBasic) {
     counter++;
     return field->full_name();
   };
-  proto2_unittest::TestAllTypes message;
-  const Descriptor* descriptor = message.GetDescriptor();
+  const Descriptor* descriptor = proto2_unittest::TestAllTypes::descriptor();
 
   const auto& name = DescriptorPoolMemoizationTest::MemoizeProjection(
-      descriptor->file()->pool(), descriptor->field(0), name_lambda);
+      descriptor->field(0), name_lambda);
   const auto& dupe_name = DescriptorPoolMemoizationTest::MemoizeProjection(
-      descriptor->file()->pool(), descriptor->field(0), name_lambda);
+      descriptor->field(0), name_lambda);
 
   ASSERT_EQ(counter, 1);
   ASSERT_EQ(name, "proto2_unittest.TestAllTypes.optional_int32");
@@ -12261,12 +12259,36 @@ TEST_F(DescriptorPoolMemoizationTest, MemoizeProjectionBasic) {
   EXPECT_EQ(&name, &dupe_name);
 
   auto other_name = DescriptorPoolMemoizationTest::MemoizeProjection(
-      descriptor->file()->pool(), descriptor->field(1), name_lambda);
+      descriptor->field(1), name_lambda);
 
   ASSERT_EQ(counter, 2);
   ASSERT_NE(other_name, "proto2_unittest.TestAllTypes.optional_int32");
 }
 
+TEST_F(DescriptorPoolMemoizationTest, SupportsDifferentDescriptorTypes) {
+  static int counter;
+  counter = 0;
+  auto name_lambda = [](const auto* field) {
+    counter++;
+    return field->full_name();
+  };
+
+  const Descriptor* descriptor = proto2_unittest::TestAllTypes::descriptor();
+
+  // Different descriptor types should be accepted and return the appropriate
+  // result, even when reusing the same lambda type.
+  EXPECT_EQ("proto2_unittest.TestAllTypes.optional_int32",
+            DescriptorPoolMemoizationTest::MemoizeProjection(
+                descriptor->field(0), name_lambda));
+  EXPECT_EQ("proto2_unittest.TestAllTypes",
+            DescriptorPoolMemoizationTest::MemoizeProjection(descriptor,
+                                                             name_lambda));
+  EXPECT_EQ("proto2_unittest.TestAllTypes.NestedMessage",
+            DescriptorPoolMemoizationTest::MemoizeProjection(
+                descriptor->nested_type(0), name_lambda));
+  EXPECT_EQ(counter, 3);
+}
+
 TEST_F(DescriptorPoolMemoizationTest, MemoizeProjectionMultithreaded) {
   auto name_lambda = [](const FieldDescriptor* field) {
     return field->full_name();
@@ -12277,9 +12299,9 @@ TEST_F(DescriptorPoolMemoizationTest, MemoizeProjectionMultithreaded) {
   for (int i = 0; i < descriptor->field_count(); ++i) {
     threads.emplace_back([this, name_lambda, descriptor, i]() {
       auto name = DescriptorPoolMemoizationTest::MemoizeProjection(
-          descriptor->file()->pool(), descriptor->field(i), name_lambda);
+          descriptor->field(i), name_lambda);
       auto first_name = DescriptorPoolMemoizationTest::MemoizeProjection(
-          descriptor->file()->pool(), descriptor->field(0), name_lambda);
+          descriptor->field(0), name_lambda);
       ASSERT_THAT(name, HasSubstr("proto2_unittest.TestAllTypes"));
       if (i != 0) {
         ASSERT_NE(name, "proto2_unittest.TestAllTypes.optional_int32");
@@ -12303,7 +12325,7 @@ TEST_F(DescriptorPoolMemoizationTest, MemoizeProjectionInsertionRace) {
     for (int j = 0; j < 3; ++j) {
       threads.emplace_back([this, name_lambda, descriptor, i]() {
         auto name = DescriptorPoolMemoizationTest::MemoizeProjection(
-            descriptor->file()->pool(), descriptor->field(i), name_lambda);
+            descriptor->field(i), name_lambda);
         ASSERT_THAT(name, HasSubstr("proto2_unittest.TestAllTypes"));
       });
     }

src/google/protobuf/text_format.cc

@@ -3037,7 +3037,7 @@ bool TextFormat::Printer::TryRedactFieldValue(
     const Message& message, const FieldDescriptor* field,
     BaseTextGenerator* generator, bool insert_value_separator) const {
   TextFormat::RedactionState redaction_state =
-      field->file()->pool()->MemoizeProjection(
+      DescriptorPool::MemoizeProjection(
           field, [](const FieldDescriptor* field) {
             return TextFormat::GetRedactionState(field);
           });