Lazy foldM for "Iterables"

core/src/main/scala/cats/Foldable.scala

@@ -173,6 +173,12 @@ import simulacrum.typeclass
 
   /**
    * Left associative monadic folding on `F`.
+   *
+   * The default implementation of this is based on `foldLeft`, and thus will
+   * always fold across the entire structure. Certain structures are able to
+   * implement this in such a way that folds can be short-circuited (not
+   * traverse the entirety of the structure), depending on the `G` result
+   * produced at a given step.
    */
   def foldM[G[_], A, B](fa: F[A], z: B)(f: (B, A) => G[B])(implicit G: Monad[G]): G[B] =
     foldLeft(fa, G.pure(z))((gb, a) => G.flatMap(gb)(f(_, a)))
@@ -372,4 +378,39 @@ object Foldable {
       Eval.defer(if (it.hasNext) f(it.next, loop()) else lb)
     loop()
   }
+
+  /**
+   * Implementation of [[Foldable.foldM]] which can short-circuit for
+   * structures with an `Iterator`.
+   *
+   * For example we can sum a `Stream` of integers and stop if
+   * the sum reaches 100 (if we reach the end of the `Stream`
+   * before getting to 100 we return the total sum) :
+   *
+   * {{{
+   * scala> import cats.implicits._
+   * scala> type LongOr[A] = Either[Long, A]
+   * scala> def sumStream(s: Stream[Int]): Long =
+   *      |   Foldable.iteratorFoldM[LongOr, Int, Long](s.toIterator, 0L){ (acc, n) =>
+   *      |     val sum = acc + n
+   *      |     if (sum < 100L) Right(sum) else Left(sum)
+   *      |   }.merge
+   *
+   * scala> sumStream(Stream.continually(1))
+   * res0: Long = 100
+   *
+   * scala> sumStream(Stream(1,2,3,4))
+   * res1: Long = 10
+   * }}}
+   *
+   * Note that `Foldable[Stream].foldM` uses this method underneath, so
+   * you wouldn't call this method explicitly like in the example above.
+   */
+  def iteratorFoldM[M[_], A, B](it: Iterator[A], z: B)(f: (B, A) => M[B])(implicit M: Monad[M]): M[B] = {
+    val go: B => M[Either[B, B]] = { b =>
+      if (it.hasNext) M.map(f(b, it.next))(Left(_))
+      else M.pure(Right(b))
+    }
+    M.tailRecM(z)(go)
+  }
 }

core/src/main/scala/cats/instances/list.scala

@@ -79,6 +79,9 @@ trait ListInstances extends cats.kernel.instances.ListInstances {
       override def isEmpty[A](fa: List[A]): Boolean = fa.isEmpty
 
       override def filter[A](fa: List[A])(f: A => Boolean): List[A] = fa.filter(f)
+
+      override def foldM[G[_], A, B](fa: List[A], z: B)(f: (B, A) => G[B])(implicit G: Monad[G]): G[B] =
+        Foldable.iteratorFoldM(fa.toIterator, z)(f)
     }
 
   implicit def catsStdShowForList[A:Show]: Show[List[A]] =

core/src/main/scala/cats/instances/map.scala

@@ -78,6 +78,9 @@ trait MapInstances extends cats.kernel.instances.MapInstances {
       override def size[A](fa: Map[K, A]): Long = fa.size.toLong
 
       override def isEmpty[A](fa: Map[K, A]): Boolean = fa.isEmpty
+
+      override def foldM[G[_], A, B](fa: Map[K, A], z: B)(f: (B, A) => G[B])(implicit G: Monad[G]): G[B] =
+        Foldable.iteratorFoldM(fa.valuesIterator, z)(f)
     }
   // scalastyle:on method.length
 }

core/src/main/scala/cats/instances/set.scala

