zio · iravid · Nov 4, 2019 · Nov 4, 2019
diff --git a/core-tests/shared/src/test/scala/zio/ZIOSpec.scala b/core-tests/shared/src/test/scala/zio/ZIOSpec.scala
@@ -1434,6 +1434,40 @@ object ZIOSpec
               """.stripMargin
             }
           }
+        ),
+        suite("doWhile")(
+          testM("doWhile repeats while condition is true") {
+            for {
+              in     <- Ref.make(10)
+              out    <- Ref.make(0)
+              _      <- (in.update(_ - 1) <* out.update(_ + 1)).doWhile(_ >= 0)
+              result <- out.get
+            } yield assert(result, equalTo(11))
+          },
+          testM("doWhile always evaluates effect once") {
+            for {
+              ref    <- Ref.make(0)
+              _      <- ref.update(_ + 1).doWhile(_ => false)
+              result <- ref.get
+            } yield assert(result, equalTo(1))
+          }
+        ),
+        suite("doUntil")(
+          testM("doUntil repeats until condition is true") {
+            for {
+              in     <- Ref.make(10)
+              out    <- Ref.make(0)
+              _      <- (in.update(_ - 1) <* out.update(_ + 1)).doUntil(_ == 0)
+              result <- out.get
+            } yield assert(result, equalTo(10))
+          },
+          testM("doUntil always evaluates effect once") {
+            for {
+              ref    <- Ref.make(0)
+              _      <- ref.update(_ + 1).doUntil(_ => true)
+              result <- ref.get
+            } yield assert(result, equalTo(1))
+          }
         )
       )
     )

diff --git a/core/shared/src/main/scala/zio/ZIO.scala b/core/shared/src/main/scala/zio/ZIO.scala
@@ -286,6 +286,18 @@ sealed trait ZIO[-R, +E, +A] extends Serializable { self =>
   final def delay(duration: Duration): ZIO[R with Clock, E, A] =
     clock.sleep(duration) *> self
 
+  /**
+   * Repeats this effect until its result satisfies the specified predicate.
+   */
+  final def doUntil(f: A => Boolean): ZIO[R, E, A] =
+    repeat(Schedule.doUntil(f))
+
+  /**
+   * Repeats this effect while its result satisfies the specified predicate.
+   */
+  final def doWhile(f: A => Boolean): ZIO[R, E, A] =
+    repeat(Schedule.doWhile(f))
+
   /**
    * Returns an effect whose failure and success have been lifted into an
    * `Either`.The resulting effect cannot fail, because the failure case has

diff --git a/test/shared/src/main/scala/zio/test/Gen.scala b/test/shared/src/main/scala/zio/test/Gen.scala
@@ -18,7 +18,7 @@ package zio.test
 
 import scala.collection.immutable.SortedMap
 
-import zio.{ Schedule, UIO, ZIO }
+import zio.{ UIO, ZIO }
 import zio.random._
 import zio.stream.{ Stream, ZStream }
 
@@ -290,7 +290,7 @@ object Gen extends GenZIO with FunctionVariants {
       val difference = max - min + 1
       val effect =
         if (difference > 0) nextLong(difference).map(min + _)
-        else nextLong.repeat(Schedule.doUntil(n => min <= n && n <= max))
+        else nextLong.doUntil(n => min <= n && n <= max)
       effect.map(Sample.shrinkIntegral(min))
     }