Provide an option for DescriptorPoolDatabase to preserve source code info.

PiperOrigin-RevId: 647452898

src/google/protobuf/BUILD.bazel

@@ -1255,6 +1255,7 @@ cc_test(
     srcs = ["descriptor_database_unittest.cc"],
     deps = [
         ":protobuf",
+        ":test_textproto",
         "//src/google/protobuf/testing",
         "//src/google/protobuf/testing:file",
         "@com_google_googletest//:gtest",

src/google/protobuf/descriptor_database.cc

@@ -881,8 +881,9 @@ EncodedDescriptorDatabase::~EncodedDescriptorDatabase() {
 
 // ===================================================================
 
-DescriptorPoolDatabase::DescriptorPoolDatabase(const DescriptorPool& pool)
-    : pool_(pool) {}
+DescriptorPoolDatabase::DescriptorPoolDatabase(
+    const DescriptorPool& pool, DescriptorPoolDatabaseOptions options)
+    : pool_(pool), options_(std::move(options)) {}
 DescriptorPoolDatabase::~DescriptorPoolDatabase() {}
 
 bool DescriptorPoolDatabase::FindFileByName(const std::string& filename,
@@ -891,6 +892,9 @@ bool DescriptorPoolDatabase::FindFileByName(const std::string& filename,
   if (file == nullptr) return false;
   output->Clear();
   file->CopyTo(output);
+  if (options_.preserve_source_code_info) {
+    file->CopySourceCodeInfoTo(output);
+  }
   return true;
 }
 
@@ -900,6 +904,9 @@ bool DescriptorPoolDatabase::FindFileContainingSymbol(
   if (file == nullptr) return false;
   output->Clear();
   file->CopyTo(output);
+  if (options_.preserve_source_code_info) {
+    file->CopySourceCodeInfoTo(output);
+  }
   return true;
 }
 
@@ -915,6 +922,9 @@ bool DescriptorPoolDatabase::FindFileContainingExtension(
 
   output->Clear();
   extension->file()->CopyTo(output);
+  if (options_.preserve_source_code_info) {
+    extension->file()->CopySourceCodeInfoTo(output);
+  }
   return true;
 }
 

src/google/protobuf/descriptor_database.h

@@ -311,10 +311,17 @@ class PROTOBUF_EXPORT EncodedDescriptorDatabase : public DescriptorDatabase {
                   FileDescriptorProto* output);
 };
 
+struct PROTOBUF_EXPORT DescriptorPoolDatabaseOptions {
+  // If true, the database will preserve source code info when returning
+  // descriptors.
+  bool preserve_source_code_info = false;
+};
+
 // A DescriptorDatabase that fetches files from a given pool.
 class PROTOBUF_EXPORT DescriptorPoolDatabase : public DescriptorDatabase {
  public:
-  explicit DescriptorPoolDatabase(const DescriptorPool& pool);
+  explicit DescriptorPoolDatabase(const DescriptorPool& pool,
+                                  DescriptorPoolDatabaseOptions options = {});
   DescriptorPoolDatabase(const DescriptorPoolDatabase&) = delete;
   DescriptorPoolDatabase& operator=(const DescriptorPoolDatabase&) = delete;
   ~DescriptorPoolDatabase() override;
@@ -332,6 +339,7 @@ class PROTOBUF_EXPORT DescriptorPoolDatabase : public DescriptorDatabase {
 
  private:
   const DescriptorPool& pool_;
+  DescriptorPoolDatabaseOptions options_;
 };
 
 // A DescriptorDatabase that wraps two or more others.  It first searches the

src/google/protobuf/descriptor_database_unittest.cc

@@ -21,6 +21,7 @@
 #include "google/protobuf/testing/googletest.h"
 #include <gtest/gtest.h>
 #include "google/protobuf/descriptor.h"
+#include "google/protobuf/test_textproto.h"
 #include "google/protobuf/text_format.h"
 
 
@@ -46,8 +47,6 @@ static void ExpectContainsType(const FileDescriptorProto& proto,
 
 // ===================================================================
 
-#if GTEST_HAS_PARAM_TEST
-
 // SimpleDescriptorDatabase, EncodedDescriptorDatabase, and
 // DescriptorPoolDatabase call for very similar tests.  Instead of writing
 // three nearly-identical sets of tests, we use parameterized tests to apply
@@ -456,8 +455,6 @@ INSTANTIATE_TEST_CASE_P(
 INSTANTIATE_TEST_CASE_P(Pool, DescriptorDatabaseTest,
                         testing::Values(&DescriptorPoolDatabaseTestCase::New));
 
-#endif  // GTEST_HAS_PARAM_TEST
-
 TEST(EncodedDescriptorDatabaseExtraTest, FindNameOfFileContainingSymbol) {
   // Create two files, one of which is in two parts.
   FileDescriptorProto file1, file2a, file2b;
@@ -570,6 +567,77 @@ TEST(SimpleDescriptorDatabaseExtraTest, AddUnowned) {
   EXPECT_THAT(messages, ::testing::UnorderedElementsAre("foo.Foo", "Bar"));
 }
 
+TEST(DescriptorPoolDatabaseTest, PreserveSourceCodeInfo) {
+  SimpleDescriptorDatabase original_db;
+  AddToDatabase(&original_db, R"pb(
+    name: "foo.proto"
+    package: "foo"
+    message_type {
+      name: "Foo"
+      extension_range { start: 1 end: 100 }
+    }
+    extension {
+      name: "foo_ext"
+      extendee: ".foo.Foo"
+      number: 3
+      label: LABEL_OPTIONAL
+      type: TYPE_INT32
+    }
+    source_code_info { location { leading_detached_comments: "comment" } }
+  )pb");
+  DescriptorPool pool(&original_db);
+  DescriptorPoolDatabase db(
+      pool, DescriptorPoolDatabaseOptions{/*preserve_source_code_info=*/true});
+
+  FileDescriptorProto file;
+  ASSERT_TRUE(db.FindFileByName("foo.proto", &file));
+  EXPECT_THAT(
+      file.source_code_info(),
+      EqualsProto(R"pb(location { leading_detached_comments: "comment" })pb"));
+
+  ASSERT_TRUE(db.FindFileContainingExtension("foo.Foo", 3, &file));
+  EXPECT_THAT(
+      file.source_code_info(),
+      EqualsProto(R"pb(location { leading_detached_comments: "comment" })pb"));
+
+  ASSERT_TRUE(db.FindFileContainingSymbol("foo.Foo", &file));
+  EXPECT_THAT(
+      file.source_code_info(),
+      EqualsProto(R"pb(location { leading_detached_comments: "comment" })pb"));
+}
+
+TEST(DescriptorPoolDatabaseTest, StripSourceCodeInfo) {
+  SimpleDescriptorDatabase original_db;
+  AddToDatabase(&original_db, R"pb(
+    name: "foo.proto"
+    package: "foo"
+    message_type {
+      name: "Foo"
+      extension_range { start: 1 end: 100 }
+    }
+    extension {
+      name: "foo_ext"
+      extendee: ".foo.Foo"
+      number: 3
+      label: LABEL_OPTIONAL
+      type: TYPE_INT32
+    }
+    source_code_info { location { leading_detached_comments: "comment" } }
+  )pb");
+  DescriptorPool pool(&original_db);
+  DescriptorPoolDatabase db(pool);
+
+  FileDescriptorProto file;
+  ASSERT_TRUE(db.FindFileByName("foo.proto", &file));
+  EXPECT_FALSE(file.has_source_code_info());
+
+  ASSERT_TRUE(db.FindFileContainingExtension("foo.Foo", 3, &file));
+  EXPECT_FALSE(file.has_source_code_info());
+
+  ASSERT_TRUE(db.FindFileContainingSymbol("foo.Foo", &file));
+  EXPECT_FALSE(file.has_source_code_info());
+}
+
 // ===================================================================
 
 class MergedDescriptorDatabaseTest : public testing::Test {