[clang] [clang-format] Allow array alignment on non-rectangular arrays (PR #143781)

[Björn Schäpers via cfe-commits]

================
@@ -1318,281 +1343,249 @@ void WhitespaceManager::alignArrayInitializers(unsigned Start, unsigned End) {
 
 void WhitespaceManager::alignArrayInitializersRightJustified(
     CellDescriptions &&CellDescs) {
-  if (!CellDescs.isRectangular())
+
+  // If there are less than two rows, there is nothing to align.
+  if (CellDescs.Rows.size() < 2)
+    return;
+
+  // If there are less than 2 columns, there is nothing to align.
+  const int ColumnCount = CellDescs.ColumnStartingCellIndices.size();
+  if (ColumnCount < 2)
     return;
 
   const int BracePadding = Style.Cpp11BracedListStyle ? 0 : 1;
+  auto &ColumnStartingIndices = CellDescs.ColumnStartingCellIndices;
   auto &Cells = CellDescs.Cells;
-  // Now go through and fixup the spaces.
-  auto *CellIter = Cells.begin();
-  for (auto i = 0U; i < CellDescs.CellCounts[0]; ++i, ++CellIter) {
-    unsigned NetWidth = 0U;
-    if (isSplitCell(*CellIter))
-      NetWidth = getNetWidth(Cells.begin(), CellIter, CellDescs.InitialSpaces);
-    auto CellWidth = getMaximumCellWidth(CellIter, NetWidth);
-
-    if (Changes[CellIter->Index].Tok->is(tok::r_brace)) {
-      // So in here we want to see if there is a brace that falls
-      // on a line that was split. If so on that line we make sure that
-      // the spaces in front of the brace are enough.
-      const auto *Next = CellIter;
-      do {
-        const FormatToken *Previous = Changes[Next->Index].Tok->Previous;
-        if (Previous && Previous->isNot(TT_LineComment)) {
-          Changes[Next->Index].Spaces = BracePadding;
-          Changes[Next->Index].NewlinesBefore = 0;
-        }
-        Next = Next->NextColumnElement;
-      } while (Next);
-      // Unless the array is empty, we need the position of all the
-      // immediately adjacent cells
-      if (CellIter != Cells.begin()) {
-        auto ThisNetWidth =
-            getNetWidth(Cells.begin(), CellIter, CellDescs.InitialSpaces);
-        auto MaxNetWidth = getMaximumNetWidth(
-            Cells.begin(), CellIter, CellDescs.InitialSpaces,
-            CellDescs.CellCounts[0], CellDescs.CellCounts.size());
-        if (ThisNetWidth < MaxNetWidth)
-          Changes[CellIter->Index].Spaces = (MaxNetWidth - ThisNetWidth);
-        auto RowCount = 1U;
-        auto Offset = std::distance(Cells.begin(), CellIter);
-        for (const auto *Next = CellIter->NextColumnElement; Next;
-             Next = Next->NextColumnElement) {
-          if (RowCount >= CellDescs.CellCounts.size())
-            break;
-          auto *Start = (Cells.begin() + RowCount * CellDescs.CellCounts[0]);
-          auto *End = Start + Offset;
-          ThisNetWidth = getNetWidth(Start, End, CellDescs.InitialSpaces);
-          if (ThisNetWidth < MaxNetWidth)
-            Changes[Next->Index].Spaces = (MaxNetWidth - ThisNetWidth);
-          ++RowCount;
-        }
-      }
-    } else {
-      auto ThisWidth =
-          calculateCellWidth(CellIter->Index, CellIter->EndIndex, true) +
-          NetWidth;
-      if (Changes[CellIter->Index].NewlinesBefore == 0) {
-        Changes[CellIter->Index].Spaces = (CellWidth - (ThisWidth + NetWidth));
-        Changes[CellIter->Index].Spaces += (i > 0) ? 1 : BracePadding;
-      }
-      alignToStartOfCell(CellIter->Index, CellIter->EndIndex);
-      for (const auto *Next = CellIter->NextColumnElement; Next;
-           Next = Next->NextColumnElement) {
-        ThisWidth =
-            calculateCellWidth(Next->Index, Next->EndIndex, true) + NetWidth;
-        if (Changes[Next->Index].NewlinesBefore == 0) {
-          Changes[Next->Index].Spaces = (CellWidth - ThisWidth);
-          Changes[Next->Index].Spaces += (i > 0) ? 1 : BracePadding;
-        }
-        alignToStartOfCell(Next->Index, Next->EndIndex);
-      }
+
+  // Calculate column widths.
+  SmallVector<unsigned> ColumnWidths;       // Widths from the previous column
+  SmallVector<unsigned> SummedColumnWidths; // Widths from the start of the row
+
+  unsigned CurrentWidth = 0;
+  for (unsigned CellIndex : ColumnStartingIndices) {
+    const CellDescription *current = &Cells[CellIndex];
+
+    unsigned MaxWidth = 0;
+    while (current != nullptr) {
+      unsigned CellWidth = calculateCellWidth(*current);
+
+      // +1 for the space after the comma in the previous column in all but
+      // the first column which has brace padding from the opening
+      // brace instead.
+      CellWidth += (current->Cell > 0) ? 1 : BracePadding;
+
+      MaxWidth = std::max(MaxWidth, CellWidth);
+      current = current->NextColumnElement;
     }
+
+    ColumnWidths.push_back(MaxWidth);
+
+    CurrentWidth += MaxWidth;
+    SummedColumnWidths.push_back(CurrentWidth);
+
+    // +1 for the comma between cells.
+    CurrentWidth++;
+  }
+
+  // Fixup spaces.
+  for (RowDescription &Row : CellDescs.Rows) {
+    unsigned WidthSoFarInRow = 0;
+    for (unsigned i = Row.StartCellIndex; i < Row.EndCellIndex; i++) {
+      const CellDescription &Cell = Cells[i];
+
+      const unsigned CellWidth = calculateCellWidth(Cell);
+      const unsigned AlignmentSpaces = ColumnWidths[Cell.Cell] - CellWidth;
+      setChangeSpaces(Cell.Index, AlignmentSpaces);
+
+      WidthSoFarInRow = SummedColumnWidths[Cell.Cell];
+
+      // +1 for the comma after columns in all but the last column
+      // Note: this can't check Cell.Cell because a row may not have a full
+      // set of columns.
+      if (i < Row.EndCellIndex - 1)
+        WidthSoFarInRow++;
+    }
+
+    // Align the end brace.
+    const unsigned AlignmentSpaces =
+        (SummedColumnWidths.back() - WidthSoFarInRow) + BracePadding;
+    setChangeSpaces(Row.ClosingBraceChangeIndex, AlignmentSpaces);
   }
 }
 
 void WhitespaceManager::alignArrayInitializersLeftJustified(
----------------
HazardyKnusperkeks wrote:

Can it be combined with the `RightJustified` version? Maybe a template on a `bool` so one can `constexpr if` the differences?

https://github.com/llvm/llvm-project/pull/143781