Router: check formatInfix using the infix tree

To upgrade to scalameta v4.5.1 (which deprecates Type.{And,Or} trees in
favour of regular Type.ApplyInfix trees), we will need to add special
handling of Type.ApplyInfix trees.

For that, checking newlines.afterInfix will need to be done selectively,
based on the tree, hence we need to pass the tree parameter around.

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