RedundantBraces: handle infix expressions

scalameta · Apr 13, 2020 · 3d84236 · 3d84236
1 parent 94b1f30
commit 3d84236
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 10 deletions.
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/RedundantBraces.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/RedundantBraces.scala
@@ -10,6 +10,7 @@ import scala.meta._
 import scala.meta.tokens.Token.LeftBrace
 import scala.meta.tokens.Token.RightBrace
 
+import org.scalafmt.util.InfixApp
 import org.scalafmt.util.TreeOps._
 
 object RedundantBraces extends Rewrite {
@@ -169,19 +170,36 @@ class RedundantBraces(implicit ctx: RewriteCtx) extends RewriteSession {
       val close = b.tokens.last
       if (close.is[RightBrace] && check(b)) {
         implicit val builder = Seq.newBuilder[TokenPatch]
-        val ok =
-          if (b.parent.exists(_.is[Term.ApplyInfix])) {
+        val ok = b.parent match {
+          case Some(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.
              * we shouldn't join with the previous line (which might also end
              * in a comment), and if we keep the break before the right brace
              * we are removing, that will likely invalidate the expression. */
-            !ctx.onlyWhitespaceBefore(close)
-          } else {
+            val canRemove: ((Token, Option[Token.LF])) => Boolean =
+              if (!ctx.style.newlines.formatInfix) {
+                case (_, lfOpt) => lfOpt.isEmpty
+              }
+              else {
+                case (nonWs, lfOpt) =>
+                  if (nonWs.is[Token.Comment]) lfOpt.isEmpty
+                  else {
+                    lfOpt.foreach(builder += TokenPatch.Remove(_))
+                    true
+                  }
+              }
+            Seq(
+              // ensure can remove opening brace
+              ctx.tokenTraverser.findAfter(open) _,
+              // ensure can remove closing brace
+              ctx.tokenTraverser.findBefore(close) _
+            ).forall(ctx.findNonWhitespaceWith(_).exists(canRemove))
+          case _ =>
             ctx.removeLFToAvoidEmptyLine(close)
             true
-          }
+        }
         if (ok) {
           builder += TokenPatch.Remove(open)
           builder += TokenPatch.Remove(close)

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/Rewrite.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/Rewrite.scala
@@ -52,6 +52,20 @@ case class RewriteCtx(
   def onlyWhitespaceBefore(token: Token): Boolean =
     tokenTraverser.findBefore(token)(RewriteCtx.isLFSkipWhitespace).isDefined
 
+  def findNonWhitespaceWith(
+      f: (Token => Option[Boolean]) => Option[Token]
+  ): Option[(Token, Option[LF])] = {
+    var lf: Option[LF] = None
+    val nonWs = f {
+      case t: LF =>
+        if (lf.nonEmpty) Some(false)
+        else { lf = Some(t); None }
+      case Whitespace() => None
+      case _ => Some(true)
+    }
+    nonWs.map((_, lf))
+  }
+
   def getPatchToAvoidEmptyLine(token: Token): Option[TokenPatch] =
     if (!onlyWhitespaceBefore(token)) None
     else

diff --git a/scalafmt-tests/src/test/resources/rewrite/RedundantBraces.stat b/scalafmt-tests/src/test/resources/rewrite/RedundantBraces.stat
@@ -1,6 +1,7 @@
 rewrite.rules = [RedundantBraces]
 rewrite.redundantBraces.maxLines = 1
 rewrite.redundantBraces.generalExpressions = true
+newlines.afterInfix = some
 <<< basic
 object a {
   def x(i: Int): Int = {
@@ -321,6 +322,8 @@ object a {
   b.c(d => e => { f; g })
 }
 <<< #1027 2.1: parens to braces no
+newlines.afterInfix = keep
+===
 object a {
   b.c (d =>
     e
@@ -607,9 +610,7 @@ object a {
 }
 >>>
 object a {
-  1 + {
-    2
-  } + 3
+  1 + 2 + 3
 }
 <<< #1633 4.2: infix with literal and no break before
 object a {
@@ -618,8 +619,7 @@ object a {
 }
 >>>
 object a {
-  1 +
-    2 + 3
+  1 + 2 + 3
 }
 <<< #1633 4.3: infix with literal and comment
 object a {