Composition that have less chance to blow the stack.

composition.scala

object StackTest {
  def main(args : Array[String]) : Unit = {
    val f1 = (1 to 100000).foldLeft(ComposeQueue(identity[Int])) { (f, _) => f.andThen(_ + 1) }
    println(f1(0))

    val f2 = (1 to 1000000).foldLeft(ComposeQueue(identity[Int])){ (f, _) => f.andThen(_ + 1) }
    println(f2(0))

    val f3 = (1 to 1000000).foldLeft(ComposeQueue(identity[Int])) { (f, _) => f.compose(_ + 1) }
    println(f3(0))
  }
}


object ComposeQueue {
  final val MAX_DEPTH = 1000

  @inline def apply[A,B](f : A => B) : ComposeQueue[A,B] =
    f match {
      case cq : ComposeQueue[A,B] => cq
      case _                      => BootStrap(f,1)
    }
}

import ComposeQueue._

sealed abstract class ComposeQueue[-A, +B] extends (A => B) {
  /** Compose this function by a BootStrap on the right
    *
    * represents (c: C) => this(bs(c))
    *
    * @param bs
    * @tparam C
    * @tparam A2
    * @return
    */
  def append[C,A2 <: A](bs : BootStrap[C,A2]) : ComposeQueue[C,B]

  /** Compose this function by a BootStrap on the left
    *
    * represents (a: A) => bs(this(a))
    *
    * @param bs
    * @tparam C
    * @tparam B2
    * @return
    */
  def prepend[C,B2 >: B](bs : BootStrap[B2,C]) : ComposeQueue[A,C]


  /** Compose this function by a ComposeQueue on the right
    *
    * represents (c: C) => this(cq(c))
    *
    * @param cq
    * @tparam C
    * @tparam A2
    * @return
    */
  def concat[C,A2 <: A](cq : ComposeQueue[C,A2]) : ComposeQueue[C,B]

  /** Normal application */
  def apply(a : A) : B


  @inline final override def compose[C](g: C => A): ComposeQueue[C,B] =
    g match {
      case cq : ComposeQueue[C,A] => concat(cq)
      case _                      => append(BootStrap(g,1))
    }


  @inline final override def andThen[C](g: B => C): ComposeQueue[A,C] =
    g match {
      case cq : ComposeQueue[B,C] => cq.concat(this)
      case _                      => prepend(BootStrap(g,1))
    }
}

/** Base case of ComposeQueue, just an ordinary function.
  * depth indicates the number of functions composed to form this one.
  * If depth is too big, stack overflow occurs.
  *
  * @param function a function
  * @param depth    the number of function composition made to form this one.
  * @tparam A
  * @tparam B
  */
final case class BootStrap[A,B](function : A => B, depth : Int) extends ComposeQueue[A,B] {

  // Just a function ...
  def apply(a : A) : B = function(a)

  @inline def append[C,A2 <: A](bs : BootStrap[C,A2]) : ComposeQueue[C,B] =
    if (depth + bs.depth <= MAX_DEPTH)
      unsafeAppend(bs)
    else
      Deep[C,A2,A,B](this, Vector.empty, bs)


  @inline def prepend[C,B2 >: B](bs : BootStrap[B2,C]) : ComposeQueue[A,C] =
    if (depth + bs.depth <= MAX_DEPTH)
      bs.unsafeAppend(this)
    else
      Deep[A,B,B2,C](bs, Vector.empty, this)


  @inline def concat[C,A2 <: A](cq : ComposeQueue[C,A2]) : ComposeQueue[C,B] =
    cq.prepend(this)

  /** Compose two BootStrap. It is unsafe because it may break the invariant that
    * depth <= MAX_DEPTH . Used as a helper for append and prepend.
    *
    * @param bs
    * @tparam C
    * @tparam A2
    * @return
    */
  @inline def unsafeAppend[C,A2 <: A](bs : BootStrap[C,A2]) : BootStrap[C,B] =
    BootStrap[C,B]((c:C) => function(bs.function(c)), depth + bs.depth)
}

/** Too much ordinary function composition blows the stack.
  * To prevent this, Deep stores composition into the heap as
  * a Vector of functions.
  *
  * It represents the function:
  *
  * (x : A) => outer(queue(n-1)(queue(n-2)( .....(queue(0)(inner(x))))))
  *
  * where n is the size of the queue.
  *
  *
  * @param outer the outer function this composition represents
  * @param queue the functions that compose this composition between outer and inner
  * @param inner the inner function this composition represents
  * @tparam A
  * @tparam X
  * @tparam Y
  * @tparam B
  */
final case class Deep[A,X,Y,B](outer : BootStrap[Y,B],
                               queue : Vector[Any => Any],
                               inner : BootStrap[A, X]
                              ) extends ComposeQueue[A,B] {
  type inner = X
  type outer = Y

  def apply(a : A) : B = outer {
    var i : Any = inner(a)
    var j = queue.size - 1
    while (j >= 0) {
      i  = queue(j)(i)
      j -= 1
    }
    i.asInstanceOf[outer]
  }

  @inline def append[C,A2 <: A](bs : BootStrap[C,A2]) : ComposeQueue[C,B] =
    if (inner.depth + bs.depth <= MAX_DEPTH)
      Deep(outer, queue, inner.unsafeAppend(bs))
    else
      Deep[C,A2,Y,B](outer, queue :+ inner.function.asInstanceOf[Any => Any] , bs)

  @inline def prepend[C,B2 >: B](bs : BootStrap[B2,C]) : ComposeQueue[A,C] =
    if (outer.depth + bs.depth <= MAX_DEPTH)
      Deep(bs.unsafeAppend(outer), queue, inner)
    else
      Deep[A,X,B2,C](bs, queue.+:(outer.function.asInstanceOf[Any => Any]) , inner)


  def concat[C,A2 <: A](cq : ComposeQueue[C,A2]) : ComposeQueue[C,B] =
    cq match {
      case bs@BootStrap(_,_) =>
        append(bs)

      case cq2 : Deep[C,_,_,A] =>
        if (queue.isEmpty)
          cq2.prepend(inner).prepend(outer.asInstanceOf[BootStrap[inner, B]]) // Because queue.isEmpty => inner = outer
        else if (cq2.queue.isEmpty)
          append(cq2.outer.asInstanceOf[BootStrap[cq2.inner, A]]).append(cq2.inner)
        else if (inner.depth + cq2.outer.depth <= MAX_DEPTH) {
          val f: Any => X = (x : Any) => inner.function(cq2.outer.function(x.asInstanceOf[cq2.outer]))
          Deep(outer, (queue :+ f) ++ cq2.queue, cq2.inner)
        } else
          Deep(outer, ((queue :+ inner.function.asInstanceOf[Any => Any]) :+ cq2.outer.function.asInstanceOf[Any => Any]) ++ cq2.queue, cq2.inner)
    }
}