Implement optimized traverseFilter in terms of Chain as well

typelevel · Dec 13, 2023 · 2d5f4d7 · 2d5f4d7
1 parent c5d90a6
commit 2d5f4d7
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 8 deletions.
diff --git a/core/src/main/scala-2.13+/cats/instances/arraySeq.scala b/core/src/main/scala-2.13+/cats/instances/arraySeq.scala
@@ -234,7 +234,7 @@ private[cats] object ArraySeqInstances {
         G match {
           case x: StackSafeMonad[G] =>
             x.map(TraverseFilter.traverseFilterDirectly(fa.iterator)(f)(x))(
-              _.to(ArraySeq.untagged)
+              _.iterator.to(ArraySeq.untagged)
             )
           case _ =>
             fa.foldRight(Eval.now(G.pure(ArraySeq.untagged.empty[B]))) { case (x, xse) =>

diff --git a/core/src/main/scala/cats/TraverseFilter.scala b/core/src/main/scala/cats/TraverseFilter.scala
@@ -21,7 +21,7 @@
 
 package cats
 
-import cats.data.State
+import cats.data.{Chain, State}
 import cats.kernel.compat.scalaVersionSpecific._
 
 import scala.collection.immutable.{IntMap, TreeSet}
@@ -206,8 +206,8 @@ object TraverseFilter {
 
   private[cats] def traverseFilterDirectly[G[_], A, B](
     fa: IterableOnce[A]
-  )(f: A => G[Option[B]])(implicit G: StackSafeMonad[G]): G[Vector[B]] = {
-    fa.iterator.foldLeft(G.pure(Vector.empty[B])) { case (bldrG, a) =>
+  )(f: A => G[Option[B]])(implicit G: StackSafeMonad[G]): G[Chain[B]] = {
+    fa.iterator.foldLeft(G.pure(Chain.empty[B])) { case (bldrG, a) =>
       G.map2(bldrG, f(a)) {
         case (acc, Some(b)) => acc :+ b
         case (acc, None)    => acc

diff --git a/core/src/main/scala/cats/data/Chain.scala b/core/src/main/scala/cats/data/Chain.scala
@@ -1372,7 +1372,7 @@ sealed abstract private[data] class ChainInstances extends ChainInstances1 {
       else
         G match {
           case x: StackSafeMonad[G] =>
-            G.map(TraverseFilter.traverseFilterDirectly(fa.iterator)(f)(x))(Chain.fromSeq(_))
+            TraverseFilter.traverseFilterDirectly(fa.iterator)(f)(x)
           case _ =>
             traverseFilterViaChain {
               val as = collection.mutable.ArrayBuffer[A]()

diff --git a/core/src/main/scala/cats/instances/queue.scala b/core/src/main/scala/cats/instances/queue.scala
@@ -241,7 +241,7 @@ private object QueueInstances {
       else
         G match {
           case x: StackSafeMonad[G] =>
-            x.map(TraverseFilter.traverseFilterDirectly(fa)(f)(x))(traverse.fromIterableOnce(_))
+            x.map(TraverseFilter.traverseFilterDirectly(fa)(f)(x))(c => traverse.fromIterableOnce(c.iterator))
           case _ =>
             G.map(Chain.traverseFilterViaChain {
               val as = collection.mutable.ArrayBuffer[A]()

diff --git a/core/src/main/scala/cats/instances/seq.scala b/core/src/main/scala/cats/instances/seq.scala
@@ -210,7 +210,7 @@ trait SeqInstances extends cats.kernel.instances.SeqInstances {
 
     def traverseFilter[G[_], A, B](fa: Seq[A])(f: (A) => G[Option[B]])(implicit G: Applicative[G]): G[Seq[B]] =
       G match {
-        case x: StackSafeMonad[G] => x.map(TraverseFilter.traverseFilterDirectly(fa)(f)(x))(_.toSeq)
+        case x: StackSafeMonad[G] => x.map(TraverseFilter.traverseFilterDirectly(fa)(f)(x))(_.toVector)
         case _ =>
           G.map(Chain.traverseFilterViaChain(fa.toIndexedSeq)(f))(_.toVector)
       }

diff --git a/core/src/main/scala/cats/instances/vector.scala b/core/src/main/scala/cats/instances/vector.scala
@@ -271,7 +271,7 @@ private[instances] trait VectorInstancesBinCompat0 {
 
     def traverseFilter[G[_], A, B](fa: Vector[A])(f: (A) => G[Option[B]])(implicit G: Applicative[G]): G[Vector[B]] =
       G match {
-        case x: StackSafeMonad[G] => TraverseFilter.traverseFilterDirectly(fa)(f)(x)
+        case x: StackSafeMonad[G] => x.map(TraverseFilter.traverseFilterDirectly(fa)(f)(x))(_.toVector)
         case _ =>
           G.map(Chain.traverseFilterViaChain(fa)(f))(_.toVector)
       }