From e33243c950ac40d027ad8facbf7ccf0624604a16 Mon Sep 17 00:00:00 2001
From: Owen Pan <owenpiano@gmail.com>
Date: Sat, 3 Dec 2022 07:16:32 -0800
Subject: [PATCH] [clang-format] Link the braces of a block in
 UnwrappedLineParser

This includes TT_MacroBlockBegin and TT_MacroBlockEnd as well.

We can no longer use MatchingParen of FormatToken as an indicator
to mark optional braces. Instead, we directly set Optional of an
l_brace first and reset it later if it turns out that the braces
are not optional.

Also added a test case for deeply nested loops.

Differential Revision: https://reviews.llvm.org/D139257
---
 clang/lib/Format/UnwrappedLineParser.cpp     | 45 +++++++++++++++-------------
 clang/unittests/Format/BracesRemoverTest.cpp | 14 +++++++++
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/clang/lib/Format/UnwrappedLineParser.cpp b/clang/lib/Format/UnwrappedLineParser.cpp
index e6bb9c4..d665a4b 100644
--- a/clang/lib/Format/UnwrappedLineParser.cpp
+++ b/clang/lib/Format/UnwrappedLineParser.cpp
@@ -934,6 +934,9 @@ FormatToken *UnwrappedLineParser::parseBlock(
     return IfLBrace;
   }
 
+  Tok->MatchingParen = FormatTok;
+  FormatTok->MatchingParen = Tok;
+
   const bool IsFunctionRBrace =
       FormatTok->is(tok::r_brace) && Tok->is(TT_FunctionLBrace);
 
@@ -967,10 +970,7 @@ FormatToken *UnwrappedLineParser::parseBlock(
     }
     return mightFitOnOneLine((*CurrentLines)[Index], Tok);
   };
-  if (RemoveBraces()) {
-    Tok->MatchingParen = FormatTok;
-    FormatTok->MatchingParen = Tok;
-  }
+  Tok->Optional = RemoveBraces();
 
   size_t PPEndHash = computePPHash();
 
@@ -2707,10 +2707,20 @@ static void markOptionalBraces(FormatToken *LeftBrace) {
 
   assert(RightBrace->is(tok::r_brace));
   assert(RightBrace->MatchingParen == LeftBrace);
-  assert(LeftBrace->Optional == RightBrace->Optional);
 
-  LeftBrace->Optional = true;
-  RightBrace->Optional = true;
+  RightBrace->Optional = LeftBrace->Optional;
+}
+
+static void resetOptional(FormatToken *LeftBrace) {
+  if (!LeftBrace)
+    return;
+
+  const auto *RightBrace = LeftBrace->MatchingParen;
+  const bool IsOptionalRightBrace = RightBrace && RightBrace->Optional;
+  assert(LeftBrace->Optional || !IsOptionalRightBrace);
+
+  if (!IsOptionalRightBrace)
+    LeftBrace->Optional = false;
 }
 
 void UnwrappedLineParser::handleAttributes() {
@@ -2770,8 +2780,7 @@ FormatToken *UnwrappedLineParser::parseIfThenElse(IfStmtKind *IfKind,
 
   if (Style.RemoveBracesLLVM) {
     assert(!NestedTooDeep.empty());
-    KeepIfBraces = KeepIfBraces ||
-                   (IfLeftBrace && !IfLeftBrace->MatchingParen) ||
+    KeepIfBraces = KeepIfBraces || (IfLeftBrace && !IfLeftBrace->Optional) ||
                    NestedTooDeep.back() || IfBlockKind == IfStmtKind::IfOnly ||
                    IfBlockKind == IfStmtKind::IfElseIf;
   }
@@ -2802,8 +2811,9 @@ FormatToken *UnwrappedLineParser::parseIfThenElse(IfStmtKind *IfKind,
                          ElseBlockKind == IfStmtKind::IfElseIf;
       } else if (FollowedByIf && IfLBrace && !IfLBrace->Optional) {
         KeepElseBraces = true;
+        assert(ElseLeftBrace->Optional);
         assert(ElseLeftBrace->MatchingParen);
-        markOptionalBraces(ElseLeftBrace);
+        ElseLeftBrace->MatchingParen->Optional = true;
       }
       addUnwrappedLine();
     } else if (FormatTok->is(tok::kw_if)) {
@@ -2838,7 +2848,7 @@ FormatToken *UnwrappedLineParser::parseIfThenElse(IfStmtKind *IfKind,
 
   assert(!NestedTooDeep.empty());
   KeepElseBraces = KeepElseBraces ||
-                   (ElseLeftBrace && !ElseLeftBrace->MatchingParen) ||
+                   (ElseLeftBrace && !ElseLeftBrace->Optional) ||
                    NestedTooDeep.back();
 
   NestedTooDeep.pop_back();
@@ -2846,17 +2856,11 @@ FormatToken *UnwrappedLineParser::parseIfThenElse(IfStmtKind *IfKind,
   if (!KeepIfBraces && !KeepElseBraces) {
     markOptionalBraces(IfLeftBrace);
     markOptionalBraces(ElseLeftBrace);
-  } else if (IfLeftBrace) {
-    FormatToken *IfRightBrace = IfLeftBrace->MatchingParen;
-    if (IfRightBrace) {
-      assert(IfRightBrace->MatchingParen == IfLeftBrace);
-      assert(!IfLeftBrace->Optional);
-      assert(!IfRightBrace->Optional);
-      IfLeftBrace->MatchingParen = nullptr;
-      IfRightBrace->MatchingParen = nullptr;
-    }
   }
 
+  resetOptional(IfLeftBrace);
+  resetOptional(ElseLeftBrace);
+
   if (IfKind)
     *IfKind = Kind;
 
@@ -3071,6 +3075,7 @@ void UnwrappedLineParser::parseLoopBody(bool KeepBraces, bool WrapRightBrace) {
       if (!NestedTooDeep.back())
         markOptionalBraces(LeftBrace);
     }
+    resetOptional(LeftBrace);
     if (WrapRightBrace)
       addUnwrappedLine();
   } else {
diff --git a/clang/unittests/Format/BracesRemoverTest.cpp b/clang/unittests/Format/BracesRemoverTest.cpp
index 9c33a70..dea551c9 100644
--- a/clang/unittests/Format/BracesRemoverTest.cpp
+++ b/clang/unittests/Format/BracesRemoverTest.cpp
@@ -204,6 +204,20 @@ TEST_F(BracesRemoverTest, RemoveBraces) {
                "while (j < 0) { ++j; }",
                Style);
 
+  verifyFormat("for (;;) {\n"
+               "  for (;;)\n"
+               "    for (;;)\n"
+               "      a;\n"
+               "}",
+               "for (;;) {\n"
+               "  for (;;) {\n"
+               "    for (;;) {\n"
+               "      a;\n"
+               "    }\n"
+               "  }\n"
+               "}",
+               Style);
+
   verifyFormat("if (a)\n"
                "  b; // comment\n"
                "else if (c)\n"
-- 
2.7.4