RedundantBraces: refactor replacing ) with }

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

@@ -120,35 +120,38 @@ class RedundantBraces(implicit val ftoks: FormatTokens)
   private def onLeftParen(implicit ft: FT, style: ScalafmtConfig): Replacement = {
     val rt = ft.right
     val rtOwner = ft.meta.rightOwner
-    def lpFunction = okToReplaceFunctionInSingleArgApply(rtOwner).map {
+
+    def replaceWithBrace(owner: Option[Tree] = None, claim: List[Int] = Nil) = {
+      val lb = new T.LeftBrace(rt.input, rt.dialect, rt.start)
+      Some(replaceToken("{", owner = owner, claim = claim)(lb))
+    }
+    def replaceFrom(lbm: FT.Meta) =
+      replaceWithBrace(owner = Some(lbm.rightOwner), claim = lbm.idx :: Nil)
+
+    def lpFunction = okToReplaceFunctionInSingleArgApply(rtOwner).flatMap {
       case (`rt`, f) => f.body match {
           case b: Term.Block => ftoks.getHead(b) match {
               case FT(_: T.LeftBrace, _, lbm) =>
-                val lb = new T.LeftBrace(rt.input, rt.dialect, rt.start)
-                replaceToken("{", claim = lbm.idx - 1 :: Nil)(lb)
-              case _ => null
+                replaceWithBrace(claim = lbm.idx - 1 :: Nil)
+              case _ => None
             }
-          case _ => null
+          case _ => None
         }
-      case _ => null
+      case _ => None
     }
     // single-arg apply of a partial function
     // a({ case b => c; d }) change to a { case b => c; d }
     def lpPartialFunction = rtOwner match {
-      case ta @ Term.ArgClause(arg :: Nil, _) if !ta.parent.is[Init] =>
-        getOpeningParen(ta).map { lp =>
-          if (lp.ne(rt) || getBlockNestedPartialFunction(arg).isEmpty) null
-          else ftoks.nextNonCommentAfter(ft) match {
-            case FT(lt, _: T.LeftBrace, lbm) =>
-              if (lt eq rt) removeToken(claim = lbm.idx :: Nil)
-              else {
-                val lbo = Some(lbm.rightOwner)
-                val lb = new T.LeftBrace(rt.input, rt.dialect, rt.start)
-                replaceToken("{", owner = lbo, claim = lbm.idx :: Nil)(lb)
-              }
-            case _ => null
-          }
+      case ta @ Term.ArgClause(arg :: Nil, _)
+          if !ta.parent.is[Init] && getOpeningParen(ta).contains(rt) =>
+        def replaceWithNextBrace = ftoks.nextNonCommentAfter(ft) match {
+          case FT(lt, _: T.LeftBrace, lbm) =>
+            if (lt ne rt) replaceFrom(lbm)
+            else Some(removeToken(claim = lbm.idx :: Nil))
+          case _ => replaceWithBrace(Some(arg))
         }
+        if (getBlockNestedPartialFunction(arg).isEmpty) None
+        else replaceWithNextBrace
       case _ => None
     }
 
@@ -181,18 +184,22 @@ class RedundantBraces(implicit val ftoks: FormatTokens)
       }.orNull
 
     case ReplacementType.Replace if left.ft.right.is[T.LeftBrace] =>
+      // this also removes any trailing commas, so can't be conditionally invoked
       val pftOpt = getRightBraceBeforeRightParen(shouldBeRemoved = true)
+      def replaceWithBrace(rb: T, rtype: ReplacementType, startOff: Int) = {
+        val rbo = Some(left.ft.rightOwner)
+        left -> replaceToken("}", owner = rbo, rtype = rtype) {
+          new T.RightBrace(rb.input, rb.dialect, rb.start + startOff)
+        }
+      }
       def replaceIfAfterRightBrace = pftOpt.map { pft =>
         val rb = pft.left
         // move right to the end of the function
         val rType = new ReplacementType.RemoveAndResurrect(ftoks.prev(pft))
-        val rbo = Some(left.ft.rightOwner)
-        left -> replaceToken("}", owner = rbo, rtype = rType) {
-          // create a shifted token so that any child tree wouldn't own it
-          new T.RightBrace(rb.input, rb.dialect, rb.start + 1)
-        }
+        // create a shifted token so that any child tree wouldn't own it
+        replaceWithBrace(rb, rType, startOff = 1)
       }
-      (ft.meta.rightOwner match {
+      (ft.rightOwner match {
         case ac: Term.ArgClause => session.rule[RemoveScala3OptionalBraces]
             .flatMap { r =>
               val repl = r.onLeftForArgClause(ac)(left.ft, left.style)