From c73bc12373c5679d05207e127f14413d998beeb6 Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Sat, 18 May 2024 11:10:08 -0700
Subject: [PATCH] Policy: add a Map policy, use instead of Proxy

---
 .../scala/org/scalafmt/internal/Policy.scala  | 28 +++++++++++++++++++
 .../scala/org/scalafmt/util/PolicyOps.scala   | 21 ++------------
 2 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Policy.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Policy.scala
index 5f69e7669..e94e9fbc0 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Policy.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Policy.scala
@@ -263,6 +263,34 @@ object Policy {
     override def toString: String = s"*($policy)$endPolicy"
   }
 
+  final class Map(
+      val endPolicy: Policy.End.WithPos,
+      val noDequeue: Boolean = false,
+      val rank: Int = 0,
+      desc: => String = "",
+  )(pred: Split => Split)(implicit fileLine: FileLine)
+      extends Policy.Clause {
+    private object PredicateDecision {
+      def unapply(d: Decision): Option[Seq[Split]] = {
+        var replaced = false
+        def applyMap(s: Split): Option[Split] = Option(pred(s)).filter { ss =>
+          (s eq ss) || {
+            replaced = true
+            !ss.isIgnored
+          }
+        }
+        val splits = d.splits.flatMap(applyMap)
+        if (replaced) Some(splits) else None
+      }
+    }
+    override val f: Policy.Pf = { case PredicateDecision(ss) => ss }
+    override def toString: String = {
+      val evalDesc = desc
+      val descStr = if (evalDesc.isEmpty) "" else s"[$evalDesc]"
+      s"MAP$descStr:" + super.toString
+    }
+  }
+
   sealed trait End extends (Token => End.WithPos) {
     def apply(endPos: Int): End.WithPos
     final def apply(token: Token): End.WithPos = apply(token.end)
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/util/PolicyOps.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/util/PolicyOps.scala
index ae044075e..711374183 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/util/PolicyOps.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/util/PolicyOps.scala
@@ -144,26 +144,11 @@ object PolicyOps {
     }
   }
 
-  private def delayedBreakPolicyFactory(onBreakPolicy: Policy): Policy.Pf = {
-    object OnBreakDecision {
-      def unapply(d: Decision): Option[Seq[Split]] = {
-        var replaced = false
-        def decisionPf(s: Split): Split =
-          if (!s.isNL) s
-          else {
-            replaced = true
-            s.orPolicy(onBreakPolicy)
-          }
-        val splits = d.splits.map(decisionPf)
-        if (replaced) Some(splits) else None
-      }
-    }
-    { case OnBreakDecision(d) => d }
-  }
-
   def delayedBreakPolicy(end: Policy.End.WithPos)(onBreakPolicy: Policy)(
       implicit fileLine: FileLine,
-  ): Policy = Policy.Proxy(onBreakPolicy, end)(delayedBreakPolicyFactory)
+  ): Policy = new Policy.Map(endPolicy = end, desc = onBreakPolicy.toString)({
+    s => if (s.isNL) s.orPolicy(onBreakPolicy) else s
+  })
 
   def delayedBreakPolicyBefore(token: T)(onBreakPolicy: Policy)(implicit
       fileLine: FileLine,