[clang-tools-extra] [clang-reorder-fields] Move trailing comments. (PR #122918)

[Clement Courbet via cfe-commits]

https://github.com/legrosbuffle updated https://github.com/llvm/llvm-project/pull/122918

>From 5180e6ee06531623adc35ac418d2dddbe58586c3 Mon Sep 17 00:00:00 2001
From: Clement Courbet <courbet at google.com>
Date: Tue, 14 Jan 2025 13:38:00 +0000
Subject: [PATCH] [clang-reorder-fields] Move trailing comments.

Currently, trailing comments get mixed up:

```
struct Foo {
  int a; // This one is the cool field
         // within the struct.
  int b;
};
```

becomes:

```
struct Foo {
  int b; // This one is the cool field
         // within the struct.
  int a;
};
```

This should be:

```
struct Foo {
  int b;
  int a; // This one is the cool field
         // within the struct.
};
```
---
 .../ReorderFieldsAction.cpp                   | 69 ++++++++++++++++---
 .../test/clang-reorder-fields/Comments.cpp    | 23 +++++++
 2 files changed, 84 insertions(+), 8 deletions(-)
 create mode 100644 clang-tools-extra/test/clang-reorder-fields/Comments.cpp

diff --git a/clang-tools-extra/clang-reorder-fields/ReorderFieldsAction.cpp b/clang-tools-extra/clang-reorder-fields/ReorderFieldsAction.cpp
index dc3a3b6211b7e4..80ee31368fe9a5 100644
--- a/clang-tools-extra/clang-reorder-fields/ReorderFieldsAction.cpp
+++ b/clang-tools-extra/clang-reorder-fields/ReorderFieldsAction.cpp
@@ -63,7 +63,9 @@ getNewFieldsOrder(const RecordDecl *Definition,
     NameToIndex[Field->getName()] = Field->getFieldIndex();
 
   if (DesiredFieldsOrder.size() != NameToIndex.size()) {
-    llvm::errs() << "Number of provided fields doesn't match definition.\n";
+    llvm::errs() << "Number of provided fields (" << DesiredFieldsOrder.size()
+                 << ") doesn't match definition (" << NameToIndex.size()
+                 << ").\n";
     return {};
   }
   SmallVector<unsigned, 4> NewFieldsOrder;
@@ -116,26 +118,77 @@ findMembersUsedInInitExpr(const CXXCtorInitializer *Initializer,
   return Results;
 }
 
-/// Returns the full source range for the field declaration up to (not
-/// including) the trailing semicolumn, including potential macro invocations,
-/// e.g. `int a GUARDED_BY(mu);`.
+/// Returns the next token after `Loc` (including comment tokens).
+static std::optional<Token> getTokenAfter(SourceLocation Loc,
+                                          const SourceManager &SM,
+                                          const LangOptions &LangOpts) {
+  if (Loc.isMacroID()) {
+    return std::nullopt;
+  }
+  Loc = Lexer::getLocForEndOfToken(Loc, 0, SM, LangOpts);
+
+  // Break down the source location.
+  std::pair<FileID, unsigned> LocInfo = SM.getDecomposedLoc(Loc);
+
+  // Try to load the file buffer.
+  bool InvalidTemp = false;
+  StringRef File = SM.getBufferData(LocInfo.first, &InvalidTemp);
+  if (InvalidTemp)
+    return std::nullopt;
+
+  const char *TokenBegin = File.data() + LocInfo.second;
+
+  Lexer lexer(SM.getLocForStartOfFile(LocInfo.first), LangOpts, File.begin(),
+              TokenBegin, File.end());
+  lexer.SetCommentRetentionState(true);
+  // Find the token.
+  Token Tok;
+  lexer.LexFromRawLexer(Tok);
+  return Tok;
+}
+
+/// Returns the end of the trailing comments after `Loc`.
+static SourceLocation getEndOfTrailingComment(SourceLocation Loc,
+                                              const SourceManager &SM,
+                                              const LangOptions &LangOpts) {
+  // We consider any following comment token that is indented more than the
+  // first comment to be part of the trailing comment.
+  const unsigned Column = SM.getPresumedColumnNumber(Loc);
+  std::optional<Token> Tok = getTokenAfter(Loc, SM, LangOpts);
+  while (Tok && Tok->is(tok::comment) &&
+         SM.getPresumedColumnNumber(Tok->getLocation()) > Column) {
+    Loc = Tok->getEndLoc();
+    Tok = getTokenAfter(Loc, SM, LangOpts);
+  }
+  return Loc;
+}
+
+/// Returns the full source range for the field declaration up to (including)
+/// the trailing semicolumn, including potential macro invocations,
+/// e.g. `int a GUARDED_BY(mu);`. If there is a trailing comment, include it.
 static SourceRange getFullFieldSourceRange(const FieldDecl &Field,
                                            const ASTContext &Context) {
-  SourceRange Range = Field.getSourceRange();
+  const SourceRange Range = Field.getSourceRange();
+  SourceLocation Begin = Range.getBegin();
   SourceLocation End = Range.getEnd();
   const SourceManager &SM = Context.getSourceManager();
   const LangOptions &LangOpts = Context.getLangOpts();
   while (true) {
     std::optional<Token> CurrentToken = Lexer::findNextToken(End, SM, LangOpts);
 
-    if (!CurrentToken || CurrentToken->is(tok::semi))
-      break;
+    if (!CurrentToken)
+      return SourceRange(Begin, End);
 
     if (CurrentToken->is(tok::eof))
       return Range; // Something is wrong, return the original range.
+
     End = CurrentToken->getLastLoc();
+
+    if (CurrentToken->is(tok::semi))
+      break;
   }
-  return SourceRange(Range.getBegin(), End);
+  End = getEndOfTrailingComment(End, SM, LangOpts);
+  return SourceRange(Begin, End);
 }
 
 /// Reorders fields in the definition of a struct/class.
diff --git a/clang-tools-extra/test/clang-reorder-fields/Comments.cpp b/clang-tools-extra/test/clang-reorder-fields/Comments.cpp
new file mode 100644
index 00000000000000..a31b6692c9ac73
--- /dev/null
+++ b/clang-tools-extra/test/clang-reorder-fields/Comments.cpp
@@ -0,0 +1,23 @@
+// RUN: clang-reorder-fields -record-name Foo -fields-order e1,e3,e2,a,c,b %s -- | FileCheck %s
+
+class Foo {
+  int a; // Trailing comment for a.
+  int b; // Multiline
+         // trailing for b.
+  // Prefix comments for c.
+  int c;
+
+  /*c-like*/ int e1;
+  int /*c-like*/ e2;
+  int e3 /*c-like*/;
+};
+
+// CHECK:       /*c-like*/ int e1;
+// CHECK-NEXT:  int e3 /*c-like*/;
+// CHECK-NEXT:  int /*c-like*/ e2;
+// CHECK-NEXT:  int a; // Trailing comment for a.
+// CHECK-NEXT:  // Prefix comments for c.
+// CHECK-NEXT:  int c;
+// CHECK-NEXT:  int b; // Multiline
+// CHECK-NEXT:         // trailing for b.
+