3 files changed, 131 insertions, 19 deletions

diff --git a/clang/lib/Format/Format.cpp b/clang/lib/Format/Format.cpp
index 3df1b60d..1d20824 100644
--- a/clang/lib/Format/Format.cpp
+++ b/clang/lib/Format/Format.cpp
@@ -112,6 +112,17 @@ template <> struct MappingTraits<FormatStyle::AlignConsecutiveStyle> {
 };
 
 template <>
+struct MappingTraits<FormatStyle::ShortCaseStatementsAlignmentStyle> {
+  static void mapping(IO &IO,
+                      FormatStyle::ShortCaseStatementsAlignmentStyle &Value) {
+    IO.mapOptional("Enabled", Value.Enabled);
+    IO.mapOptional("AcrossEmptyLines", Value.AcrossEmptyLines);
+    IO.mapOptional("AcrossComments", Value.AcrossComments);
+    IO.mapOptional("AlignCaseColons", Value.AlignCaseColons);
+  }
+};
+
+template <>
 struct ScalarEnumerationTraits<FormatStyle::AttributeBreakingStyle> {
   static void enumeration(IO &IO, FormatStyle::AttributeBreakingStyle &Value) {
     IO.enumCase(Value, "Always", FormatStyle::ABS_Always);
@@ -857,6 +868,8 @@ template <> struct MappingTraits<FormatStyle> {
     IO.mapOptional("AlignConsecutiveDeclarations",
                    Style.AlignConsecutiveDeclarations);
     IO.mapOptional("AlignConsecutiveMacros", Style.AlignConsecutiveMacros);
+    IO.mapOptional("AlignConsecutiveShortCaseStatements",
+                   Style.AlignConsecutiveShortCaseStatements);
     IO.mapOptional("AlignEscapedNewlines", Style.AlignEscapedNewlines);
     IO.mapOptional("AlignOperands", Style.AlignOperands);
     IO.mapOptional("AlignTrailingComments", Style.AlignTrailingComments);
@@ -1333,6 +1346,7 @@ FormatStyle getLLVMStyle(FormatStyle::LanguageKind Language) {
   LLVMStyle.AlignConsecutiveBitFields = {};
   LLVMStyle.AlignConsecutiveDeclarations = {};
   LLVMStyle.AlignConsecutiveMacros = {};
+  LLVMStyle.AlignConsecutiveShortCaseStatements = {};
   LLVMStyle.AlignTrailingComments = {};
   LLVMStyle.AlignTrailingComments.Kind = FormatStyle::TCAS_Always;
   LLVMStyle.AlignTrailingComments.OverEmptyLines = 0;
diff --git a/clang/lib/Format/WhitespaceManager.cpp b/clang/lib/Format/WhitespaceManager.cpp
index 03dcc31..668ca38 100644
--- a/clang/lib/Format/WhitespaceManager.cpp
+++ b/clang/lib/Format/WhitespaceManager.cpp
@@ -102,6 +102,7 @@ const tooling::Replacements &WhitespaceManager::generateReplacements() {
   llvm::sort(Changes, Change::IsBeforeInFile(SourceMgr));
   calculateLineBreakInformation();
   alignConsecutiveMacros();
+  alignConsecutiveShortCaseStatements();
   alignConsecutiveDeclarations();
   alignConsecutiveBitFields();
   alignConsecutiveAssignments();
@@ -675,14 +676,12 @@ static unsigned AlignTokens(const FormatStyle &Style, F &&Matches,
 //
 // 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.
-static void AlignMacroSequence(
+static void AlignMatchingTokenSequence(
     unsigned &StartOfSequence, unsigned &EndOfSequence, unsigned &MinColumn,
-    unsigned &MaxColumn, bool &FoundMatchOnLine,
-    std::function<bool(const WhitespaceManager::Change &C)> AlignMacrosMatches,
+    std::function<bool(const WhitespaceManager::Change &C)> Matches,
     SmallVector<WhitespaceManager::Change, 16> &Changes) {
   if (StartOfSequence > 0 && StartOfSequence < EndOfSequence) {
-
-    FoundMatchOnLine = false;
+    bool FoundMatchOnLine = false;
     int Shift = 0;
 
     for (unsigned I = StartOfSequence; I != EndOfSequence; ++I) {
@@ -693,8 +692,8 @@ static void AlignMacroSequence(
 
       // 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 && AlignMacrosMatches(Changes[I])) {
+      // shifted by the same amount.
+      if (!FoundMatchOnLine && Matches(Changes[I])) {
         FoundMatchOnLine = true;
         Shift = MinColumn - Changes[I].StartOfTokenColumn;
         Changes[I].Spaces += Shift;
@@ -708,7 +707,6 @@ static void AlignMacroSequence(
   }
 
   MinColumn = 0;
-  MaxColumn = UINT_MAX;
   StartOfSequence = 0;
   EndOfSequence = 0;
 }
@@ -747,7 +745,6 @@ void WhitespaceManager::alignConsecutiveMacros() {
   };
 
   unsigned MinColumn = 0;
-  unsigned MaxColumn = UINT_MAX;
 
   // Start and end of the token sequence we're processing.
   unsigned StartOfSequence = 0;
@@ -775,8 +772,8 @@ void WhitespaceManager::alignConsecutiveMacros() {
           !(LineIsComment && Style.AlignConsecutiveMacros.AcrossComments);
 
       if (EmptyLineBreak || NoMatchBreak) {
-        AlignMacroSequence(StartOfSequence, EndOfSequence, MinColumn, MaxColumn,
-                           FoundMatchOnLine, AlignMacrosMatches, Changes);
+        AlignMatchingTokenSequence(StartOfSequence, EndOfSequence, MinColumn,
+                                   AlignMacrosMatches, Changes);
       }
 
       // A new line starts, re-initialize line status tracking bools.
@@ -796,18 +793,12 @@ void WhitespaceManager::alignConsecutiveMacros() {
       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;
-
     MinColumn = std::max(MinColumn, ChangeMinColumn);
-    MaxColumn = std::min(MaxColumn, ChangeMaxColumn);
   }
 
   EndOfSequence = I;
-  AlignMacroSequence(StartOfSequence, EndOfSequence, MinColumn, MaxColumn,
-                     FoundMatchOnLine, AlignMacrosMatches, Changes);
+  AlignMatchingTokenSequence(StartOfSequence, EndOfSequence, MinColumn,
+                             AlignMacrosMatches, Changes);
 }
 
 void WhitespaceManager::alignConsecutiveAssignments() {
@@ -863,6 +854,110 @@ void WhitespaceManager::alignConsecutiveBitFields() {
       Changes, /*StartAt=*/0, Style.AlignConsecutiveBitFields);
 }
 
+void WhitespaceManager::alignConsecutiveShortCaseStatements() {
+  if (!Style.AlignConsecutiveShortCaseStatements.Enabled ||
+      !Style.AllowShortCaseLabelsOnASingleLine) {
+    return;
+  }
+
+  auto Matches = [&](const Change &C) {
+    if (Style.AlignConsecutiveShortCaseStatements.AlignCaseColons)
+      return C.Tok->is(TT_CaseLabelColon);
+
+    // Ignore 'IsInsideToken' to allow matching trailing comments which
+    // need to be reflowed as that causes the token to appear in two
+    // different changes, which will cause incorrect alignment as we'll
+    // reflow early due to detecting multiple aligning tokens per line.
+    return !C.IsInsideToken && C.Tok->Previous &&
+           C.Tok->Previous->is(TT_CaseLabelColon);
+  };
+
+  unsigned MinColumn = 0;
+
+  // Empty case statements don't break the alignment, but don't necessarily
+  // match our predicate, so we need to track their column so they can push out
+  // our alignment.
+  unsigned MinEmptyCaseColumn = 0;
+
+  // Start and end of the token sequence we're processing.
+  unsigned StartOfSequence = 0;
+  unsigned EndOfSequence = 0;
+
+  // Whether a matching token has been found on the current line.
+  bool FoundMatchOnLine = false;
+
+  bool LineIsComment = true;
+  bool LineIsEmptyCase = false;
+
+  unsigned I = 0;
+  for (unsigned E = Changes.size(); I != E; ++I) {
+    if (Changes[I].NewlinesBefore != 0) {
+      // Whether to break the alignment sequence because of an empty line.
+      bool EmptyLineBreak =
+          (Changes[I].NewlinesBefore > 1) &&
+          !Style.AlignConsecutiveShortCaseStatements.AcrossEmptyLines;
+
+      // Whether to break the alignment sequence because of a line without a
+      // match.
+      bool NoMatchBreak =
+          !FoundMatchOnLine &&
+          !(LineIsComment &&
+            Style.AlignConsecutiveShortCaseStatements.AcrossComments) &&
+          !LineIsEmptyCase;
+
+      if (EmptyLineBreak || NoMatchBreak) {
+        AlignMatchingTokenSequence(StartOfSequence, EndOfSequence, MinColumn,
+                                   Matches, Changes);
+        MinEmptyCaseColumn = 0;
+      }
+
+      // A new line starts, re-initialize line status tracking bools.
+      FoundMatchOnLine = false;
+      LineIsComment = true;
+      LineIsEmptyCase = false;
+    }
+
+    if (Changes[I].Tok->isNot(tok::comment))
+      LineIsComment = false;
+
+    if (Changes[I].Tok->is(TT_CaseLabelColon)) {
+      LineIsEmptyCase =
+          !Changes[I].Tok->Next || Changes[I].Tok->Next->isTrailingComment();
+
+      if (LineIsEmptyCase) {
+        if (Style.AlignConsecutiveShortCaseStatements.AlignCaseColons) {
+          MinEmptyCaseColumn =
+              std::max(MinEmptyCaseColumn, Changes[I].StartOfTokenColumn);
+        } else {
+          MinEmptyCaseColumn =
+              std::max(MinEmptyCaseColumn, Changes[I].StartOfTokenColumn + 2);
+        }
+      }
+    }
+
+    if (!Matches(Changes[I]))
+      continue;
+
+    if (LineIsEmptyCase)
+      continue;
+
+    FoundMatchOnLine = true;
+
+    if (StartOfSequence == 0)
+      StartOfSequence = I;
+
+    EndOfSequence = I + 1;
+
+    MinColumn = std::max(MinColumn, Changes[I].StartOfTokenColumn);
+
+    // Allow empty case statements to push out our alignment.
+    MinColumn = std::max(MinColumn, MinEmptyCaseColumn);
+  }
+
+  AlignMatchingTokenSequence(StartOfSequence, EndOfSequence, MinColumn, Matches,
+                             Changes);
+}
+
 void WhitespaceManager::alignConsecutiveDeclarations() {
   if (!Style.AlignConsecutiveDeclarations.Enabled)
     return;
diff --git a/clang/lib/Format/WhitespaceManager.h b/clang/lib/Format/WhitespaceManager.h
index cadfc1e..df7e9ad 100644
--- a/clang/lib/Format/WhitespaceManager.h
+++ b/clang/lib/Format/WhitespaceManager.h
@@ -232,6 +232,9 @@ private:
   /// Align consecutive declarations over all \c Changes.
   void alignChainedConditionals();
 
+  /// Align consecutive short case statements over all \c Changes.
+  void alignConsecutiveShortCaseStatements();
+
   /// Align trailing comments over all \c Changes.
   void alignTrailingComments();