foldStep to Free, remove Functor in FoldableFree

typelevel · Aug 22, 2017 · 622aec2 · 622aec2
1 parent 98879bd
commit 622aec2
Showing 1 changed file with 37 additions and 12 deletions.
diff --git a/free/src/main/scala/cats/free/Free.scala b/free/src/main/scala/cats/free/Free.scala
@@ -54,6 +54,26 @@ sealed abstract class Free[S[_], A] extends Product with Serializable {
       }
   }
 
+  /**
+   * Re-associate any left-nested binds to the right, pull the first suspension to the top
+   * and then pass the result to one of the callbacks. Borrowed from scalaz.
+   */
+  @tailrec
+  final def foldStep[B](
+    onPure: A => B,
+    onSuspend: S[A] => B,
+    onFlatMapped: ((S[X], X => Free[S, A]) forSome { type X }) => B
+  ): B = this match {
+    case Pure(a) => onPure(a)
+    case Suspend(a) => onSuspend(a)
+    case FlatMapped(c,f) => c match {
+      case Pure(a) => f(a).foldStep(onPure, onSuspend, onFlatMapped)
+      case Suspend(a) => onFlatMapped((a,f))
+      case FlatMapped(fmc,fmf) =>
+        fmc.flatMap(y => fmf(y).flatMap(f)).foldStep(onPure, onSuspend, onFlatMapped)
+    }
+  }
+
   /**
    * Run to completion, using a function that extracts the resumption
    * from its suspension functor.
@@ -253,26 +273,33 @@ object Free extends FreeInstances {
 
 private trait FreeFoldable[F[_]] extends Foldable[Free[F, ?]] {
 
-  implicit def FoldableF: Foldable[F]
-  implicit def FunctorF: Functor[F]
+  implicit def F: Foldable[F]
 
   override final def foldLeft[A, B](fa: Free[F, A], b: B)(f: (B, A) => B): B =
-    fa.fold(f(b, _), FoldableF.foldLeft(_, b)((bb, freeA) => foldLeft(freeA, bb)(f)))
+    fa.foldStep(
+      a => f(b,a),
+      fa => F.foldLeft(fa, b)(f),
+      { case (fx, g) => F.foldLeft(fx, b)((bb, x) => foldLeft(g(x), bb)(f)) }
+    )
 
   override final def foldRight[A, B](fa: Free[F, A], lb: Eval[B])(f: (A, Eval[B]) => Eval[B]): Eval[B] =
-    fa.fold(f(_, lb), FoldableF.foldRight(_, lb)( (freeA, eb) => foldRight(freeA, eb)(f)))
+    fa.foldStep(
+      a => f(a, lb),
+      fa => F.foldRight(fa, lb)(f),
+      { case (fx, g) => F.foldRight(fx, lb)( (a, lbb) => foldRight(g(a), lbb)(f)) }
+    )
 }
 
 private trait FreeTraverse[F[_]] extends Traverse[Free[F, ?]] with FreeFoldable[F] {
   implicit def TraversableF: Traverse[F]
 
-  def FoldableF: Traverse[F] = TraversableF
+  def F: Foldable[F] = TraversableF
 
   override final def traverse[G[_], A, B](fa: Free[F, A])(f: A => G[B])(implicit G: Applicative[G]): G[Free[F, B]] =
-    fa.fold(
-      a => G.map(f(a))(Free.pure(_)),
-      ffreeA => G.map(TraversableF.traverse(ffreeA)(traverse(_)(f)))(Free.roll(_))
-    )
+    fa.resume match {
+      case Right(a) => G.map(f(a))(Free.pure(_))
+      case Left(ffreeA) => G.map(TraversableF.traverse(ffreeA)(traverse(_)(f)))(Free.roll(_))
+    }
 
   // Override Traverse's map to use Free's map for better performance
   override final def map[A, B](fa: Free[F, A])(f: A => B): Free[F, B] = fa.map(f)
@@ -282,12 +309,10 @@ sealed private[free] abstract class FreeInstances {
 
   implicit def catsFreeFoldableForFree[F[_]](
     implicit
-    functorF: Functor[F],
     foldableF: Foldable[F]
   ): Foldable[Free[F, ?]] =
     new FreeFoldable[F] {
-      val FoldableF = foldableF
-      val FunctorF = functorF
+      val F = foldableF
     }
 
   implicit def catsFreeTraverseForFree[F[_]](