From a17f606888fe1dccf3c8c47629cace83f1954a08 Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Sat, 6 Apr 2024 14:34:29 -0700
Subject: [PATCH] FormatWriter: add style, get overflow in AlignLine

---
 .../org/scalafmt/internal/FormatWriter.scala  | 64 +++++++++++--------
 1 file changed, 36 insertions(+), 28 deletions(-)

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatWriter.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatWriter.scala
index 043e028abe..1d889d4fd1 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatWriter.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatWriter.scala
@@ -1185,8 +1185,11 @@ class FormatWriter(formatOps: FormatOps) {
             val block = getOrCreateBlock(alignContainer)
             val blockWasEmpty = block.isEmpty
             if (!blockWasEmpty || !isBlankLine) {
-              val alignLine =
-                new AlignLine(candidates, floc.state.prev.column + columnShift)
+              val alignLine = new AlignLine(
+                candidates,
+                floc.state.prev.column + columnShift,
+                floc.style,
+              )
               val appendToEmptyBlock = blockWasEmpty || {
                 val matches = columnMatches(
                   wasSameContainer(alignContainer),
@@ -1326,11 +1329,13 @@ class FormatWriter(formatOps: FormatOps) {
     )(implicit builder: mutable.Builder[(Int, Int), Map[Int, Int]]): Unit = {
       val endIndex = locations.length - 1
       block.foreach { x =>
-        val headStop = x.stops.head
-        if (headStop.floc.style.align.multiline) {
-          val offset = block.stopColumns.head - headStop.column
+        if (x.style.align.multiline) {
+          val headStop = x.stops.head
           val ftIndex = headStop.floc.formatToken.meta.idx
-          if (ftIndex < endIndex) shiftStateColumnIndent(ftIndex + 1, offset)
+          if (ftIndex < endIndex) shiftStateColumnIndent(
+            ftIndex + 1,
+            block.stopColumns.head - headStop.column,
+          )
         }
         var previousShift = 0
         x.stops.zip(block.stopColumns).foreach { case (stop, blockStop) =>
@@ -1653,7 +1658,14 @@ object FormatWriter {
 
   class AlignStop(val column: Int, val floc: FormatLocation, val hashKey: Int)
 
-  class AlignLine(var stops: IndexedSeq[AlignStop], val eolColumn: Int)
+  class AlignLine(
+      var stops: IndexedSeq[AlignStop],
+      val eolColumn: Int,
+      val style: ScalafmtConfig,
+  ) {
+    @inline
+    def noOverflow(shift: Int) = eolColumn + shift <= style.maxColumn
+  }
 
   class AlignBlock(
       buffer: mutable.ArrayBuffer[AlignLine] =
@@ -1697,29 +1709,25 @@ object FormatWriter {
       oldStops.drop(newStops.length).foreach(newStopCols += oldShift + _)
       newStops.drop(oldStops.length).foreach(newStopCols += newShift + _.column)
 
-      // check overflow
-      val style = line.stops.last.floc.style
-      if (!style.align.allowOverflow) {
-        val overflow = (line.eolColumn + newShift) > style.maxColumn ||
-          buffer.exists { x =>
-            val idx = -1 + {
-              if (truncate < 0) math.min(matches, x.stops.length)
-              else x.stops.length
+      def finalize() =
+        (line.style.align.allowOverflow || // check overflow
+          line.noOverflow(newShift) && buffer.forall { bl =>
+            bl.style.align.allowOverflow || {
+              val len = bl.stops.length
+              val idx = -1 + (if (truncate < 0) math.min(matches, len) else len)
+              bl.noOverflow(newStopCols(idx) - bl.stops(idx).column)
             }
-            val lastStop = x.stops(idx)
-            val totalShift = newStopCols(idx) - lastStop.column
-            x.eolColumn + totalShift > lastStop.floc.style.maxColumn
-          }
-        if (overflow) return false // RETURNING
-      }
+          }) && {
+          // now we mutate
+          line.stops = newStops
+          if (truncate < 0) foreach(x => x.stops = x.stops.take(matches))
+          if (newStopCols.length == newStops.length) refStops = newStops
+          buffer += line
+          stopColumns = newStopCols.toIndexedSeq
+          true
+        }
 
-      // now we mutate
-      line.stops = newStops
-      if (truncate < 0) foreach(x => x.stops = x.stops.take(matches))
-      if (newStopCols.length == newStops.length) refStops = newStops
-      buffer += line
-      stopColumns = newStopCols.toIndexedSeq
-      true
+      finalize()
     }
 
     // <0 old, 0 neither, >0 new