How can I make this Scala function (a "flatMap" variant) tail recursive?

Question

I'm having a look at the following code

http://aperiodic.net/phil/scala/s-99/p26.scala

Specifically

def flatMapSublists[A,B](ls: List[A])(f: (List[A]) => List[B]): List[B] = 
    ls match {
      case Nil => Nil
      case sublist@(_ :: tail) => f(sublist) ::: flatMapSublists(tail)(f)
    }

I'm getting a StackOverflowError for large values presumably because the function is not tail recursive. Is there a way to transform the function to accommodate large numbers?

Answer 1

It is definitely not tail recursive. The f(sublist) ::: is modifying the results of the recursive call, making it a plain-old-stack-blowing recursion instead of a tail recursion.

One way to ensure that your functions are tail recursive is to put the @annotation.tailrec on any function that you expect to be tail recursive. The compiler will report an error if it fails to perform the tail call optimization.

For this, I would add a small helper function that's actually tail recursive:

def flatMapSublistsTR[A,B](ls: List[A])(f: (List[A]) => List[B]): List[B] = {
  @annotation.tailrec
  def helper(r: List[B], ls: List[A]): List[B] = {
    ls match {
      case Nil => r
      case sublist@(_ :: tail) => helper(r ::: f(sublist), tail)
    }
  }
  helper(Nil, ls)
}

For reasons not immediately obvious to me, the results come out in a different order than the original function. But, it looks like it works :-) Fixed.

Answer 2

Here is another way to implement the function:

scala> def flatMapSublists[A,B](ls: List[A])(f: (List[A]) => List[B]): List[B] =
     |   List.iterate(ls, ls.size)(_.tail).flatMap(f)
flatMapSublists: [A, B](ls: List[A])(f: List[A] => List[B])List[B]

A simply comparison between dave's flatMapSublistsTR and mine:

scala> def time(count: Int)(call : => Unit):Long = {
     |    val start = System.currentTimeMillis
     |    var cnt =  count
     |    while(cnt > 0) {
     |       cnt -= 1
     |       call
     |    }
     |    System.currentTimeMillis - start
     | }
time: (count: Int)(call: => Unit)Long

scala> val xs = List.range(0,100)

scala> val fn = identity[List[Int]] _
fn: List[Int] => List[Int] = <function1>

scala> time(10000){ flatMapSublists(xs)(fn) }
res1: Long = 5732

scala> time(10000){ flatMapSublistsTR(xs)(fn) }
res2: Long = 347232

Where the method flatMapSublistsTR is implemented as:

def flatMapSublistsTR[A,B](ls: List[A])(f: (List[A]) => List[B]): List[B] = {
  @annotation.tailrec
  def helper(r: List[B], ls: List[A]): List[B] = {
    ls match {
      case Nil => r
      case sublist@(_ :: tail) => helper(r ::: f(sublist), tail)
    }
  }
  helper(Nil, ls)
}