[clang-format] Supress aligning of trailing namespace comments

Fixes https://github.com/llvm/llvm-project/issues/57504.

Differential Revision: https://reviews.llvm.org/D138263

5 files changed, 125 insertions, 7 deletions

diff --git a/clang/lib/Format/FormatToken.h b/clang/lib/Format/FormatToken.h
index 4e45478..6ce9249 100644
--- a/clang/lib/Format/FormatToken.h
+++ b/clang/lib/Format/FormatToken.h
@@ -102,7 +102,9 @@ namespace format {
   TYPE(MacroBlockBegin)                                                        \
   TYPE(MacroBlockEnd)                                                          \
   TYPE(ModulePartitionColon)                                                   \
+  TYPE(NamespaceLBrace)                                                        \
   TYPE(NamespaceMacro)                                                         \
+  TYPE(NamespaceRBrace)                                                        \
   TYPE(NonNullAssertion)                                                       \
   TYPE(NullCoalescingEqual)                                                    \
   TYPE(NullCoalescingOperator)                                                 \
diff --git a/clang/lib/Format/UnwrappedLineParser.cpp b/clang/lib/Format/UnwrappedLineParser.cpp
index 5c6e7e2..22fa651 100644
--- a/clang/lib/Format/UnwrappedLineParser.cpp
+++ b/clang/lib/Format/UnwrappedLineParser.cpp
@@ -806,6 +806,9 @@ FormatToken *UnwrappedLineParser::parseBlock(
     return IfLBrace;
   }
 
+  if (FormatTok->is(tok::r_brace) && Tok->is(TT_NamespaceLBrace))
+    FormatTok->setFinalizedType(TT_NamespaceRBrace);
+
   const bool IsFunctionRBrace =
       FormatTok->is(tok::r_brace) && Tok->is(TT_FunctionLBrace);
 
@@ -2959,6 +2962,8 @@ void UnwrappedLineParser::parseNamespace() {
     }
   }
   if (FormatTok->is(tok::l_brace)) {
+    FormatTok->setFinalizedType(TT_NamespaceLBrace);
+
     if (ShouldBreakBeforeBrace(Style, InitialToken))
       addUnwrappedLine();
 
diff --git a/clang/lib/Format/WhitespaceManager.cpp b/clang/lib/Format/WhitespaceManager.cpp
index 668ca38..61e59d0 100644
--- a/clang/lib/Format/WhitespaceManager.cpp
+++ b/clang/lib/Format/WhitespaceManager.cpp
@@ -1081,11 +1081,10 @@ void WhitespaceManager::alignTrailingComments() {
 
     if (i + 1 != e && Changes[i + 1].ContinuesPPDirective)
       ChangeMaxColumn -= 2;
-    // If this comment follows an } in column 0, it probably documents the
-    // closing of a namespace and we don't want to align it.
-    bool FollowsRBraceInColumn0 = i > 0 && Changes[i].NewlinesBefore == 0 &&
-                                  Changes[i - 1].Tok->is(tok::r_brace) &&
-                                  Changes[i - 1].StartOfTokenColumn == 0;
+
+    // We don't want to align namespace end comments.
+    bool DontAlignThisComment = i > 0 && Changes[i].NewlinesBefore == 0 &&
+                                Changes[i - 1].Tok->is(TT_NamespaceRBrace);
     bool WasAlignedWithStartOfNextLine = false;
     if (Changes[i].NewlinesBefore >= 1) { // A comment on its own line.
       unsigned CommentColumn = SourceMgr.getSpellingColumnNumber(
@@ -1105,7 +1104,7 @@ void WhitespaceManager::alignTrailingComments() {
       }
     }
     if (Style.AlignTrailingComments.Kind == FormatStyle::TCAS_Never ||
-        FollowsRBraceInColumn0) {
+        DontAlignThisComment) {
       alignTrailingComments(StartOfSequence, i, MinColumn);
       MinColumn = ChangeMinColumn;
       MaxColumn = ChangeMinColumn;
diff --git a/clang/unittests/Format/FormatTest.cpp b/clang/unittests/Format/FormatTest.cpp
index ed0c5e6..3ed97c0 100644
--- a/clang/unittests/Format/FormatTest.cpp
+++ b/clang/unittests/Format/FormatTest.cpp
@@ -20316,7 +20316,7 @@ TEST_F(FormatTest, WhitesmithsBraceBreaking) {
                "      int x;\n"
                "      };\n"
                "    } // namespace b\n"
-               "  }   // namespace a",
+               "  } // namespace a",
                WhitesmithsBraceStyle);
 
   verifyFormat("void f()\n"
diff --git a/clang/unittests/Format/FormatTestComments.cpp b/clang/unittests/Format/FormatTestComments.cpp
index 2406fd6..001c5bf 100644
--- a/clang/unittests/Format/FormatTestComments.cpp
+++ b/clang/unittests/Format/FormatTestComments.cpp
@@ -3087,6 +3087,118 @@ TEST_F(FormatTestComments, AlignTrailingCommentsLeave) {
                    Style));
 }
 
+TEST_F(FormatTestComments, DontAlignNamespaceComments) {
+  FormatStyle Style = getLLVMStyle();
+  Style.NamespaceIndentation = FormatStyle::NI_All;
+  Style.NamespaceMacros.push_back("TESTSUITE");
+  Style.ShortNamespaceLines = 0;
+
+  StringRef Input = "namespace A {\n"
+                    "  TESTSUITE(B) {\n"
+                    "    namespace C {\n"
+                    "      namespace D {} // namespace D\n"
+                    "      std::string Foo = Bar; // Comment\n"
+                    "      std::string BazString = Baz;   // C2\n"
+                    "    }          // namespace C\n"
+                    "  }\n"
+                    "} // NaMeSpAcE A";
+
+  EXPECT_TRUE(Style.FixNamespaceComments);
+  EXPECT_EQ(Style.AlignTrailingComments.Kind, FormatStyle::TCAS_Always);
+  verifyFormat("namespace A {\n"
+               "  TESTSUITE(B) {\n"
+               "    namespace C {\n"
+               "      namespace D {} // namespace D\n"
+               "      std::string Foo = Bar;       // Comment\n"
+               "      std::string BazString = Baz; // C2\n"
+               "    } // namespace C\n"
+               "  } // TESTSUITE(B)\n"
+               "} // NaMeSpAcE A",
+               Input, Style);
+
+  Style.AlignTrailingComments.Kind = FormatStyle::TCAS_Never;
+  verifyFormat("namespace A {\n"
+               "  TESTSUITE(B) {\n"
+               "    namespace C {\n"
+               "      namespace D {} // namespace D\n"
+               "      std::string Foo = Bar; // Comment\n"
+               "      std::string BazString = Baz; // C2\n"
+               "    } // namespace C\n"
+               "  } // TESTSUITE(B)\n"
+               "} // NaMeSpAcE A",
+               Input, Style);
+
+  Style.AlignTrailingComments.Kind = FormatStyle::TCAS_Leave;
+  verifyFormat("namespace A {\n"
+               "  TESTSUITE(B) {\n"
+               "    namespace C {\n"
+               "      namespace D {} // namespace D\n"
+               "      std::string Foo = Bar; // Comment\n"
+               "      std::string BazString = Baz;   // C2\n"
+               "    }          // namespace C\n"
+               "  } // TESTSUITE(B)\n"
+               "} // NaMeSpAcE A",
+               Input, Style);
+
+  Style.FixNamespaceComments = false;
+  Style.AlignTrailingComments.Kind = FormatStyle::TCAS_Always;
+  verifyFormat("namespace A {\n"
+               "  TESTSUITE(B) {\n"
+               "    namespace C {\n"
+               "      namespace D {} // namespace D\n"
+               "      std::string Foo = Bar;       // Comment\n"
+               "      std::string BazString = Baz; // C2\n"
+               "    } // namespace C\n"
+               "  }\n"
+               "} // NaMeSpAcE A",
+               Input, Style);
+
+  Style.AlignTrailingComments.Kind = FormatStyle::TCAS_Never;
+  verifyFormat("namespace A {\n"
+               "  TESTSUITE(B) {\n"
+               "    namespace C {\n"
+               "      namespace D {} // namespace D\n"
+               "      std::string Foo = Bar; // Comment\n"
+               "      std::string BazString = Baz; // C2\n"
+               "    } // namespace C\n"
+               "  }\n"
+               "} // NaMeSpAcE A",
+               Input, Style);
+
+  Style.AlignTrailingComments.Kind = FormatStyle::TCAS_Leave;
+  verifyFormat("namespace A {\n"
+               "  TESTSUITE(B) {\n"
+               "    namespace C {\n"
+               "      namespace D {} // namespace D\n"
+               "      std::string Foo = Bar; // Comment\n"
+               "      std::string BazString = Baz;   // C2\n"
+               "    }          // namespace C\n"
+               "  }\n"
+               "} // NaMeSpAcE A",
+               Input, Style);
+
+  Style.AlignTrailingComments.Kind = FormatStyle::TCAS_Always;
+  Style.FixNamespaceComments = true;
+  Input = "namespace A {\n"
+          "  int Foo;\n"
+          "  int Bar;\n"
+          "}\n"
+          "// Comment";
+
+#if 0
+  // FIXME: The following comment is aligned with the namespace comment.
+  verifyFormat("namespace A {\n"
+               "  int Foo;\n"
+               "  int Bar;\n"
+               "} // namespace A\n"
+               " // Comment",
+               Input, Style);
+#endif
+
+  Style.FixNamespaceComments = false;
+  verifyFormat(Input, Style);
+}
+
 TEST_F(FormatTestComments, AlignsBlockCommentDecorations) {
   EXPECT_EQ("/*\n"
             " */",