Merge pull request #702 from svalaskevicius/fix-foldmap-for-free

fix foldMap stack safety

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)
+      }
     }
 
   /**

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))
+  }
 }