Merge pull request #2924 from ybasket/more-stack-safety-for-indexedre…

…aderwriterstatet More stack safety for IndexedReaderWriterStateT & IndexedStateT
typelevel · Jul 19, 2019 · ed99e17 · ed99e17
2 parents aaed15c + 3c671ad
commit ed99e17
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 3 deletions.
diff --git a/core/src/main/scala/cats/data/IndexedReaderWriterStateT.scala b/core/src/main/scala/cats/data/IndexedReaderWriterStateT.scala
@@ -90,7 +90,7 @@ final class IndexedReaderWriterStateT[F[_], E, L, SA, SB, A](val runF: F[(E, SA)
   def flatMap[SC, B](
     f: A => IndexedReaderWriterStateT[F, E, L, SB, SC, B]
   )(implicit F: FlatMap[F], L: Semigroup[L]): IndexedReaderWriterStateT[F, E, L, SA, SC, B] =
-    IndexedReaderWriterStateT.applyF {
+    IndexedReaderWriterStateT.shift {
       F.map(runF) { rwsfa => (e: E, sa: SA) =>
         F.flatMap(rwsfa(e, sa)) {
           case (la, sb, a) =>
@@ -108,7 +108,7 @@ final class IndexedReaderWriterStateT[F[_], E, L, SA, SB, A](val runF: F[(E, SA)
    * Like [[map]], but allows the mapping function to return an effectful value.
    */
   def flatMapF[B](faf: A => F[B])(implicit F: FlatMap[F]): IndexedReaderWriterStateT[F, E, L, SA, SB, B] =
-    IndexedReaderWriterStateT.applyF {
+    IndexedReaderWriterStateT.shift {
       F.map(runF) { rwsfa => (e: E, sa: SA) =>
         F.flatMap(rwsfa(e, sa)) {
           case (l, sb, a) =>
@@ -390,6 +390,25 @@ object IndexedReaderWriterStateT extends IRWSTInstances with CommonIRWSTConstruc
   def modifyF[F[_], E, L, SA, SB](f: SA => F[SB])(implicit F: Applicative[F],
                                                   L: Monoid[L]): IndexedReaderWriterStateT[F, E, L, SA, SB, Unit] =
     IndexedReaderWriterStateT((_, s) => F.map(f(s))((L.empty, _, ())))
+
+  /**
+   * Internal API — shifts the execution of `run` in the `F` context.
+   *
+   * Used to build IndexedReaderWriterStateT values for `F[_]` data types that implement `Monad`,
+   * in which case it is safer to trigger the `F[_]` context earlier.
+   *
+   * This is needed for [[IndexedReaderWriterStateT.flatMap]] to be stack-safe when the underlying F[_] is,
+   * for further explanation see [[Kleisli.shift]].
+   */
+  private[data] def shift[F[_], E, L, SA, SB, A](
+    runF: F[(E, SA) => F[(L, SB, A)]]
+  )(implicit F: FlatMap[F]): IndexedReaderWriterStateT[F, E, L, SA, SB, A] =
+    F match {
+      case ap: Applicative[F] @unchecked =>
+        IndexedReaderWriterStateT.apply[F, E, L, SA, SB, A]((e: E, sa: SA) => F.flatMap(runF)(f => f(e, sa)))(ap)
+      case _ =>
+        IndexedReaderWriterStateT.applyF(runF)
+    }
 }
 
 abstract private[data] class RWSTFunctions extends CommonIRWSTConstructors {

diff --git a/core/src/main/scala/cats/data/IndexedStateT.scala b/core/src/main/scala/cats/data/IndexedStateT.scala
@@ -101,7 +101,7 @@ final class IndexedStateT[F[_], SA, SB, A](val runF: F[SA => F[(SB, A)]]) extend
    */
   def transform[B, SC](f: (SB, A) => (SC, B))(implicit F: Functor[F]): IndexedStateT[F, SA, SC, B] =
     IndexedStateT.applyF(F.map(runF) { sfsa =>
-      sfsa.andThen { fsa =>
+      AndThen(sfsa).andThen { fsa =>
         F.map(fsa) { case (s, a) => f(s, a) }
       }
     })

diff --git a/tests/src/test/scala/cats/tests/IndexedReaderWriterStateTSuite.scala b/tests/src/test/scala/cats/tests/IndexedReaderWriterStateTSuite.scala
@@ -29,6 +29,22 @@ class ReaderWriterStateTSuite extends CatsSuite {
     rws.runS("context", 0).value should ===(70001)
   }
 
+  test("flatMap is stack-safe on repeated left binds when F is") {
+    val ns = (0 to 70000).toList
+    val one = addLogUnit(1)
+    val rws = ns.foldLeft(one)((acc, _) => acc.flatMap(_ => one))
+
+    rws.runS("context", 0).value should ===(70002)
+  }
+
+  test("flatMap is stack-safe on repeated right binds when F is") {
+    val ns = (0 to 70000).toList
+    val one = addLogUnit(1)
+    val rws = ns.foldLeft(one)((acc, _) => one.flatMap(_ => acc))
+
+    rws.runS("context", 0).value should ===(70002)
+  }
+
   test("map2 combines logs") {
     forAll {
       (rwsa: ReaderWriterState[String, Vector[Int], Int, Int],

diff --git a/tests/src/test/scala/cats/tests/IndexedStateTSuite.scala b/tests/src/test/scala/cats/tests/IndexedStateTSuite.scala
@@ -263,6 +263,15 @@ class IndexedStateTSuite extends CatsSuite {
   private val stackSafeTestSize =
     if (Platform.isJvm) 100000 else 100
 
+  test("repeated map is stack safe") {
+    val unit = StateT.pure[Eval, Unit, Int](0)
+    val count = stackSafeTestSize
+    val result = (0 until count).foldLeft(unit) { (acc, _) =>
+      acc.map(_ + 1)
+    }
+    result.run(()).value should ===(((), count))
+  }
+
   test("flatMap is stack safe on repeated left binds when F is") {
     val unit = StateT.pure[Eval, Unit, Unit](())
     val count = stackSafeTestSize