[clang-format] Rework removeBraces() in Format.cpp

[owenpan]

Summary

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

[owenpan]

clang/lib/Format/Format.cpp

1914

Deleted by mistake. Will add it back before landing.

[MyDeveloperDay]

LGTM

[owenpan]

clang/lib/Format/Format.cpp

1914

Actually, it can (and should) be deleted as it was related to line 1918 below.

clang/lib/Format/Format.cpp

analyze(TokenAnnotator &Annotator,
     removeBraces(AnnotatedLines, Result);
     return {Result, 0};
   }
 
 private:
   void removeBraces(SmallVectorImpl<AnnotatedLine *> &Lines,
                     tooling::Replacements &Result) {
     const auto &SourceMgr = Env.getSourceManager();
-    bool EndsWithComment = false;
-    for (AnnotatedLine *Line : Lines) {
+    const auto End = Lines.end();
+    for (auto I = Lines.begin(); I != End; ++I) {
+      const auto Line = *I;
       removeBraces(Line->Children, Result);
-      if (Line->Affected) {
-        for (FormatToken *Token = Line->First; Token && !Token->Finalized;
-             Token = Token->Next) {
-          if (!Token->Optional)
-            continue;
-          assert(Token->isOneOf(tok::l_brace, tok::r_brace));
-          assert(Token->Previous || Token == Line->First);
-          const FormatToken *Next = Token->Next;
-          assert(Next || Token == Line->Last);
-          const auto Start =
-              (!Token->Previous && EndsWithComment) ||
-                      (Next && !(Next->isOneOf(tok::kw_else, tok::comment) &&
-                                 Next->NewlinesBefore > 0))
-                  ? Token->Tok.getLocation()
-                  : Token->WhitespaceRange.getBegin();
-          const auto Range =
-              CharSourceRange::getCharRange(Start, Token->Tok.getEndLoc());
-          cantFail(Result.add(tooling::Replacement(SourceMgr, Range, "")));
-        }
-      }
-      assert(Line->Last);
-      EndsWithComment = Line->Last->is(tok::comment);
-    }
+      if (!Line->Affected)
+        continue;
+      const auto NextLine = I + 1 == End ? nullptr : I[1];
+      for (auto Token = Line->First; Token && !Token->Finalized;
+        auto Next = Token->Next;
+        if (!Next && NextLine)
+          Next = NextLine->First;
+            Next && Next->NewlinesBefore == 0 && Next->isNot(tok::eof)
   }
 };
 
 class JavaScriptRequoter : public TokenAnalyzer {
 public:
   JavaScriptRequoter(const Environment &Env, const FormatStyle &Style)
       : TokenAnalyzer(Env, Style) {}
 

clang/unittests/Format/FormatTest.cpp

verifyFormat("if (a)\n"
                "{ //\n"
                "  foo();\n"
                "}",
                "if (a) { //\n"
                "  foo();\n"
                "}",
                Style);
 
+  verifyFormat("if (a) // comment\n"
+               "  b = 1;",
+               "if (a) // comment\n"
+               "{\n"
+               "  b = 1;\n"
+               "}",
+               Style);
+
   Style.ColumnLimit = 20;
 
   verifyFormat("int ab = [](int i) {\n"
                "  if (i > 0)\n"
                "  {\n"
                "    i = 12345678 -\n"
                "        i;\n"
                "  }\n"