file_manager: ensure includes are added at the beginning of the next line

Summary: To avoid breaking a line between an existing include and trailing comments (See file_manager_test.cc)

Reviewed By: iahs

Differential Revision: D90194275

fbshipit-source-id: cb7d685c758eb992c54b38ce03eadc88a59a9b06

Author: aristidisp

third-party/thrift/src/thrift/compiler/codemod/annotate_allow_legacy_missing_uris_test.py

@@ -101,6 +101,26 @@ def test_adds_annotation_existing_package_and_other_includes(self):
                 """,
         )
 
+    def test_adds_annotation_existing_package_and_other_includes_with_commend(self):
+        self._check(
+            before="""\
+                include "thrift/annotation/cpp.thrift" // foo
+
+                package;
+
+                struct S { }
+                """,
+            after="""\
+                include "thrift/annotation/cpp.thrift" // foo
+                include "thrift/annotation/thrift.thrift"
+
+                @thrift.AllowLegacyMissingUris
+                package;
+
+                struct S { }
+                """,
+        )
+
     def test_adds_annotation_existing_package_before_include(self):
         self._check(
             before="""\
@@ -160,6 +180,29 @@ def test_adds_annotation_with_namespaces_and_other_include(self):
                 package;
 
 
+                struct S { }
+                """,
+        )
+
+    def test_adds_annotation_with_namespaces_and_other_include_with_comment(self):
+        self._check(
+            before="""\
+                namespace cpp2 foo
+
+                include "thrift/annotation/cpp.thrift" // foo
+
+                struct S { }
+                """,
+            after="""\
+                namespace cpp2 foo
+
+                include "thrift/annotation/cpp.thrift" // foo
+                include "thrift/annotation/thrift.thrift"
+
+                @thrift.AllowLegacyMissingUris
+                package;
+
+
                 struct S { }
                 """,
         )

third-party/thrift/src/thrift/compiler/codemod/file_manager.cc

@@ -157,7 +157,15 @@ std::optional<size_t> file_manager::add_include(std::string thrift_include) {
   const size_t new_include_offset = [&]() -> size_t {
     // If there are includes, add new include after the last one.
     if (!program_->includes().empty()) {
-      return to_offset(program_->includes().back()->src_range().end) + 1;
+      size_t offset = to_offset(program_->includes().back()->src_range().end);
+      // Move to the beginning of the next line to handle any trailing comments
+      // on the include line (e.g., include "foo.thrift" // comment)
+      size_t newline_pos = old_content_.find('\n', offset);
+      if (newline_pos != std::string::npos) {
+        return newline_pos + 1;
+      }
+      // If no newline found, place at end of content
+      return old_content_.length();
     }
 
     // Otherwise, add the include before the first of: package, namespace or

third-party/thrift/src/thrift/compiler/codemod/test/file_manager_test.cc

@@ -15,6 +15,7 @@
  */
 
 #include <string>
+#include <glog/logging.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <thrift/compiler/ast/t_program_bundle.h>
@@ -56,9 +57,10 @@ Fixture prepare_fixture(std::string_view contents) {
   diagnostics_engine diags = make_diagnostics_printer(*sourceManager);
   std::unique_ptr<t_program_bundle> programBundle =
       parse_ast(*sourceManager, diags, kTestFileName, {});
+  CHECK_NOTNULL(programBundle.get());
 
   auto fileManager = std::make_unique<file_manager>(
-      *sourceManager, *programBundle->root_program());
+      *sourceManager, *CHECK_NOTNULL(programBundle->root_program()));
 
   return Fixture{
       .source_manager_ = std::move(sourceManager),
@@ -187,6 +189,19 @@ struct S {}
 )");
 }
 
+TEST_F(FileManagerTest, add_include_trailing_comment) {
+  Fixture fixture = prepare_fixture(R"(
+include "thrift/annotation/thrift.thrift" // foo
+)");
+
+  fixture.file_manager_->add_include("some/other/file.thrift");
+
+  EXPECT_EQ(fixture.file_manager_->get_new_content(), R"(
+include "thrift/annotation/thrift.thrift" // foo
+include "some/other/file.thrift"
+)");
+}
+
 TEST_F(FileManagerTest, has_thrift_include) {
   Fixture fixture = prepare_fixture(R"(
 include "thrift/annotation/thrift.thrift"