Optimize traverse

bench/src/main/scala/cats/bench/TraverseBench.scala

@@ -24,15 +24,19 @@ package cats.bench
 import cats.{Eval, Traverse, TraverseFilter}
 import org.openjdk.jmh.annotations.{Benchmark, Param, Scope, Setup, State}
 import org.openjdk.jmh.infra.Blackhole
+import cats.data.Chain
 
 @State(Scope.Benchmark)
 class TraverseBench {
   val listT: Traverse[List] = Traverse[List]
   val listTFilter: TraverseFilter[List] = TraverseFilter[List]
+  val chainTFilter: TraverseFilter[Chain] = TraverseFilter[Chain]
 
   val vectorT: Traverse[Vector] = Traverse[Vector]
   val vectorTFilter: TraverseFilter[Vector] = TraverseFilter[Vector]
 
+  val chainT: Traverse[Chain] = Traverse[Chain]
+
   // the unit of CPU work per iteration
   private[this] val Work: Long = 10
 
@@ -43,11 +47,13 @@ class TraverseBench {
 
   var list: List[Int] = _
   var vector: Vector[Int] = _
+  var chain: Chain[Int] = _
 
   @Setup
   def setup(): Unit = {
     list = 0.until(length).toList
     vector = 0.until(length).toVector
+    chain = Chain.fromSeq(0.until(length))
   }
 
   @Benchmark
@@ -83,6 +89,18 @@ class TraverseBench {
     }
   }
 
+  @Benchmark
+  def traverse_List(bh: Blackhole) = {
+    val result = listT.traverse_(list) { i =>
+      Eval.later {
+        Blackhole.consumeCPU(Work)
+        i * 2
+      }
+    }
+
+    bh.consume(result.value)
+  }
+
   @Benchmark
   def traverseFilterList(bh: Blackhole) = {
     val result = listTFilter.traverseFilter(list) { i =>
@@ -137,6 +155,18 @@ class TraverseBench {
     bh.consume(result.value)
   }
 
+  @Benchmark
+  def traverse_Vector(bh: Blackhole) = {
+    val result = vectorT.traverse_(vector) { i =>
+      Eval.later {
+        Blackhole.consumeCPU(Work)
+        i * 2
+      }
+    }
+
+    bh.consume(result.value)
+  }
+
   @Benchmark
   def traverseVectorError(bh: Blackhole) = {
     val result = vectorT.traverse(vector) { i =>
@@ -199,4 +229,61 @@ class TraverseBench {
 
     bh.consume(results)
   }
+
+  @Benchmark
+  def traverseChain(bh: Blackhole) = {
+    val result = chainT.traverse(chain) { i =>
+      Eval.later {
+        Blackhole.consumeCPU(Work)
+        i * 2
+      }
+    }
+
+    bh.consume(result.value)
+  }
+
+  @Benchmark
+  def traverse_Chain(bh: Blackhole) = {
+    val result = chainT.traverse_(chain) { i =>
+      Eval.later {
+        Blackhole.consumeCPU(Work)
+        i * 2
+      }
+    }
+
+    bh.consume(result.value)
+  }
+
+  @Benchmark
+  def traverseChainError(bh: Blackhole) = {
+    val result = chainT.traverse(chain) { i =>
+      Eval.later {
+        Blackhole.consumeCPU(Work)
+
+        if (i == length * 0.3) {
+          throw Failure
+        }
+
+        i * 2
+      }
+    }
+
+    try {
+      bh.consume(result.value)
+    } catch {
+      case Failure => ()
+    }
+  }
+
+  @Benchmark
+  def traverseFilterChain(bh: Blackhole) = {
+    val result = chainTFilter.traverseFilter(chain) { i =>
+      Eval.later {
+        Blackhole.consumeCPU(Work)
+        if (i % 2 == 0) Some(i * 2) else None
+      }
+    }
+
+    bh.consume(result.value)
+  }
 }

core/src/main/scala-2.13+/cats/instances/arraySeq.scala

@@ -102,7 +102,24 @@ private[cats] object ArraySeqInstances {
         B.combineAll(fa.iterator.map(f))
 
       def traverse[G[_], A, B](fa: ArraySeq[A])(f: A => G[B])(implicit G: Applicative[G]): G[ArraySeq[B]] =
-        G.map(Chain.traverseViaChain(fa)(f))(_.iterator.to(ArraySeq.untagged))
+        G match {
+          case x: StackSafeMonad[G] =>
+            x.map(Traverse.traverseDirectly(Vector.newBuilder[B])(fa.iterator)(f)(x))(_.to(ArraySeq.untagged))
+          case _ =>
+            G.map(Chain.traverseViaChain(fa)(f))(_.iterator.to(ArraySeq.untagged))
+
+        }
+
+      override def traverse_[G[_], A, B](fa: ArraySeq[A])(f: A => G[B])(implicit G: Applicative[G]): G[Unit] =
+        G match {
+          case x: StackSafeMonad[G] => Traverse.traverse_Directly(fa)(f)(x)
+          case _ =>
+            foldRight(fa, Always(G.pure(()))) { (a, acc) =>
+              G.map2Eval(f(a), acc) { (_, _) =>
+                ()
+              }
+            }.value
+        }
 
       override def mapAccumulate[S, A, B](init: S, fa: ArraySeq[A])(f: (S, A) => (S, B)): (S, ArraySeq[B]) =
         StaticMethods.mapAccumulateFromStrictFunctor(init, fa, f)(this)
@@ -214,9 +231,17 @@ private[cats] object ArraySeqInstances {
       def traverseFilter[G[_], A, B](
         fa: ArraySeq[A]
       )(f: (A) => G[Option[B]])(implicit G: Applicative[G]): G[ArraySeq[B]] =
-        fa.foldRight(Eval.now(G.pure(ArraySeq.untagged.empty[B]))) { case (x, xse) =>
-          G.map2Eval(f(x), xse)((i, o) => i.fold(o)(_ +: o))
-        }.value
+        G match {
+          case x: StackSafeMonad[G] =>
+            x.map(TraverseFilter.traverseFilterDirectly(Vector.newBuilder[B])(fa.iterator)(f)(x))(
+              _.to(ArraySeq.untagged)
+            )
+          case _ =>
+            fa.foldRight(Eval.now(G.pure(ArraySeq.untagged.empty[B]))) { case (x, xse) =>
+              G.map2Eval(f(x), xse)((i, o) => i.fold(o)(_ +: o))
+            }.value
+
+        }
 
       override def filterA[G[_], A](fa: ArraySeq[A])(f: (A) => G[Boolean])(implicit G: Applicative[G]): G[ArraySeq[A]] =
         fa.foldRight(Eval.now(G.pure(ArraySeq.untagged.empty[A]))) { case (x, xse) =>

core/src/main/scala/cats/Traverse.scala

@@ -23,6 +23,9 @@ package cats
 
 import cats.data.State
 import cats.data.StateT
+import cats.kernel.compat.scalaVersionSpecific._
+import cats.StackSafeMonad
+import scala.collection.mutable
 
 /**
  * Traverse, also known as Traversable.
@@ -284,4 +287,35 @@ object Traverse {
   @deprecated("Use cats.syntax object imports", "2.2.0")
   object nonInheritedOps extends ToTraverseOps
 
+  private[cats] def traverseDirectly[Coll[x] <: IterableOnce[x], G[_], A, B](
+    builder: mutable.Builder[B, Coll[B]]
+  )(fa: IterableOnce[A])(f: A => G[B])(implicit G: StackSafeMonad[G]): G[Coll[B]] = {
+    val size = fa.knownSize
+    if (size >= 0) {
+      builder.sizeHint(size)
+    }
+    G.map(fa.iterator.foldLeft(G.pure(builder)) { case (accG, a) =>
+      G.flatMap(accG) { acc =>
+        G.map(f(a)) { a =>
+          acc += a
+          acc
+        }
+      }
+    })(_.result())
+  }
+
+  private[cats] def traverse_Directly[G[_], A, B](
+    fa: IterableOnce[A]
+  )(f: A => G[B])(implicit G: StackSafeMonad[G]): G[Unit] = {
+    val iter = fa.iterator
+    if (iter.hasNext) {
+      val first = iter.next()
+      G.map(iter.foldLeft(f(first)) { case (g, a) =>
+        G.flatMap(g) { _ =>
+          f(a)
+        }
+      })(_ => ())
+    } else G.unit
+  }
+
 }

core/src/main/scala/cats/TraverseFilter.scala

@@ -22,8 +22,10 @@
 package cats
 
 import cats.data.State
+import cats.kernel.compat.scalaVersionSpecific._
 
 import scala.collection.immutable.{IntMap, TreeSet}
+import scala.collection.mutable
 
 /**
  * `TraverseFilter`, also known as `Witherable`, represents list-like structures
@@ -203,4 +205,21 @@ object TraverseFilter {
   @deprecated("Use cats.syntax object imports", "2.2.0")
   object nonInheritedOps extends ToTraverseFilterOps
 
+  private[cats] def traverseFilterDirectly[Coll[x] <: IterableOnce[x], G[_], A, B](
+    builder: mutable.Builder[B, Coll[B]]
+  )(fa: IterableOnce[A])(f: A => G[Option[B]])(implicit G: StackSafeMonad[G]): G[Coll[B]] = {
+    val size = fa.knownSize
+    if (size >= 0) {
+      builder.sizeHint(size)
+    }
+    G.map(fa.iterator.foldLeft(G.pure(builder)) { case (bldrG, a) =>
+      G.flatMap(bldrG) { bldr =>
+        G.map(f(a)) {
+          case Some(b) => bldr += b
+          case None    => bldr
+        }
+      }
+    })(_.result())
+  }
+
 }

core/src/main/scala/cats/data/Chain.scala

@@ -1241,11 +1241,27 @@ sealed abstract private[data] class ChainInstances extends ChainInstances1 {
       def traverse[G[_], A, B](fa: Chain[A])(f: A => G[B])(implicit G: Applicative[G]): G[Chain[B]] =
         if (fa.isEmpty) G.pure(Chain.nil)
         else
-          traverseViaChain {
-            val as = collection.mutable.ArrayBuffer[A]()
-            as ++= fa.iterator
-            KernelStaticMethods.wrapMutableIndexedSeq(as)
-          }(f)
+          G match {
+            case x: StackSafeMonad[G] =>
+              x.map(Traverse.traverseDirectly(List.newBuilder[B])(fa.iterator)(f)(x))(Chain.fromSeq(_))
+            case _ =>
+              traverseViaChain {
+                val as = collection.mutable.ArrayBuffer[A]()
+                as ++= fa.iterator
+                KernelStaticMethods.wrapMutableIndexedSeq(as)
+              }(f)
+          }
+
+      override def traverse_[G[_], A, B](fa: Chain[A])(f: A => G[B])(implicit G: Applicative[G]): G[Unit] =
+        G match {
+          case x: StackSafeMonad[G] => Traverse.traverse_Directly(fa.iterator)(f)(x)
+          case _ =>
+            foldRight(fa, Always(G.pure(()))) { (a, acc) =>
+              G.map2Eval(f(a), acc) { (_, _) =>
+                ()
+              }
+            }.value
+        }
 
       override def mapAccumulate[S, A, B](init: S, fa: Chain[A])(f: (S, A) => (S, B)): (S, Chain[B]) =
         StaticMethods.mapAccumulateFromStrictFunctor(init, fa, f)(this)
@@ -1354,11 +1370,16 @@ sealed abstract private[data] class ChainInstances extends ChainInstances1 {
     def traverseFilter[G[_], A, B](fa: Chain[A])(f: A => G[Option[B]])(implicit G: Applicative[G]): G[Chain[B]] =
       if (fa.isEmpty) G.pure(Chain.nil)
       else
-        traverseFilterViaChain {
-          val as = collection.mutable.ArrayBuffer[A]()
-          as ++= fa.iterator
-          KernelStaticMethods.wrapMutableIndexedSeq(as)
-        }(f)
+        G match {
+          case x: StackSafeMonad[G] =>
+            G.map(TraverseFilter.traverseFilterDirectly(List.newBuilder[B])(fa.iterator)(f)(x))(Chain.fromSeq(_))
+          case _ =>
+            traverseFilterViaChain {
+              val as = collection.mutable.ArrayBuffer[A]()
+              as ++= fa.iterator
+              KernelStaticMethods.wrapMutableIndexedSeq(as)
+            }(f)
+        }
 
     override def filterA[G[_], A](fa: Chain[A])(f: A => G[Boolean])(implicit G: Applicative[G]): G[Chain[A]] =
       traverse