@@ -27,6 +27,9 @@ trait SetInstances extends cats.kernel.instances.SetInstances {
         fa.forall(p)
 
       override def isEmpty[A](fa: Set[A]): Boolean = fa.isEmpty
+
+      override def foldM[G[_], A, B](fa: Set[A], z: B)(f: (B, A) => G[B])(implicit G: Monad[G]): G[B] =
+        Foldable.iteratorFoldM(fa.toIterator, z)(f)
     }
 
   implicit def catsStdShowForSet[A:Show]: Show[Set[A]] = new Show[Set[A]] {

core/src/main/scala/cats/instances/stream.scala

@@ -106,6 +106,9 @@ trait StreamInstances extends cats.kernel.instances.StreamInstances {
       override def filter[A](fa: Stream[A])(f: A => Boolean): Stream[A] = fa.filter(f)
 
       override def collect[A, B](fa: Stream[A])(f: PartialFunction[A, B]): Stream[B] = fa.collect(f)
+
+      override def foldM[G[_], A, B](fa: Stream[A], z: B)(f: (B, A) => G[B])(implicit G: Monad[G]): G[B] =
+        Foldable.iteratorFoldM(fa.toIterator, z)(f)
     }
 
   implicit def catsStdShowForStream[A: Show]: Show[Stream[A]] =

core/src/main/scala/cats/instances/vector.scala

@@ -86,6 +86,9 @@ trait VectorInstances extends cats.kernel.instances.VectorInstances {
       override def filter[A](fa: Vector[A])(f: A => Boolean): Vector[A] = fa.filter(f)
 
       override def collect[A, B](fa: Vector[A])(f: PartialFunction[A, B]): Vector[B] = fa.collect(f)
+
+      override def foldM[G[_], A, B](fa: Vector[A], z: B)(f: (B, A) => G[B])(implicit G: Monad[G]): G[B] =
+        Foldable.iteratorFoldM(fa.toIterator, z)(f)
     }
 
   implicit def catsStdShowForVector[A:Show]: Show[Vector[A]] =

tests/src/test/scala/cats/tests/FoldableTests.scala

@@ -107,14 +107,35 @@ class FoldableTestsAdditional extends CatsSuite {
     larger.value should === (large.map(_ + 1))
   }
 
-  test("Foldable[List].foldM stack safety") {
-    def nonzero(acc: Long, x: Long): Option[Long] =
+  def checkFoldMStackSafety[F[_]](fromRange: Range => F[Int])(implicit F: Foldable[F]): Unit = {
+    def nonzero(acc: Long, x: Int): Option[Long] =
       if (x == 0) None else Some(acc + x)
 
-    val n = 100000L
-    val expected = n*(n+1)/2
-    val actual = Foldable[List].foldM((1L to n).toList, 0L)(nonzero)
-    assert(actual.get == expected)
+    val n = 100000
+    val expected = n.toLong*(n.toLong+1)/2
+    val foldMResult = F.foldM(fromRange(1 to n), 0L)(nonzero)
+    assert(foldMResult.get == expected)
+    ()
+  }
+
+  test("Foldable[List].foldM stack safety") {
+    checkFoldMStackSafety[List](_.toList)
+  }
+
+  test("Foldable[Stream].foldM stack safety") {
+    checkFoldMStackSafety[Stream](_.toStream)
+  }
+
+  test("Foldable[Vector].foldM stack safety") {
+    checkFoldMStackSafety[Vector](_.toVector)
+  }
+
+  test("Foldable[Set].foldM stack safety") {
+    checkFoldMStackSafety[Set](_.toSet)
+  }
+
+  test("Foldable[Map[String, ?]].foldM stack safety") {
+    checkFoldMStackSafety[Map[String, ?]](_.map(x => x.toString -> x).toMap)
   }
 
   test("Foldable[Stream]") {
@@ -141,6 +162,9 @@ class FoldableTestsAdditional extends CatsSuite {
     // test trampolining
     val large = Stream((1 to 10000): _*)
     assert(contains(large, 10000).value)
+
+    // test laziness of foldM
+    dangerous.foldM(0)((acc, a) => if (a < 2) Some(acc + a) else None) should === (None)
   }
 }