Add & test monadic FState-based AsyncStream generators. Refactor few methods and fix flatMap in FState

Author: danslapman

src/main/scala/asyncstreams/AsyncStreamMonad.scala

@@ -24,7 +24,7 @@ class AsyncStreamMonad(implicit executor: ExecutionContext) extends MonadPlus[As
 }
 
 trait AsyncStreamMonadFunctions {
-  def foreach[A, S](stream: AsyncStream[A])(f: A => FState[S, Any])
+  def foreach[A, S](stream: AsyncStream[A])(f: A => FState[S, _])
     (implicit ex: ExecutionContext): FState[S, Unit] =
     FState(s => {
       stream.foldLeft(Future(s))((futureS, a) => futureS.flatMap(s2 => f(a)(s2).map(_._2)))
@@ -35,17 +35,17 @@ trait AsyncStreamMonadFunctions {
     FState(s => stream.data.map(pair => (pair eq END, s)))
 
   def isEmpty[A, S : FStateMonad](f: S => AsyncStream[A])(implicit fsm: FStateMonad[S], ex: ExecutionContext): FState[S, Boolean] =
-    fsm.fconds((s: S) => isEmpty(f(s)))
+    fsm.fcondS((s: S) => isEmpty(f(s)))
 
   def notEmpty[A, S](stream: AsyncStream[A])(implicit ex: ExecutionContext): FState[S, Boolean] =
     FState(s => stream.data map (pair => (!(pair eq END), s)))
 
   def notEmpty[A, S](f: S => AsyncStream[A])(implicit fsm: FStateMonad[S], ex: ExecutionContext): FState[S, Boolean] =
-    fsm.fconds(s => notEmpty(f(s)))
+    fsm.fcondS(s => notEmpty(f(s)))
 
   def get[A, S](stream: AsyncStream[A])(implicit ex: ExecutionContext): FState[S, (A, AsyncStream[A])] =
     FState(s => stream.data.map(pair => ((pair.first, pair.second), s)))
 
-  def generate[S,A](start: S)(gen: FState[S, A])(implicit ex: ExecutionContext) =
+  def generateS[S,A](start: S)(gen: FState[S, A])(implicit ex: ExecutionContext) =
     AsyncStream.generate(start)(gen.func)
 }

src/main/scala/asyncstreams/FState.scala

@@ -2,18 +2,35 @@ package asyncstreams
 
 import scala.concurrent.{ExecutionContext, Future}
 
-class FState[S, A](val func: S => Future[(A, S)]) extends ((S) => Future[(A, S)]) {
+class FState[S, A](val func: S => Future[(A, S)]) {
+  import FState._
+
   def apply(s: S) = func(s)
 
   def flatMap[B](f: A => FState[S, B])(implicit ex: ExecutionContext): FState[S, B] = FState[S, B](
-    (s: S) => func(s) flatMap ((fst: A, snd: S) => f(fst)(snd)).tupled
+    (s: S) => func(s).flatMap {
+      case END => ENDF
+      case (fst, snd) => f(fst)(snd)
+    }
   )
 
   def map[B](f: A => B)(implicit ex: ExecutionContext): FState[S, B] =
     flatMap((a: A) => FState.unit(f(a)))
+
+  def bind[B](f: A => FState[S, B])(implicit ex: ExecutionContext): FState[S, B] =
+    FState((s: S) => func(s) flatMap {
+      case END => ENDF
+      case (fst, snd) => f(fst)(snd)
+    })
+
+  def filter(p: A => Boolean)(implicit executor: ExecutionContext): FState[S, A] =
+    bind(a => if (p(a)) unit(a) else empty[S, A])
+
+  def withFilter(p: A => Boolean)(implicit executor: ExecutionContext): FState[S, A] = filter(p)
 }
 
 object FState {
   def apply[S, A](f: S => Future[(A, S)]) = new FState[S, A](f)
   def unit[S, A](a: => A)(implicit ex: ExecutionContext) = FState[S, A]((s: S) => Future((a, s)))
+  def empty[S, A](implicit ex: ExecutionContext) = FState[S, A]((s: S) => ENDF)
 }

src/main/scala/asyncstreams/FStateMonad.scala

@@ -7,38 +7,30 @@ class FStateMonad[S](implicit ex: ExecutionContext)
   extends MonadPlus[({ type f[X] = FState[S, X]})#f] with FStateMonadFunctions {
   type F[X] = FState[S, X]
 
-  override def empty[A]: F[A] = FState((s: S) => ENDF)
+  override def empty[A]: F[A] = FState.empty[S, A]
 
   override def point[A](a: => A): F[A] = FState.unit(a)
 
-  override def bind[A, B](m: F[A])(f: A => F[B]): F[B] =
-    FState((s: S) => m(s) flatMap {
-      case END => ENDF
-      case (fst, snd) => f(fst)(snd)
-    })
+  override def bind[A, B](m: F[A])(f: A => F[B]): F[B] = m.bind(f)
 
   override def plus[A](a: F[A],b: => F[A]): F[A] = bind(a)(_ => b)
 
-  def conds(f: S => Boolean): F[Boolean] = bind(gets[S])(vs => point(f(vs)))
-  def fconds(f: S => F[Boolean]): F[Boolean] = bind(gets[S])(f)
-  def mods(f: S => S): F[S] = bind(gets[S])(vs => puts(f(vs)))
+  def condS(f: S => Boolean): F[Boolean] = bind(getS[S])(vs => point(f(vs)))
+  def fcondS(f: S => F[Boolean]): F[Boolean] = bind(getS[S])(f)
+  def modS(f: S => S): F[S] = bind(getS[S])(vs => putS(f(vs)))
 
   def forM_[A](cond: S => Boolean, mod: S => S)(action: => F[A]): F[Unit] =
-    whileM_(conds(cond), bind(action)(va => mods(mod)))
+    whileM_(condS(cond), bind(action)(va => modS(mod)))
 }
 
 trait FStateMonadFunctions {
-  def gets[S](implicit ex: ExecutionContext): FState[S, S] = FState((s: S) => Future((s, s)))
-  def puts[S](news: S)(implicit ex: ExecutionContext): FState[S, S] = FState((_: S) => Future((news, news)))
+  def getS[S](implicit ex: ExecutionContext): FState[S, S] = FState((s: S) => Future((s, s)))
 
-  def conds[S](f: S => Boolean)(implicit m: FStateMonad[S]): FStateMonad[S]#F[Boolean] =
-    m.conds(f)
+  def putS[S](news: S)(implicit ex: ExecutionContext): FState[S, S] = FState((_: S) => Future((news, news)))
 
-  /*
-  def fconds[S](f: S => FState[S, Boolean])(implicit m: FStateMonad[S]): FStateMonad[S]#F[Boolean] =
-    m.fconds(f)
-    */
+  def condS[S](f: S => Boolean)(implicit m: FStateMonad[S]): FStateMonad[S]#F[Boolean] = m.condS(f)
 
-  def mods[S : FStateMonad](f: S => S)(implicit m: FStateMonad[S]): FStateMonad[S]#F[S] =
-    m.mods(f)
+  def fconds[S](f: S => FState[S, Boolean])(implicit m: FStateMonad[S]): FStateMonad[S]#F[Boolean] = m.fconds(f)
+
+  def modS[S](f: S => S)(implicit m: FStateMonad[S]): FStateMonad[S]#F[S] = m.modS(f)
 }

src/test/scala/asyncstreams/AsyncStreamMonadicOperationsTests.scala

@@ -0,0 +1,72 @@
+package asyncstreams
+
+import monadops._
+import AsyncStream._
+
+import scala.concurrent.duration._
+import scala.concurrent.{Await, Future}
+import scala.concurrent.ExecutionContext.Implicits.global
+
+class AsyncStreamMonadicOperationsTests extends BaseSuite {
+  def makeStream(l: List[Int]) = generate(l)(l => if (l.isEmpty) ENDF else Future((l.head, l.tail)))
+
+  private def wait[T](f: Future[T]): T = Await.result(f, 10.seconds)
+
+  test("foreach") {
+    implicit val fsm = fStateInstance[Int]
+
+    val fstate = for {
+      _ <- foreach(makeStream(0 :: 1 :: 2 :: Nil)) {
+        v => modS[Int](_ + 1)
+      }
+      v2 <- getS[Int]
+    } yield v2
+
+    wait(fstate(0)) shouldBe (3, 3)
+  }
+
+  test("get, isEmpty") {
+    case class State(counter: Int, stream: AsyncStream[Int])
+    implicit val fsm = fStateInstance[State]
+
+    val stream = makeStream(0 :: 1 :: 2 :: 3 :: Nil)
+
+    val fstate = for {
+      _ <- fsm.whileM_(notEmpty(_.stream), for {
+        s <- getS[State]
+        (el, newStream) <- get[Int, State](s.stream)
+        _ <- putS[State](State(s.counter + el, newStream))
+      } yield ())
+      v <- getS[State]
+    } yield v.counter
+
+    wait(fstate(State(0, stream)))._1 shouldBe 6
+  }
+
+  test("FState as generator") {
+    implicit val fsm = fStateInstance[Int]
+
+    val stream = generateS(0) {
+      for {
+        s <- getS[Int]
+        _ <- putS[Int](s + 1)
+      } yield s
+    } take 3
+
+    wait(stream.to[List]) shouldBe (0 :: 1 :: 2 :: Nil)
+  }
+
+  test("Generate finite stream") {
+    implicit val fsm = fStateInstance[Int]
+
+    val stream = generateS(0) {
+      for {
+        s <- getS[Int]
+        if s < 3
+        _ <- putS(s + 1)
+      } yield s
+    }
+
+    wait(stream.to[List]) shouldBe (0 :: 1 :: 2 :: Nil)
+  }
+}

src/test/scala/asyncstreams/FStateTests.scala

@@ -24,22 +24,23 @@ class FStateTests extends BaseSuite {
     val fsm = fStateInstance[Int]
 
     val t = for {
-      _ <- fsm.whileM_(gets[Int] map (_ < 10), for {
-        i <- gets[Int]
-        _ <- puts(i + 1)
-      } yield (()))
-      v1 <- gets[Int]
+      _ <- fsm.whileM_(getS[Int] map (_ < 10), for {
+        i <- getS[Int]
+        _ <- putS(i + 1)
+      } yield ())
+      v1 <- getS[Int]
     } yield v1
 
     wait(t(0)) shouldBe (10, 10)
   }
 
   test("conds & mods") {
     implicit val fsm = fStateInstance[Int]
+    import fsm.whileM_
 
     val t = for {
-      _ <- fsm.whileM_(conds(_ < 10), mods(_ + 1))
-      v1 <- gets[Int]
+      _ <- whileM_(condS(_ < 10), modS(_ + 1))
+      v1 <- getS[Int]
     } yield v1
 
     wait(t(0)) shouldBe (10, 10)
@@ -52,7 +53,7 @@ class FStateTests extends BaseSuite {
       _ <- fsm.forM_(_ < 10, _ + 1) {
         fsm.point("AAAAAA")
       }
-      v1 <- gets[Int]
+      v1 <- getS[Int]
     } yield v1
 
     wait(t(0)) shouldBe (10, 10)