1 files changed, 45 insertions, 23 deletions

diff --git a/clang/lib/Format/MacroCallReconstructor.cpp b/clang/lib/Format/MacroCallReconstructor.cpp
index cbdd168..101acefd 100644
--- a/clang/lib/Format/MacroCallReconstructor.cpp
+++ b/clang/lib/Format/MacroCallReconstructor.cpp
@@ -33,7 +33,7 @@ void forEachToken(const UnwrappedLine &Line, const T &Call,
                   FormatToken *Parent = nullptr) {
   bool First = true;
   for (const auto &N : Line.Tokens) {
-    Call(N.Tok, Parent, First);
+    Call(N.Tok, Parent, First, Line.Level);
     First = false;
     for (const auto &Child : N.Children)
       forEachToken(Child, Call, N.Tok);
@@ -44,7 +44,7 @@ MacroCallReconstructor::MacroCallReconstructor(
     unsigned Level,
     const llvm::DenseMap<FormatToken *, std::unique_ptr<UnwrappedLine>>
         &ActiveExpansions)
-    : Level(Level), IdToReconstructed(ActiveExpansions) {
+    : Result(Level), IdToReconstructed(ActiveExpansions) {
   Result.Tokens.push_back(std::make_unique<LineNode>());
   ActiveReconstructedLines.push_back(&Result);
 }
@@ -52,9 +52,8 @@ MacroCallReconstructor::MacroCallReconstructor(
 void MacroCallReconstructor::addLine(const UnwrappedLine &Line) {
   assert(State != Finalized);
   LLVM_DEBUG(llvm::dbgs() << "MCR: new line...\n");
-  forEachToken(Line, [&](FormatToken *Token, FormatToken *Parent, bool First) {
-    add(Token, Parent, First);
-  });
+  forEachToken(Line, [&](FormatToken *Token, FormatToken *Parent, bool First,
+                         unsigned Level) { add(Token, Parent, First, Level); });
   assert(InProgress || finished());
 }
 
@@ -62,8 +61,8 @@ UnwrappedLine MacroCallReconstructor::takeResult() && {
   finalize();
   assert(Result.Tokens.size() == 1 &&
          Result.Tokens.front()->Children.size() == 1);
-  UnwrappedLine Final =
-      createUnwrappedLine(*Result.Tokens.front()->Children.front(), Level);
+  UnwrappedLine Final = createUnwrappedLine(
+      *Result.Tokens.front()->Children.front(), Result.Level);
   assert(!Final.Tokens.empty());
   return Final;
 }
@@ -72,7 +71,8 @@ UnwrappedLine MacroCallReconstructor::takeResult() && {
 // ExpandedParent in the incoming unwrapped line. \p First specifies whether it
 // is the first token in a given unwrapped line.
 void MacroCallReconstructor::add(FormatToken *Token,
-                                 FormatToken *ExpandedParent, bool First) {
+                                 FormatToken *ExpandedParent, bool First,
+                                 unsigned Level) {
   LLVM_DEBUG(
       llvm::dbgs() << "MCR: Token: " << Token->TokenText << ", Parent: "
                    << (ExpandedParent ? ExpandedParent->TokenText : "<null>")
@@ -102,7 +102,7 @@ void MacroCallReconstructor::add(FormatToken *Token,
       First = true;
   }
 
-  prepareParent(ExpandedParent, First);
+  prepareParent(ExpandedParent, First, Level);
 
   if (Token->MacroCtx) {
     // If this token was generated by a macro call, add the reconstructed
@@ -129,7 +129,7 @@ void MacroCallReconstructor::add(FormatToken *Token,
 // is the parent of ActiveReconstructedLines.back() in the reconstructed
 // unwrapped line.
 void MacroCallReconstructor::prepareParent(FormatToken *ExpandedParent,
-                                           bool NewLine) {
+                                           bool NewLine, unsigned Level) {
   LLVM_DEBUG({
     llvm::dbgs() << "ParentMap:\n";
     debugParentMap();
@@ -172,7 +172,7 @@ void MacroCallReconstructor::prepareParent(FormatToken *ExpandedParent,
     }
     assert(!ActiveReconstructedLines.empty());
     ActiveReconstructedLines.back()->Tokens.back()->Children.push_back(
-        std::make_unique<ReconstructedLine>());
+        std::make_unique<ReconstructedLine>(Level));
     ActiveReconstructedLines.push_back(
         &*ActiveReconstructedLines.back()->Tokens.back()->Children.back());
   } else if (parentLine().Tokens.back()->Tok != Parent) {
@@ -424,7 +424,8 @@ bool MacroCallReconstructor::processNextReconstructed() {
       SpelledParentToReconstructedParent[MacroCallStructure.back()
                                              .ParentLastToken] = Token;
       appendToken(Token);
-      prepareParent(Token, /*NewLine=*/true);
+      prepareParent(Token, /*NewLine=*/true,
+                    MacroCallStructure.back().Line->Level);
       Token->MacroParent = true;
       return false;
     }
@@ -435,7 +436,8 @@ bool MacroCallReconstructor::processNextReconstructed() {
             [MacroCallStructure.back().Line->Tokens.back()->Tok] = Token;
         Token->MacroParent = true;
         appendToken(Token, MacroCallStructure.back().Line);
-        prepareParent(Token, /*NewLine=*/true);
+        prepareParent(Token, /*NewLine=*/true,
+                      MacroCallStructure.back().Line->Level);
         return true;
       }
       if (Token->is(tok::r_paren)) {
@@ -509,16 +511,36 @@ MacroCallReconstructor::createUnwrappedLine(const ReconstructedLine &Line,
   for (const auto &N : Line.Tokens) {
     Result.Tokens.push_back(N->Tok);
     UnwrappedLineNode &Current = Result.Tokens.back();
-    for (const auto &Child : N->Children) {
-      if (Child->Tokens.empty())
-        continue;
-      Current.Children.push_back(createUnwrappedLine(*Child, Level + 1));
-    }
-    if (Current.Children.size() == 1 &&
-        Current.Tok->isOneOf(tok::l_paren, tok::comma)) {
-      Result.Tokens.splice(Result.Tokens.end(),
-                           Current.Children.front().Tokens);
-      Current.Children.clear();
+    auto NumChildren =
+        std::count_if(N->Children.begin(), N->Children.end(),
+                      [](const auto &Child) { return !Child->Tokens.empty(); });
+    if (NumChildren == 1 && Current.Tok->isOneOf(tok::l_paren, tok::comma)) {
+      // If we only have one child, and the child is due to a macro expansion
+      // (either attached to a left parenthesis or comma), merge the child into
+      // the current line to prevent forced breaks for macro arguments.
+      auto *Child = std::find_if(
+          N->Children.begin(), N->Children.end(),
+          [](const auto &Child) { return !Child->Tokens.empty(); });
+      auto Line = createUnwrappedLine(**Child, Level);
+      Result.Tokens.splice(Result.Tokens.end(), Line.Tokens);
+    } else if (NumChildren > 0) {
+      // When there are multiple children with different indent, make sure that
+      // we indent them:
+      // 1. One level below the current line's level.
+      // 2. At the correct level relative to each other.
+      unsigned MinChildLevel =
+          std::min_element(N->Children.begin(), N->Children.end(),
+                           [](const auto &E1, const auto &E2) {
+                             return E1->Level < E2->Level;
+                           })
+              ->get()
+              ->Level;
+      for (const auto &Child : N->Children) {
+        if (Child->Tokens.empty())
+          continue;
+        Current.Children.push_back(createUnwrappedLine(
+            *Child, Level + 1 + (Child->Level - MinChildLevel)));
+      }
     }
   }
   return Result;