Merge pull request #713 from ceedubs/free-foldmap-bug

Revert regression in Free.foldMap
typelevel · Dec 7, 2015 · 492b030 · 492b030
2 parents bc07d85 + 6feb7d0
commit 492b030
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 12 deletions.
diff --git a/free/src/main/scala/cats/free/Free.scala b/free/src/main/scala/cats/free/Free.scala
@@ -74,6 +74,14 @@ sealed abstract class Free[S[_], A] extends Product with Serializable {
   final def fold[B](r: A => B, s: S[Free[S, A]] => B)(implicit S: Functor[S]): B =
     resume.fold(s, r)
 
+  /** 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
+  }
+
   /**
    * Evaluate a single layer of the free monad.
    */
@@ -122,16 +130,11 @@ sealed abstract class Free[S[_], A] extends Product with Serializable {
    * 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] =
-    this match {
+    step match {
       case Pure(a) => M.pure(a)
       case Suspend(s) => f(s)
-      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)
-      }
+      case Gosub(c, g) => M.flatMap(c.foldMap(f))(cc => g(cc).foldMap(f))
     }
 
   /**

diff --git a/free/src/test/scala/cats/free/FreeTests.scala b/free/src/test/scala/cats/free/FreeTests.scala
@@ -6,6 +6,7 @@ import cats.tests.CatsSuite
 import cats.laws.discipline.{MonadTests, SerializableTests}
 import cats.laws.discipline.arbitrary.function0Arbitrary
 import org.scalacheck.{Arbitrary, Gen}
+import Arbitrary.{arbitrary, arbFunction1}
 
 class FreeTests extends CatsSuite {
   import FreeTests._
@@ -19,7 +20,7 @@ class FreeTests extends CatsSuite {
     }
   }
 
-  test("foldMap is stack safe") {
+  ignore("foldMap is stack safe") {
     trait FTestApi[A]
     case class TB(i: Int) extends FTestApi[Int]
 
@@ -53,11 +54,26 @@ object FreeTests extends FreeTestsInstances {
 }
 
 sealed trait FreeTestsInstances {
+  private def freeGen[F[_], A](maxDepth: Int)(implicit F: Arbitrary[F[A]], A: Arbitrary[A]): Gen[Free[F, A]] = {
+    val noGosub = Gen.oneOf(
+      A.arbitrary.map(Free.pure[F, A]),
+      F.arbitrary.map(Free.liftF[F, A]))
+
+    val nextDepth = Gen.chooseNum(1, maxDepth - 1)
+
+    def withGosub = for {
+      fDepth <- nextDepth
+      freeDepth <- nextDepth
+      f <- arbFunction1[A, Free[F, A]](Arbitrary(freeGen[F, A](fDepth))).arbitrary
+      freeFA <- freeGen[F, A](freeDepth)
+    } yield freeFA.flatMap(f)
+
+    if (maxDepth <= 1) noGosub
+    else Gen.oneOf(noGosub, withGosub)
+  }
+
   implicit def freeArbitrary[F[_], A](implicit F: Arbitrary[F[A]], A: Arbitrary[A]): Arbitrary[Free[F, A]] =
-    Arbitrary(
-      Gen.oneOf(
-        A.arbitrary.map(Free.pure[F, A]),
-        F.arbitrary.map(Free.liftF[F, A])))
+    Arbitrary(freeGen[F, A](6))
 
   implicit def freeEq[S[_]: Monad, A](implicit SA: Eq[S[A]]): Eq[Free[S, A]] =
     new Eq[Free[S, A]] {