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

Fixes #63882.

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

Author: owenca

clang/lib/Format/NamespaceEndCommentsFixer.cpp

@@ -359,8 +359,10 @@ std::pair<tooling::Replacements, unsigned> NamespaceEndCommentsFixer::analyze(
         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,

clang/lib/Format/TokenAnnotator.h

@@ -91,6 +91,13 @@ class AnnotatedLine {
     }
   }
 
+  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;

clang/unittests/Format/NamespaceEndCommentsFixerTest.cpp

@@ -1376,6 +1376,22 @@ TEST_F(ShortNamespaceLinesTest, MultipleUnwrappedLine) {
                                     "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) {