1 files changed, 175 insertions, 80 deletions

diff --git a/lib/Format/WhitespaceManager.cpp b/lib/Format/WhitespaceManager.cpp
index 65395277f89a..725f05bcd8fc 100644
--- a/lib/Format/WhitespaceManager.cpp
+++ b/lib/Format/WhitespaceManager.cpp
@@ -26,16 +26,18 @@ operator()(const Change &C1, const Change &C2) const {
 }
 
 WhitespaceManager::Change::Change(
-    bool CreateReplacement, const SourceRange &OriginalWhitespaceRange,
+    bool CreateReplacement, SourceRange OriginalWhitespaceRange,
     unsigned IndentLevel, int Spaces, unsigned StartOfTokenColumn,
     unsigned NewlinesBefore, StringRef PreviousLinePostfix,
-    StringRef CurrentLinePrefix, tok::TokenKind Kind, bool ContinuesPPDirective)
+    StringRef CurrentLinePrefix, tok::TokenKind Kind, bool ContinuesPPDirective,
+    bool IsStartOfDeclName)
     : CreateReplacement(CreateReplacement),
       OriginalWhitespaceRange(OriginalWhitespaceRange),
       StartOfTokenColumn(StartOfTokenColumn), NewlinesBefore(NewlinesBefore),
       PreviousLinePostfix(PreviousLinePostfix),
       CurrentLinePrefix(CurrentLinePrefix), Kind(Kind),
-      ContinuesPPDirective(ContinuesPPDirective), IndentLevel(IndentLevel),
+      ContinuesPPDirective(ContinuesPPDirective),
+      IsStartOfDeclName(IsStartOfDeclName), IndentLevel(IndentLevel),
       Spaces(Spaces), IsTrailingComment(false), TokenLength(0),
       PreviousEndOfTokenColumn(0), EscapedNewlineColumn(0),
       StartOfBlockComment(nullptr), IndentationOffset(0) {}
