diff --git a/core/src/main/scala/scalaz/Free.scala b/core/src/main/scala/scalaz/Free.scala index 7667b94e2d..bbc9d1052a 100644 --- a/core/src/main/scala/scalaz/Free.scala +++ b/core/src/main/scala/scalaz/Free.scala @@ -193,12 +193,16 @@ sealed abstract class Free[S[_], A] { * Runs to completion, mapping the suspension with the given transformation at each step and * accumulating into the monad `M`. */ + @annotation.tailrec final def foldMap[M[_]](f: S ~> M)(implicit M: Monad[M]): M[A] = - step match { + this match { case Return(a) => M.pure(a) case Suspend(s) => f(s) - // This is stack safe because `step` ensures right-associativity of Gosub - case Gosub(x, g) => M.bind(x foldMap f)(c => g(c) foldMap f) + case Gosub(x, g) => x match { + case Suspend(s) => g(f(s)).foldMap(f) + case Gosub(cSub, h) => cSub.flatMap(cc => h(cc).flatMap(g)).foldMap(f) + case Return(a) => g(a).foldMap(f) + } } import Id._ diff --git a/tests/src/test/scala/scalaz/FreeTest.scala b/tests/src/test/scala/scalaz/FreeTest.scala index 907fd0ccfc..2e9523eb27 100644 --- a/tests/src/test/scala/scalaz/FreeTest.scala +++ b/tests/src/test/scala/scalaz/FreeTest.scala @@ -91,6 +91,25 @@ object FreeTest extends SpecLite { checkAll(bindRec.laws[FreeOption]) } + "foldMap is stack safe" ! { + val n = 1000000 + trait FTestApi[A] + case class TB(i: Int) extends FTestApi[Int] + + def a(i: Int): Free[FTestApi, Int] = for { + j <- Free.liftF(TB(i)) + z <- if (j < n) a(j) else Free.pure[FTestApi, Int](j) + } yield z + + val runner = new (FTestApi ~> Id.Id) { + def apply[A](fa: FTestApi[A]) = fa match { + case TB(i) => i + 1 + } + } + + a(0).foldMap(runner) must_=== n + } + "List" should { "not stack overflow with 50k binds" in { val expected = Applicative[FreeList].point(())