[clang-format] Fix a bug in inserting braces at trailing comments

[owenpan]

Summary

If the style wraps control statement braces, the opening braces should be inserted after the trailing comments if present.

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

clang/lib/Format/Format.cpp

for (AnnotatedLine *Line : Lines) {
         continue;
       for (FormatToken *Token = Line->First; Token && !Token->Finalized;
            Token = Token->Next) {
         if (Token->BraceCount == 0)
           continue;
         std::string Brace;
         if (Token->BraceCount < 0) {
           assert(Token->BraceCount == -1);
-          Brace = '{';
+          Brace = "\n{";
         } else {
           Brace = '\n' + std::string(Token->BraceCount, '}');
         }
         Token->BraceCount = 0;
         const auto Start = Token->Tok.getEndLoc();
         cantFail(Result.add(tooling::Replacement(SourceMgr, Start, 0, Brace)));
       }
     }

clang/lib/Format/UnwrappedLineParser.cpp

static FormatToken *getLastNonComment(const UnwrappedLine &Line) {
   return nullptr;
 }
 
 void UnwrappedLineParser::parseUnbracedBody(bool CheckEOF) {
   FormatToken *Tok = nullptr;
 
   if (Style.InsertBraces && !Line->InPPDirective && !Line->Tokens.empty() &&
       PreprocessorDirectives.empty()) {
-    Tok = getLastNonComment(*Line);
+    assert(!Line->Tokens.empty());
+    Tok = Style.BraceWrapping.AfterControlStatement == FormatStyle::BWACS_Never
+              ? getLastNonComment(*Line)
+              : Line->Tokens.back().Tok;
     assert(Tok);
     if (Tok->BraceCount < 0) {
       assert(Tok->BraceCount == -1);
       Tok = nullptr;
     } else {
       Tok->BraceCount = -1;
     }
   }

clang/unittests/Format/FormatTest.cpp

verifyFormat("if (a) {\n"
                "  f();\n"
                "}",
                "if (a)\n"
                "  // foo\n"
                "  // bar\n"
                "  f();",
                Style);
 
+  verifyFormat("if (a) { //\n"
+               "  b = 1;\n"
+               "}",
+               "if (a) //\n"
+               "  b = 1;",
+               Style);
+
   verifyFormat("if (a) { // comment\n"
                "  // comment\n"
                "  f();\n"
                "}",
                "if (a) // comment\n"
                "  // comment\n"
                "  f();",
                Style);
TEST_F(FormatTest, InsertBraces) {
 
   verifyFormat("if (a + b >\n"
                "    c) {\n"
                "  f();\n"
                "}",
                "if (a + b > c)\n"
                "  f();",
                Style);
+
+  Style.BreakBeforeBraces = FormatStyle::BS_Custom;
+  Style.BraceWrapping.AfterControlStatement = FormatStyle::BWACS_Always;
+  // TODO: Delete the following line after #57421 is fixed.
+  Style.BraceWrapping.AfterFunction = true;
+
+  verifyFormat("if (a) //\n"
+               "{\n"
+               "  b = 1;\n"
+               "}",
+               "if (a) //\n"
+               "  b = 1;",
+               Style);
 }
 
 TEST_F(FormatTest, RemoveBraces) {
   FormatStyle Style = getLLVMStyle();
   Style.RemoveBracesLLVM = true;
 
   // The following test cases are fully-braced versions of the examples at
   // "llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-