1 files changed, 99 insertions, 1 deletions

diff --git a/lib/Format/WhitespaceManager.cpp b/lib/Format/WhitespaceManager.cpp
index bf1207e..4baaab1 100644
--- a/lib/Format/WhitespaceManager.cpp
+++ b/lib/Format/WhitespaceManager.cpp
@@ -36,7 +36,9 @@ WhitespaceManager::Change::Change(
       PreviousLinePostfix(PreviousLinePostfix),
       CurrentLinePrefix(CurrentLinePrefix), Kind(Kind),
       ContinuesPPDirective(ContinuesPPDirective), IndentLevel(IndentLevel),
-      Spaces(Spaces) {}
+      Spaces(Spaces), IsTrailingComment(false), TokenLength(0),
+      PreviousEndOfTokenColumn(0), EscapedNewlineColumn(0),
+      StartOfBlockComment(nullptr), IndentationOffset(0) {}
 
 void WhitespaceManager::reset() {
   Changes.clear();
@@ -91,6 +93,7 @@ const tooling::Replacements &WhitespaceManager::generateReplacements() {
 
   std::sort(Changes.begin(), Changes.end(), Change::IsBeforeInFile(SourceMgr));
   calculateLineBreakInformation();
+  alignConsecutiveAssignments();
   alignTrailingComments();
   alignEscapedNewlines();
   generateChanges();
@@ -139,6 +142,96 @@ 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;
+
+  unsigned MinColumn = 0;
+  unsigned StartOfSequence = 0;
+  unsigned EndOfSequence = 0;
+  bool FoundAssignmentOnLine = false;
+  bool FoundLeftParenOnLine = false;
+  unsigned CurrentLine = 0;
+
+  auto AlignSequence = [&] {
+    alignConsecutiveAssignments(StartOfSequence, EndOfSequence, MinColumn);
+    MinColumn = 0;
+    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;
+    }
+
+    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;
+
+      unsigned ChangeMinColumn = Changes[i].StartOfTokenColumn;
+      MinColumn = std::max(MinColumn, ChangeMinColumn);
+    }
+  }
+
+  if (StartOfSequence > 0) {
+    EndOfSequence = Changes.size();
+    AlignSequence();
+  }
+}
+
+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::alignTrailingComments() {
   unsigned MinColumn = 0;
   unsigned MaxColumn = UINT_MAX;
@@ -264,6 +357,11 @@ void WhitespaceManager::alignEscapedNewlines(unsigned Start, unsigned End,
 void WhitespaceManager::generateChanges() {
   for (unsigned i = 0, e = Changes.size(); i != e; ++i) {
     const Change &C = Changes[i];
+    if (i > 0) {
+      assert(Changes[i - 1].OriginalWhitespaceRange.getBegin() !=
+                 C.OriginalWhitespaceRange.getBegin() &&
+             "Generating two replacements for the same location");
+    }
     if (C.CreateReplacement) {
       std::string ReplacementText = C.PreviousLinePostfix;
       if (C.ContinuesPPDirective)