diff --git a/free/src/main/scala/cats/free/Free.scala b/free/src/main/scala/cats/free/Free.scala index 7acd6be79a..9ed9f7c174 100644 --- a/free/src/main/scala/cats/free/Free.scala +++ b/free/src/main/scala/cats/free/Free.scala @@ -116,25 +116,22 @@ sealed abstract class Free[S[_], A] extends Product with Serializable { runM2(this) } - /** Takes one evaluation step in the Free monad, re-associating left-nested binds in the process. */ - @tailrec - final def step: Free[S, A] = this match { - case Gosub(Gosub(c, f), g) => c.flatMap(cc => f(cc).flatMap(g)).step - case Gosub(Pure(a), f) => f(a).step - case x => x - } - /** * Catamorphism for `Free`. * * Run to completion, mapping the suspension with the given transformation at each step and * accumulating into the monad `M`. */ + @tailrec final def foldMap[M[_]](f: S ~> M)(implicit M: Monad[M]): M[A] = - step match { + this match { case Pure(a) => M.pure(a) case Suspend(s) => f(s) - case Gosub(c, g) => M.flatMap(c.foldMap(f))(cc => g(cc).foldMap(f)) + case Gosub(c, g) => c match { + case Suspend(s) => g(f(s)).foldMap(f) + case Gosub(cSub, h) => cSub.flatMap(cc => h(cc).flatMap(g)).foldMap(f) + case Pure(a) => g(a).foldMap(f) + } } /** diff --git a/free/src/test/scala/cats/free/FreeTests.scala b/free/src/test/scala/cats/free/FreeTests.scala index 57f26ef781..01770147e6 100644 --- a/free/src/test/scala/cats/free/FreeTests.scala +++ b/free/src/test/scala/cats/free/FreeTests.scala @@ -28,4 +28,26 @@ class FreeTests extends CatsSuite { x.mapSuspension(NaturalTransformation.id[List]) should === (x) } } + + test("foldMap is stack safe") { + trait FTestApi[A] + case class TB(i: Int) extends FTestApi[Int] + + type FTest[A] = Free[FTestApi, A] + + def tb(i: Int): FTest[Int] = Free.liftF(TB(i)) + + def a(i: Int): FTest[Int] = for { + j <- tb(i) + z <- if (j<10000) a(j) else Free.pure[FTestApi, Int](j) + } yield z + + def runner: FTestApi ~> Id = new (FTestApi ~> Id) { + def apply[A](fa: FTestApi[A]): Id[A] = fa match { + case TB(i) => i+1 + } + } + + assert(10000 == a(0).foldMap(runner)) + } }