Merge pull request #973 from travisbrown/topic/statet-flatmapf

Add flatMapF for StateT
typelevel · Apr 11, 2016 · 26efefa · 26efefa
2 parents b3033e2 + 03a8664
commit 26efefa
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/core/src/main/scala/cats/data/StateT.scala b/core/src/main/scala/cats/data/StateT.scala
@@ -17,6 +17,15 @@ final class StateT[F[_], S, A](val runF: F[S => F[(S, A)]]) extends Serializable
         }
       })
 
+  def flatMapF[B](faf: A => F[B])(implicit F: Monad[F]): StateT[F, S, B] =
+    StateT(s =>
+      F.flatMap(runF) { fsf =>
+        F.flatMap(fsf(s)) { case (s, a) =>
+          F.map(faf(a))((s, _))
+        }
+      }
+    )
+
   def map[B](f: A => B)(implicit F: Monad[F]): StateT[F, S, B] =
     transform { case (s, a) => (s, f(a)) }
 

diff --git a/tests/src/test/scala/cats/tests/StateTTests.scala b/tests/src/test/scala/cats/tests/StateTTests.scala
@@ -40,6 +40,12 @@ class StateTTests extends CatsSuite {
     }
   }
 
+  test("flatMap and flatMapF consistent") {
+    forAll { (stateT: StateT[Option, Long, Int], f: Int => Option[Int]) =>
+      stateT.flatMap(a => StateT(s => f(a).map(b => (s, b)))) should === (stateT.flatMapF(f))
+    }
+  }
+
   test("runEmpty, runEmptyS, and runEmptyA consistent"){
     forAll { (f: StateT[List, Long, Int]) =>
       (f.runEmptyS zip f.runEmptyA) should === (f.runEmpty)