FormatOps: extract helpers to OptimizationEntities

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

@@ -57,63 +57,10 @@ class FormatOps(
 
   val (forceConfigStyle, emptyQueueSpots) = getForceConfigStyle
 
+  val optimizationEntities = OptimizationEntities(topSourceTree)
+
   @inline
   def owners(token: T): Tree = ownersMap(hash(token))
-  /*
-   * The tokens on the left hand side of Pkg
-   *
-   * For example Set(org, ., scalafmt) in:
-   *
-   * package org.scalafmt
-   *
-   * import foo.bar
-   * ...
-   *
-   */
-  val (argumentStarts, optionalNewlines) = {
-    val arguments = mutable.Map.empty[Int, Tree]
-    val optional = Set.newBuilder[Int]
-    def getHeadIndex(tree: Tree): Option[Int] = getHeadOpt(tree)
-      .map(_.meta.idx - 1)
-    def addWith(key: Tree)(value: Tree): Unit = getHeadIndex(key)
-      .foreach(arguments.getOrElseUpdate(_, value))
-    def add(tree: Tree): Unit = addWith(tree)(tree)
-    def addOptional(tree: Tree): Unit = getHeadIndex(tree).foreach(optional += _)
-    def addParam(t: Term.Param, key: Tree): Unit = {
-      addWith(key)(t)
-      t.mods.foreach(addOptional)
-      addOptional(t.name)
-    }
-
-    val queue = new mutable.ListBuffer[Seq[Tree]]
-    queue += topSourceTree :: Nil
-    while (queue.nonEmpty) queue.remove(0).foreach { tree =>
-      tree match {
-        case _: Lit.Unit =>
-        case t: Term.ParamClause =>
-          val params = t.mod match {
-            case Some(mod) =>
-              addOptional(mod)
-              t.values match {
-                case head :: rest =>
-                  addParam(head, mod)
-                  rest
-                case _ => Nil
-              }
-            case _ => t.values
-          }
-          params.foreach(x => addParam(x, x))
-        case t: Term.ArgClause => add(t)
-        case t: Member.SyntaxValuesClause => t.values.foreach(add)
-        case t: Member.Tuple => t.args.foreach(add)
-        case _: Term.Param => // covered by Term.ParamClause
-        case t: Term => add(t)
-        case _ =>
-      }
-      queue += tree.children
-    }
-    (arguments.toMap, optional.result())
-  }
 
   @inline
   final def findFirst(start: FormatToken, end: T)(

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

@@ -0,0 +1,73 @@
+package org.scalafmt.internal
+
+import scala.meta._
+
+import scala.collection.mutable
+
+class OptimizationEntities(
+    argumentStarts: Map[Int, Tree],
+    optionalNewlines: Set[Int],
+) {
+  def argumentAt(idx: Int): Option[Tree] = argumentStarts.get(idx)
+  def argument(implicit ft: FormatToken): Option[Tree] = argumentAt(ft.meta.idx)
+  def optionalNL(implicit ft: FormatToken): Boolean =
+    optionalNewlines(ft.meta.idx)
+}
+
+object OptimizationEntities {
+  def apply(tree: Tree)(implicit ftoks: FormatTokens): OptimizationEntities =
+    new Builder(tree).build()
+
+  private class Builder(topSourceTree: Tree)(implicit tokens: FormatTokens) {
+
+    private val arguments = mutable.Map.empty[Int, Tree]
+    private val optional = Set.newBuilder[Int]
+
+    def build(): OptimizationEntities = {
+      val queue = new mutable.ListBuffer[Seq[Tree]]
+      queue += topSourceTree :: Nil
+      while (queue.nonEmpty) queue.remove(0).foreach { tree =>
+        processForArguments(tree)
+        queue += tree.children
+      }
+      new OptimizationEntities(arguments.toMap, optional.result())
+    }
+
+    private def getHeadIndex(tree: Tree): Option[Int] = tokens.getHeadOpt(tree)
+      .map(_.meta.idx - 1)
+    private def addArgWith(key: Tree)(value: Tree): Unit = getHeadIndex(key)
+      .foreach(arguments.getOrElseUpdate(_, value))
+    private def addArg(tree: Tree): Unit = addArgWith(tree)(tree)
+    private def addOptional(tree: Tree): Unit = getHeadIndex(tree)
+      .foreach(optional += _)
+    private def addParam(t: Term.Param, key: Tree): Unit = {
+      addArgWith(key)(t)
+      t.mods.foreach(addOptional)
+      addOptional(t.name)
+    }
+
+    private def processForArguments(tree: Tree): Unit = tree match {
+      case _: Lit.Unit =>
+      case t: Term.ParamClause =>
+        val params = t.mod match {
+          case Some(mod) =>
+            addOptional(mod)
+            t.values match {
+              case head :: rest =>
+                addParam(head, mod)
+                rest
+              case _ => Nil
+            }
+          case _ => t.values
+        }
+        params.foreach(x => addParam(x, x))
+      case t: Term.ArgClause => addArg(t)
+      case t: Member.SyntaxValuesClause => t.values.foreach(addArg)
+      case t: Member.Tuple => t.args.foreach(addArg)
+      case _: Term.Param => // covered by Term.ParamClause
+      case t: Term => addArg(t)
+      case _ =>
+    }
+  }
+
+}

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

@@ -1123,7 +1123,7 @@ class Router(formatOps: FormatOps) {
             .map(p => PenalizeAllNewlines(close, p + 3))
 
           val binpack = style.binPack.defnSiteFor(isBracket)
-          val firstArg = argumentStarts.get(ft.meta.idx)
+          val firstArg = optimizationEntities.argument
           val nextComma = firstArg.flatMap { x =>
             val ok = isSeqMulti(getArgs(leftOwner))
             if (ok) findFirstOnRight[T.Comma](getLast(x), close) else None
@@ -1494,7 +1494,7 @@ class Router(formatOps: FormatOps) {
       case FormatToken(_: T.Comma, right, _) if !leftOwner.is[Template] =>
         def forBinPack(binPack: BinPack.Site, callSite: Boolean) =
           if (binPack eq BinPack.Site.Never) None
-          else argumentStarts.get(ft.meta.idx).map { nextArg =>
+          else optimizationEntities.argument.map { nextArg =>
             val lastFT = getLast(nextArg)
             val lastTok = lastFT.left
             val oneline = binPack.isOneline
@@ -2373,7 +2373,7 @@ class Router(formatOps: FormatOps) {
           )
         Seq(spaceSplit, Split(Newline, if (spaceSplit.isActive) 1 else 0))
 
-      case FormatToken(_, r, _) if optionalNewlines(ft.meta.idx) =>
+      case FormatToken(_, r, _) if optimizationEntities.optionalNL =>
         @tailrec
         def noAnnoLeftFor(tree: Tree): Boolean = tree.parent match {
           case Some(_: Mod.Annot) => false