Add iterateWhileM and iterateUntilM

typelevel · Aug 30, 2017 · 3f7dd4d · 3f7dd4d
1 parent 9a31560
commit 3f7dd4d
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 0 deletions.
diff --git a/core/src/main/scala/cats/Monad.scala b/core/src/main/scala/cats/Monad.scala
@@ -100,4 +100,23 @@ import simulacrum.typeclass
     }
   }
 
+  /**
+    * Apply a monadic function iteratively until its result fails
+    * to satisfy the given predicate and return that result.
+    */
+  def iterateWhileM[A](init: A)(f: A => F[A])(p: A => Boolean): F[A] =
+    tailRecM(init) { a =>
+      if (p(a))
+        map(f(a))(Left(_))
+      else
+        pure(Right(a))
+    }
+
+  /**
+    * Apply a monadic function iteratively until its result satisfies
+    * the given predicate and return that result.
+    */
+  def iterateUntilM[A](init: A)(f: A => F[A])(p: A => Boolean): F[A] =
+    iterateWhileM(init)(f)(!p(_))
+
 }
diff --git a/core/src/main/scala/cats/syntax/monad.scala b/core/src/main/scala/cats/syntax/monad.scala
@@ -3,6 +3,9 @@ package syntax
 
 trait MonadSyntax {
   implicit final def catsSyntaxMonad[F[_], A](fa: F[A]): MonadOps[F, A] = new MonadOps(fa)
+
+  implicit final def catsSyntaxMonadIdOps[A](a: A): MonadIdOps[A] =
+    new MonadIdOps[A](a)
 }
 
 final class MonadOps[F[_], A](val fa: F[A]) extends AnyVal {
@@ -13,3 +16,16 @@ final class MonadOps[F[_], A](val fa: F[A]) extends AnyVal {
   def iterateWhile(p: A => Boolean)(implicit M: Monad[F]): F[A] = M.iterateWhile(fa)(p)
   def iterateUntil(p: A => Boolean)(implicit M: Monad[F]): F[A] = M.iterateUntil(fa)(p)
 }
+
+final class MonadIdOps[A](val a: A) extends AnyVal {
+
+  /**
+    * Iterative application of `f` while `p` holds.
+    */
+  def iterateWhileM[F[_]](f: A => F[A])(p: A => Boolean)(implicit M: Monad[F]): F[A] = M.iterateWhileM(a)(f)(p)
+
+  /**
+    * Iterative application of `f` until `p` holds.
+    */
+  def iterateUntilM[F[_]](f: A => F[A])(p: A => Boolean)(implicit M: Monad[F]): F[A] = M.iterateUntilM(a)(f)(p)
+}
diff --git a/tests/src/test/scala/cats/tests/MonadTest.scala b/tests/src/test/scala/cats/tests/MonadTest.scala
@@ -76,4 +76,28 @@ class MonadTest extends CatsSuite {
     result should ===(50000)
   }
 
+  test("iterateWhileM") {
+    forAll(smallPosInt) { (max: Int) =>
+      val (n, sum) = 0.iterateWhileM(s => incrementAndGet map (_ + s))(_ < max).run(0)
+      sum should ===(n * (n + 1) / 2)
+    }
+  }
+
+  test("iterateWhileM is stack safe") {
+    val (n, sum) = 0.iterateWhileM(s => incrementAndGet map (_ + s))(_ < 50000000).run(0)
+    sum should ===(n * (n + 1) / 2)
+  }
+
+  test("iterateUntilM") {
+    forAll(smallPosInt) { (max: Int) =>
+      val (n, sum) = 0.iterateUntilM(s => incrementAndGet map (_ + s))(_ > max).run(0)
+      sum should ===(n * (n + 1) / 2)
+    }
+  }
+
+  test("iterateUntilM is stack safe") {
+    val (n, sum) = 0.iterateUntilM(s => incrementAndGet map (_ + s))(_ > 50000000).run(0)
+    sum should ===(n * (n + 1) / 2)
+  }
+
 }