[clang-format] Correctly count annoated lines in a namespace body

[owenpan]

Summary

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

[owenpan]

clang/lib/Format/TokenAnnotator.h

95

Will fix it before landing.

clang/lib/Format/NamespaceEndCommentsFixer.cpp

if (!EndCommentNextTok && I + 1 < E)
       EndCommentNextTok = AnnotatedLines[I + 1]->First;
     bool AddNewline = EndCommentNextTok &&
                       EndCommentNextTok->NewlinesBefore == 0 &&
                       EndCommentNextTok->isNot(tok::eof);
     const std::string EndCommentText =
         computeEndCommentText(NamespaceName, AddNewline, NamespaceTok,
                               Style.SpacesInLineCommentPrefix.Minimum);
     if (!hasEndComment(EndCommentPrevTok)) {
-      bool isShort = I - StartLineIndex <= Style.ShortNamespaceLines + 1;
-      if (!isShort) {
+      unsigned LineCount = 0;
+      for (auto J = StartLineIndex + 1; J < I; ++J)
+        LineCount += AnnotatedLines[J]->size();
+      if (LineCount > Style.ShortNamespaceLines) {
         addEndComment(EndCommentPrevTok,
                       std::string(Style.SpacesBeforeTrailingComments, ' ') +
                           EndCommentText,
                       SourceMgr, &Fixes);
       }
     } else if (!validEndComment(EndCommentPrevTok, NamespaceName,
                                 NamespaceTok)) {
       updateEndComment(EndCommentPrevTok, EndCommentText, SourceMgr, &Fixes);
     }
     StartLineIndex = SIZE_MAX;
   }
   return {Fixes, 0};
 }
 
 } // namespace format
 } // namespace clang

clang/lib/Format/TokenAnnotator.h

void addChildren(const UnwrappedLineNode &Node, FormatToken *Current) {
     for (const auto &Child : Node.Children) {
       Children.push_back(new AnnotatedLine(Child));
       if (Children.back()->ContainsMacroCall)
         ContainsMacroCall = true;
       Current->Children.push_back(Children.back());
     }
   }
 
+  size_t size() const {
+    size_t Size = 1;
+    for (const auto *Child : Children)
+      Size += Child->size();
+    return Size;
+  }
+
   ~AnnotatedLine() {
     for (AnnotatedLine *Child : Children)
       delete Child;
     FormatToken *Current = First;
     while (Current) {
       Current->Children.clear();
       Current->Role.reset();
       Current = Current->Next;

clang/unittests/Format/NamespaceEndCommentsFixerTest.cpp

EXPECT_EQ("namespace LongNamespace {\n"
             "int k;\n"
             "} // namespace LongNamespace\n",
             fixNamespaceEndComments("namespace LongNamespace {\n"
                                     "int i;\n"
                                     "int j;\n"
                                     "int k;\n"
                                     "}\n",
                                     Style));
+
+  // The namespace body has 5 unwrapped/annotated lines.
+  const std::string NestedLambdas{"namespace foo {\n"
+                                  "auto bar = [] {\n" // line 1
+                                  "  int i;\n"        // line 2
+                                  "  return [] {\n"   // line 3
+                                  "      int j;"      // line 4
+                                  "      return 0;\n" // line 5
+                                  "  };\n"            // part of line 3
+                                  "};\n"              // part of line 1
+                                  "}"};
+  Style.ShortNamespaceLines = 4;
+  EXPECT_EQ(NestedLambdas + " // namespace foo",
+            fixNamespaceEndComments(NestedLambdas, Style));
+  ++Style.ShortNamespaceLines;
+  EXPECT_EQ(NestedLambdas, fixNamespaceEndComments(NestedLambdas, Style));
 }
 
 TEST_F(ShortNamespaceLinesTest, NamespaceAlias) {
   auto Style = getLLVMStyle();
 
   EXPECT_EQ("namespace n = nn;\n"
             "{\n"
             "  int i;\n"