Router: check formatInfix using the infix tree

scalafmt-core/shared/src/main/scala/org/scalafmt/config/Newlines.scala

@@ -237,11 +237,8 @@ case class Newlines(
   val formatInfix: Boolean = breakAfterInfix ne AfterInfix.keep
 
   def checkInfixConfig(infixCount: Int): Newlines = {
-    val needAfterInfix: AfterInfix =
-      if (infixCount > afterInfixMaxCountPerFile) AfterInfix.keep
-      else breakAfterInfix
-    if (needAfterInfix == breakAfterInfix) this
-    else copy(afterInfix = Some(needAfterInfix))
+    if (infixCount <= afterInfixMaxCountPerFile || !formatInfix) this
+    else copy(afterInfix = Some(AfterInfix.keep))
   }
 
   lazy val forceBeforeImplicitParamListModifier: Boolean =

scalafmt-core/shared/src/main/scala/org/scalafmt/config/ScalafmtConfig.scala

@@ -5,6 +5,7 @@ import java.nio.file
 import scala.collection.mutable
 import scala.io.Codec
 import scala.meta.Dialect
+import scala.meta.Tree
 import scala.util.Try
 
 import metaconfig._
@@ -254,6 +255,12 @@ case class ScalafmtConfig(
   lazy val forceNewlineBeforeDocstring: Boolean =
     docstrings.forceBlankLineBefore
       .getOrElse(optIn.forceBlankLineBeforeDocstring)
+
+  def breakAfterInfix(tree: => Tree): Newlines.AfterInfix =
+    newlines.breakAfterInfix
+
+  def formatInfix(tree: => Tree): Boolean =
+    breakAfterInfix(tree) ne Newlines.AfterInfix.keep
 }
 
 object ScalafmtConfig {

scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala

@@ -499,25 +499,42 @@ class FormatOps(
       case t: Type.ApplyInfix
           if style.spaces.neverAroundInfixTypes.contains(t.op.value) =>
         Seq(Split(NoSplit, 0))
-      case _ if style.newlines.formatInfix =>
-        if (ft.meta.leftOwner ne app.op) Seq(Split(Space, 0))
-        else InfixSplits(app, ft).getBeforeLhsOrRhs()
-      case _ =>
-        // we don't modify line breaks generally around infix expressions
-        // TODO: if that ever changes, modify how rewrite rules handle infix
-        Seq(InfixSplits(app, ft).withNLIndent(Split(getMod(ft), 0)))
+      case t =>
+        val afterInfix = style.breakAfterInfix(t)
+        if (afterInfix ne Newlines.AfterInfix.keep) {
+          if (ft.meta.leftOwner ne app.op) Seq(Split(Space, 0))
+          else InfixSplits(app, ft).getBeforeLhsOrRhs(afterInfix)
+        } else {
+          // we don't modify line breaks generally around infix expressions
+          // TODO: if that ever changes, modify how rewrite rules handle infix
+          Seq(InfixSplits(app, ft).withNLIndent(Split(getMod(ft), 0)))
+        }
     }
 
   def getInfixSplitsBeforeLhs(
       lhsApp: InfixApp,
       ft: FormatToken,
+      afterInfix: Newlines.AfterInfix,
       newStmtMod: Option[Modification] = None
   )(implicit style: ScalafmtConfig): Seq[Split] = {
     val fullInfixTreeOpt =
       findTreeWithParentSimple(lhsApp.all, false)(isInfixApp)
     val fullInfix = fullInfixTreeOpt.flatMap(asInfixApp).getOrElse(lhsApp)
     val app = findLeftInfix(fullInfix)
-    new InfixSplits(app, ft, fullInfix, app).getBeforeLhsOrRhs(newStmtMod)
+    new InfixSplits(app, ft, fullInfix, app)
+      .getBeforeLhsOrRhs(afterInfix, newStmtMod)
+  }
+
+  final def maybeGetInfixSplitsBeforeLhs(
+      ft: FormatToken,
+      mod: => Option[Modification] = None
+  )(nonInfixSplits: => Seq[Split])(implicit
+      style: ScalafmtConfig
+  ): Seq[Split] = {
+    val tree = ft.meta.rightOwner
+    val ai = style.breakAfterInfix(tree)
+    val app = if (ai eq Newlines.AfterInfix.keep) None else asInfixApp(tree)
+    app.fold(nonInfixSplits)(getInfixSplitsBeforeLhs(_, ft, ai, mod))
   }
 
   private[internal] object InfixSplits {
@@ -649,7 +666,8 @@ class FormatOps(
           Policy.on(fullExpire) {
             case Decision(t: FormatToken, s)
                 if isInfixOp(t.meta.leftOwner) ||
-                  !style.newlines.formatInfix && isInfixOp(t.meta.rightOwner) =>
+                  isInfixOp(t.meta.rightOwner) &&
+                  !t.meta.rightOwner.parent.exists(style.formatInfix(_)) =>
               InfixSplits.switch(s, triggers: _*)
           }
 
@@ -677,6 +695,7 @@ class FormatOps(
     }
 
     def getBeforeLhsOrRhs(
+        afterInfix: Newlines.AfterInfix,
         newStmtMod: Option[Modification] = None
     ): Seq[Split] = {
       val beforeLhs = ft.meta.leftOwner ne app.op
@@ -716,8 +735,7 @@ class FormatOps(
 
       val infixTooLong = infixSequenceLength(fullInfix) >
         style.newlines.afterInfixMaxCountPerExprForSome
-      val breakMany = infixTooLong ||
-        style.newlines.breakAfterInfix == Newlines.AfterInfix.many
+      val breakMany = infixTooLong || afterInfix == Newlines.AfterInfix.many
       val rightAsInfix = asInfixApp(ft.meta.rightOwner)
 
       def breakAfterComment(t: FormatToken) = {

scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala

@@ -603,7 +603,13 @@ class Router(formatOps: FormatOps) {
         )
           Seq(Split(getMod(formatToken), 0))
         else {
-          asInfixApp(rightOwner, style.newlines.formatInfix).fold {
+          maybeGetInfixSplitsBeforeLhs(
+            formatToken,
+            Some(
+              if (left.is[T.Comment] && tok.noBreak) Space
+              else NewlineT(isDouble = tok.hasBlankLine)
+            )
+          ) {
             val spaceCouldBeOk = annoLeft && (style.newlines.source match {
               case Newlines.unfold =>
                 right.is[T.Comment] ||
@@ -624,11 +630,6 @@ class Router(formatOps: FormatOps) {
               // For some reason, this newline cannot cost 1.
               Split(NewlineT(isDouble = tok.hasBlankLine), 0)
             )
-          } { app =>
-            val mod =
-              if (left.is[T.Comment] && tok.noBreak) Space
-              else NewlineT(isDouble = tok.hasBlankLine)
-            getInfixSplitsBeforeLhs(app, tok, Some(mod))
           }
         }
 
@@ -824,13 +825,13 @@ class Router(formatOps: FormatOps) {
       case FormatToken(_: T.KwDef, _: T.Ident, _) =>
         Seq(Split(Space, 0))
       case ft @ FormatToken(_: T.Equals, _, SplitAssignIntoPartsLeft(parts)) =>
-        asInfixApp(rightOwner, style.newlines.formatInfix).fold {
+        maybeGetInfixSplitsBeforeLhs(ft) {
           val (rhs, paramss) = parts
           getSplitsDefValEquals(ft, rhs) {
             if (paramss.isDefined) getSplitsDefEquals(ft, rhs)
             else getSplitsValEquals(ft, rhs)(getSplitsValEqualsClassic(ft, rhs))
           }
-        }(getInfixSplitsBeforeLhs(_, ft))
+        }
 
       // Parameter opening for one parameter group. This format works
       // on the WHOLE defnSite (via policies)
@@ -2750,7 +2751,7 @@ class Router(formatOps: FormatOps) {
       ft: FormatToken,
       body: Tree
   )(implicit style: ScalafmtConfig): Seq[Split] =
-    asInfixApp(ft.meta.rightOwner, style.newlines.formatInfix).fold {
+    maybeGetInfixSplitsBeforeLhs(ft) {
       val expire = getLastNonTrivialToken(body)
       val spaceIndents =
         if (!style.align.arrowEnumeratorGenerator) Seq.empty
@@ -2770,6 +2771,6 @@ class Router(formatOps: FormatOps) {
           }
         }(Split(Newline, _).withIndent(style.indent.main, expire, After))
       }
-    }(getInfixSplitsBeforeLhs(_, ft))
+    }
 
 }

scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/RedundantBraces.scala

@@ -267,7 +267,7 @@ class RedundantBraces(ftoks: FormatTokens) extends FormatTokensRewrite.Rule {
         b.tokens.lastOption.contains(rb) && b.tokens.head.is[Token.LeftBrace]
       case _ => false
     }) && okToRemoveBlock(b) && (b.parent match {
-      case Some(InfixApp(_)) =>
+      case Some(p @ InfixApp(_)) =>
         /* for infix, we will preserve the block unless the closing brace
          * follows a non-whitespace character on the same line as we don't
          * break lines around infix expressions.
@@ -277,12 +277,12 @@ class RedundantBraces(ftoks: FormatTokens) extends FormatTokensRewrite.Rule {
         def checkOpen = {
           val nft = ftoks.next(ft)
           nft.noBreak ||
-          style.newlines.formatInfix && !nft.right.is[Token.Comment]
+          style.formatInfix(p) && !nft.right.is[Token.Comment]
         }
         def checkClose = {
           val nft = ftoks(ftoks.matching(ft.right), -1)
           nft.noBreak ||
-          style.newlines.formatInfix && !nft.left.is[Token.Comment]
+          style.formatInfix(p) && !nft.left.is[Token.Comment]
         }
         checkOpen && checkClose
       case _ => true

scalafmt-core/shared/src/main/scala/org/scalafmt/util/TreeOps.scala

@@ -575,11 +575,8 @@ object TreeOps {
       case _ => false
     }
 
-  final def asInfixApp(tree: Tree): Option[InfixApp] = InfixApp.unapply(tree)
-
   @inline
-  final def asInfixApp(tree: Tree, flag: Boolean = true): Option[InfixApp] =
-    if (flag) asInfixApp(tree) else None
+  final def asInfixApp(tree: Tree): Option[InfixApp] = InfixApp.unapply(tree)
 
   @inline
   final def isInfixApp(tree: Tree): Boolean = asInfixApp(tree).isDefined