From 5b85396a0c6b46fadbc2c1859fadd0e06cfeff59 Mon Sep 17 00:00:00 2001
From: Alexandru Nedelcu <alexandru@users.noreply.github.com>
Date: Wed, 14 Mar 2018 16:52:44 +0200
Subject: [PATCH] Fix #1733: make Kleisli.flatMap stack safe (#2185)

* Fix #1733: make Kleisli stack safe for Monads w/ a stack safe flatMap

* Fix tests, add test for repeated right binds too

* Discriminate between JVM and JS in tests
---
 core/src/main/scala/cats/data/Kleisli.scala   | 38 ++++++++++++++++---
 .../test/scala/cats/tests/KleisliSuite.scala  | 21 +++++++++-
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/core/src/main/scala/cats/data/Kleisli.scala b/core/src/main/scala/cats/data/Kleisli.scala
index 05857c1d44..fcdb3e8d3e 100644
--- a/core/src/main/scala/cats/data/Kleisli.scala
+++ b/core/src/main/scala/cats/data/Kleisli.scala
@@ -28,19 +28,19 @@ final case class Kleisli[F[_], A, B](run: A => F[B]) { self =>
     Kleisli[G, A, B](run andThen f.apply)
 
   def flatMap[C](f: B => Kleisli[F, A, C])(implicit F: FlatMap[F]): Kleisli[F, A, C] =
-    Kleisli((r: A) => F.flatMap[B, C](run(r))((b: B) => f(b).run(r)))
+    Kleisli.shift(a => F.flatMap[B, C](run(a))((b: B) => f(b).run(a)))
 
   def flatMapF[C](f: B => F[C])(implicit F: FlatMap[F]): Kleisli[F, A, C] =
-    Kleisli(a => F.flatMap(run(a))(f))
+    Kleisli.shift(a => F.flatMap(run(a))(f))
 
   def andThen[C](f: B => F[C])(implicit F: FlatMap[F]): Kleisli[F, A, C] =
-    Kleisli((a: A) => F.flatMap(run(a))(f))
+    Kleisli.shift(a => F.flatMap(run(a))(f))
 
   def andThen[C](k: Kleisli[F, B, C])(implicit F: FlatMap[F]): Kleisli[F, A, C] =
     this andThen k.run
 
   def compose[Z](f: Z => F[A])(implicit F: FlatMap[F]): Kleisli[F, Z, B] =
-    Kleisli((z: Z) => F.flatMap(f(z))(run))
+    Kleisli.shift((z: Z) => F.flatMap(f(z))(run))
 
   def compose[Z](k: Kleisli[F, Z, A])(implicit F: FlatMap[F]): Kleisli[F, Z, B] =
     this compose k.run
@@ -80,7 +80,35 @@ final case class Kleisli[F[_], A, B](run: A => F[B]) { self =>
   def apply(a: A): F[B] = run(a)
 }
 
-object Kleisli extends KleisliInstances with KleisliFunctions with KleisliExplicitInstances
+object Kleisli extends KleisliInstances with KleisliFunctions with KleisliExplicitInstances {
+  /**
+   * Internal API — shifts the execution of `run` in the `F` context.
+   *
+   * Used to build Kleisli values for `F[_]` data types that implement `Monad`,
+   * in which case it is safer to trigger the `F[_]` context earlier.
+   *
+   * The requirement is for `FlatMap` as this will get used in operations
+   * that invoke `F.flatMap` (e.g. in `Kleisli#flatMap`). However we are
+   * doing discrimination based on inheritance and if we detect an
+   * `Applicative`, then we use it to trigger the `F[_]` context earlier.
+   *
+   * Triggering the `F[_]` context earlier is important to avoid stack
+   * safety issues for `F` monads that have a stack safe `flatMap`
+   * implementation. For example `Eval` or `IO`. Without this the `Monad`
+   * instance is stack unsafe, even if the underlying `F` is stack safe
+   * in `flatMap`.
+   */
+  private[data] def shift[F[_], A, B](run: A => F[B])
+    (implicit F: FlatMap[F]): Kleisli[F, A, B] = {
+
+    F match {
+      case ap: Applicative[F] @unchecked =>
+        Kleisli(r => F.flatMap(ap.pure(r))(run))
+      case _ =>
+        Kleisli(run)
+    }
+  }
+}
 
 private[data] sealed trait KleisliFunctions {
 
diff --git a/tests/src/test/scala/cats/tests/KleisliSuite.scala b/tests/src/test/scala/cats/tests/KleisliSuite.scala
index b6a0b9e264..846ca6e4e5 100644
--- a/tests/src/test/scala/cats/tests/KleisliSuite.scala
+++ b/tests/src/test/scala/cats/tests/KleisliSuite.scala
@@ -9,8 +9,9 @@ import cats.laws.discipline.arbitrary._
 import cats.laws.discipline.eq._
 import org.scalacheck.Arbitrary
 import cats.kernel.laws.discipline.{MonoidTests, SemigroupTests}
-import cats.laws.discipline.{SemigroupKTests, MonoidKTests}
+import cats.laws.discipline.{MonoidKTests, SemigroupKTests}
 import Helpers.CSemi
+import catalysts.Platform
 
 class KleisliSuite extends CatsSuite {
   implicit def kleisliEq[F[_], A, B](implicit A: Arbitrary[A], FB: Eq[F[B]]): Eq[Kleisli[F, A, B]] =
@@ -256,6 +257,24 @@ class KleisliSuite extends CatsSuite {
     kconfig1.run(config) should === (kconfig2.run(config))
   }
 
+  test("flatMap is stack safe on repeated left binds when F is") {
+    val unit = Kleisli.pure[Eval, Unit, Unit](())
+    val count = if (Platform.isJvm) 10000 else 100
+    val result = (0 until count).foldLeft(unit) { (acc, _) =>
+      acc.flatMap(_ => unit)
+    }
+    result.run(()).value
+  }
+
+  test("flatMap is stack safe on repeated right binds when F is") {
+    val unit = Kleisli.pure[Eval, Unit, Unit](())
+    val count = if (Platform.isJvm) 10000 else 100
+    val result = (0 until count).foldLeft(unit) { (acc, _) =>
+      unit.flatMap(_ => acc)
+    }
+    result.run(()).value
+  }
+
   /**
    * Testing that implicit resolution works. If it compiles, the "test" passes.
    */