@@ -52,19 +54,21 @@ void WhitespaceManager::replaceWhitespace(FormatToken &Tok, unsigned Newlines,
   if (Tok.Finalized)
     return;
   Tok.Decision = (Newlines > 0) ? FD_Break : FD_Continue;
-  Changes.push_back(Change(true, Tok.WhitespaceRange, IndentLevel, Spaces,
-                           StartOfTokenColumn, Newlines, "", "",
-                           Tok.Tok.getKind(), InPPDirective && !Tok.IsFirst));
+  Changes.push_back(
+      Change(true, Tok.WhitespaceRange, IndentLevel, Spaces, StartOfTokenColumn,
+             Newlines, "", "", Tok.Tok.getKind(), InPPDirective && !Tok.IsFirst,
+             Tok.is(TT_StartOfName) || Tok.is(TT_FunctionDeclarationName)));
 }
 
 void WhitespaceManager::addUntouchableToken(const FormatToken &Tok,
                                             bool InPPDirective) {
   if (Tok.Finalized)
     return;
-  Changes.push_back(Change(false, Tok.WhitespaceRange, /*IndentLevel=*/0,
-                           /*Spaces=*/0, Tok.OriginalColumn, Tok.NewlinesBefore,
-                           "", "", Tok.Tok.getKind(),
-                           InPPDirective && !Tok.IsFirst));
+  Changes.push_back(
+      Change(false, Tok.WhitespaceRange, /*IndentLevel=*/0,
+             /*Spaces=*/0, Tok.OriginalColumn, Tok.NewlinesBefore, "", "",
+             Tok.Tok.getKind(), InPPDirective && !Tok.IsFirst,
+             Tok.is(TT_StartOfName) || Tok.is(TT_FunctionDeclarationName)));
 }
 
 void WhitespaceManager::replaceWhitespaceInToken(
@@ -84,7 +88,8 @@ void WhitespaceManager::replaceWhitespaceInToken(
       // calculate the new length of the comment and to calculate the changes
       // for which to do the alignment when aligning comments.
       Tok.is(TT_LineComment) && Newlines > 0 ? tok::comment : tok::unknown,
-      InPPDirective && !Tok.IsFirst));
+      InPPDirective && !Tok.IsFirst,
+      Tok.is(TT_StartOfName) || Tok.is(TT_FunctionDeclarationName)));
 }
 
 const tooling::Replacements &WhitespaceManager::generateReplacements() {
@@ -93,6 +98,7 @@ const tooling::Replacements &WhitespaceManager::generateReplacements() {
 
   std::sort(Changes.begin(), Changes.end(), Change::IsBeforeInFile(SourceMgr));
   calculateLineBreakInformation();
+  alignConsecutiveDeclarations();
   alignConsecutiveAssignments();
   alignTrailingComments();
   alignEscapedNewlines();
@@ -142,94 +148,183 @@ void WhitespaceManager::calculateLineBreakInformation() {
   }
 }
 
-// Walk through all of the changes and find sequences of "=" to align.  To do
-// so, keep track of the lines and whether or not an "=" was found on align. If
-// a "=" is found on a line, extend the current sequence. If the current line
-// cannot be part of a sequence, e.g. because there is an empty line before it
-// or it contains non-assignments, finalize the previous sequence.
-void WhitespaceManager::alignConsecutiveAssignments() {
-  if (!Style.AlignConsecutiveAssignments)
-    return;
+// Align a single sequence of tokens, see AlignTokens below.
+template <typename F>
+static void
+AlignTokenSequence(unsigned Start, unsigned End, unsigned Column, F &&Matches,
+                   SmallVector<WhitespaceManager::Change, 16> &Changes) {
+  bool FoundMatchOnLine = false;
+  int Shift = 0;
+  for (unsigned i = Start; i != End; ++i) {
+    if (Changes[i].NewlinesBefore > 0) {
+      FoundMatchOnLine = false;
+      Shift = 0;
+    }
+
+    // If this is the first matching token to be aligned, remember by how many
+    // spaces it has to be shifted, so the rest of the changes on the line are
+    // shifted by the same amount
+    if (!FoundMatchOnLine && Matches(Changes[i])) {
+      FoundMatchOnLine = true;
+      Shift = Column - Changes[i].StartOfTokenColumn;
+      Changes[i].Spaces += Shift;
+    }
+
+    assert(Shift >= 0);
+    Changes[i].StartOfTokenColumn += Shift;
+    if (i + 1 != Changes.size())
+      Changes[i + 1].PreviousEndOfTokenColumn += Shift;
+  }
+}
 
+// Walk through all of the changes and find sequences of matching tokens to
+// align. To do so, keep track of the lines and whether or not a matching token
+// was found on a line. If a matching token is found, extend the current
+// sequence. If the current line cannot be part of a sequence, e.g. because
+// there is an empty line before it or it contains only non-matching tokens,
+// finalize the previous sequence.
+template <typename F>
+static void AlignTokens(const FormatStyle &Style, F &&Matches,
+                        SmallVector<WhitespaceManager::Change, 16> &Changes) {
   unsigned MinColumn = 0;
+  unsigned MaxColumn = UINT_MAX;
+
+  // Line number of the start and the end of the current token sequence.
   unsigned StartOfSequence = 0;
   unsigned EndOfSequence = 0;
-  bool FoundAssignmentOnLine = false;
-  bool FoundLeftParenOnLine = false;
-  unsigned CurrentLine = 0;
 
-  auto AlignSequence = [&] {
-    alignConsecutiveAssignments(StartOfSequence, EndOfSequence, MinColumn);
+  // Keep track of the nesting level of matching tokens, i.e. the number of
+  // surrounding (), [], or {}. We will only align a sequence of matching
+  // token that share the same scope depth.
+  //
+  // FIXME: This could use FormatToken::NestingLevel information, but there is
+  // an outstanding issue wrt the brace scopes.
+  unsigned NestingLevelOfLastMatch = 0;
+  unsigned NestingLevel = 0;
+
+  // Keep track of the number of commas before the matching tokens, we will only
+  // align a sequence of matching tokens if they are preceded by the same number
+  // of commas.
+  unsigned CommasBeforeLastMatch = 0;
+  unsigned CommasBeforeMatch = 0;
+
+  // Whether a matching token has been found on the current line.
+  bool FoundMatchOnLine = false;
+
+  // Aligns a sequence of matching tokens, on the MinColumn column.
+  //
+  // Sequences start from the first matching token to align, and end at the
+  // first token of the first line that doesn't need to be aligned.
+  //
+  // We need to adjust the StartOfTokenColumn of each Change that is on a line
+  // containing any matching token to be aligned and located after such token.
+  auto AlignCurrentSequence = [&] {
+    if (StartOfSequence > 0 && StartOfSequence < EndOfSequence)
+      AlignTokenSequence(StartOfSequence, EndOfSequence, MinColumn, Matches,
+                         Changes);
     MinColumn = 0;
+    MaxColumn = UINT_MAX;
     StartOfSequence = 0;
     EndOfSequence = 0;
   };
 
   for (unsigned i = 0, e = Changes.size(); i != e; ++i) {
     if (Changes[i].NewlinesBefore != 0) {
-      CurrentLine += Changes[i].NewlinesBefore;
-      if (StartOfSequence > 0 &&
-          (Changes[i].NewlinesBefore > 1 || !FoundAssignmentOnLine)) {
-        EndOfSequence = i;
-        AlignSequence();
-      }
-      FoundAssignmentOnLine = false;
-      FoundLeftParenOnLine = false;
+      CommasBeforeMatch = 0;
+      EndOfSequence = i;
+      // If there is a blank line, or if the last line didn't contain any
+      // matching token, the sequence ends here.
+      if (Changes[i].NewlinesBefore > 1 || !FoundMatchOnLine)
+        AlignCurrentSequence();
+
+      FoundMatchOnLine = false;
     }
 
-    if ((Changes[i].Kind == tok::equal &&
-         (FoundAssignmentOnLine || ((Changes[i].NewlinesBefore > 0 ||
-                                     Changes[i + 1].NewlinesBefore > 0)))) ||
-        (!FoundLeftParenOnLine && Changes[i].Kind == tok::r_paren)) {
-      if (StartOfSequence > 0)
-        AlignSequence();
-    } else if (Changes[i].Kind == tok::l_paren) {
-      FoundLeftParenOnLine = true;
-      if (!FoundAssignmentOnLine && StartOfSequence > 0)
-        AlignSequence();
-    } else if (!FoundAssignmentOnLine && !FoundLeftParenOnLine &&
-               Changes[i].Kind == tok::equal) {
-      FoundAssignmentOnLine = true;
-      EndOfSequence = i;
-      if (StartOfSequence == 0)
-        StartOfSequence = i;
+    if (Changes[i].Kind == tok::comma) {
+      ++CommasBeforeMatch;
+    } else if (Changes[i].Kind == tok::r_brace ||
+               Changes[i].Kind == tok::r_paren ||
+               Changes[i].Kind == tok::r_square) {
+      --NestingLevel;
+    } else if (Changes[i].Kind == tok::l_brace ||
+               Changes[i].Kind == tok::l_paren ||
+               Changes[i].Kind == tok::l_square) {
+      // We want sequences to skip over child scopes if possible, but not the
+      // other way around.
+      NestingLevelOfLastMatch = std::min(NestingLevelOfLastMatch, NestingLevel);
+      ++NestingLevel;
+    }
 
-      unsigned ChangeMinColumn = Changes[i].StartOfTokenColumn;
-      MinColumn = std::max(MinColumn, ChangeMinColumn);
+    if (!Matches(Changes[i]))
+      continue;
+
+    // If there is more than one matching token per line, or if the number of
+    // preceding commas, or the scope depth, do not match anymore, end the
+    // sequence.
+    if (FoundMatchOnLine || CommasBeforeMatch != CommasBeforeLastMatch ||
+        NestingLevel != NestingLevelOfLastMatch)
+      AlignCurrentSequence();
+
+    CommasBeforeLastMatch = CommasBeforeMatch;
+    NestingLevelOfLastMatch = NestingLevel;
+    FoundMatchOnLine = true;
+
+    if (StartOfSequence == 0)
+      StartOfSequence = i;
+
+    unsigned ChangeMinColumn = Changes[i].StartOfTokenColumn;
+    int LineLengthAfter = -Changes[i].Spaces;
+    for (unsigned j = i; j != e && Changes[j].NewlinesBefore == 0; ++j)
+      LineLengthAfter += Changes[j].Spaces + Changes[j].TokenLength;
+    unsigned ChangeMaxColumn = Style.ColumnLimit - LineLengthAfter;
+
+    // If we are restricted by the maximum column width, end the sequence.
+    if (ChangeMinColumn > MaxColumn || ChangeMaxColumn < MinColumn ||
+        CommasBeforeLastMatch != CommasBeforeMatch) {
+      AlignCurrentSequence();
+      StartOfSequence = i;
     }
-  }
 
-  if (StartOfSequence > 0) {
-    EndOfSequence = Changes.size();
-    AlignSequence();
+    MinColumn = std::max(MinColumn, ChangeMinColumn);
+    MaxColumn = std::min(MaxColumn, ChangeMaxColumn);
   }
+
+  EndOfSequence = Changes.size();
+  AlignCurrentSequence();
 }
 
-void WhitespaceManager::alignConsecutiveAssignments(unsigned Start,
-                                                    unsigned End,
-                                                    unsigned Column) {
-  bool AlignedAssignment = false;
-  int PreviousShift = 0;
-  for (unsigned i = Start; i != End; ++i) {
-    int Shift = 0;
-    if (Changes[i].NewlinesBefore > 0)
-      AlignedAssignment = false;
-    if (!AlignedAssignment && Changes[i].Kind == tok::equal) {
-      Shift = Column - Changes[i].StartOfTokenColumn;
-      AlignedAssignment = true;
-      PreviousShift = Shift;
-    }
-    assert(Shift >= 0);
-    Changes[i].Spaces += Shift;
-    if (i + 1 != Changes.size())
-      Changes[i + 1].PreviousEndOfTokenColumn += Shift;
-    Changes[i].StartOfTokenColumn += Shift;
-    if (AlignedAssignment) {
-      Changes[i].StartOfTokenColumn += PreviousShift;
-      if (i + 1 != Changes.size())
-        Changes[i + 1].PreviousEndOfTokenColumn += PreviousShift;
-    }
-  }
+void WhitespaceManager::alignConsecutiveAssignments() {
+  if (!Style.AlignConsecutiveAssignments)
+    return;
+
+  AlignTokens(Style,
+              [&](const Change &C) {
+                // Do not align on equal signs that are first on a line.
+                if (C.NewlinesBefore > 0)
+                  return false;
+
+                // Do not align on equal signs that are last on a line.
+                if (&C != &Changes.back() && (&C + 1)->NewlinesBefore > 0)
+                  return false;
+
+                return C.Kind == tok::equal;
+              },
+              Changes);
+}
+
+void WhitespaceManager::alignConsecutiveDeclarations() {
+  if (!Style.AlignConsecutiveDeclarations)
+    return;
+
+  // FIXME: Currently we don't handle properly the PointerAlignment: Right
+  // The * and & are not aligned and are left dangling. Something has to be done
+  // about it, but it raises the question of alignment of code like:
+  //   const char* const* v1;
+  //   float const* v2;
+  //   SomeVeryLongType const& v3;
+
+  AlignTokens(Style, [](Change const &C) { return C.IsStartOfDeclName; },
+              Changes);
 }
 
 void WhitespaceManager::alignTrailingComments() {
@@ -377,7 +472,7 @@ void WhitespaceManager::generateChanges() {
   }
 }
 
-void WhitespaceManager::storeReplacement(const SourceRange &Range,
+void WhitespaceManager::storeReplacement(SourceRange Range,
                                          StringRef Text) {
   unsigned WhitespaceLength = SourceMgr.getFileOffset(Range.getEnd()) -
                               SourceMgr.getFileOffset(Range.getBegin());