Tail recursion : filter nested lists in scala

Question

The following function filters (removes the 1's ) in a nested list:

def filter(list : List[Any], acc : List[Any] = Nil) : List[Any] = {
  list match {
    case Nil => acc
    case (l : List[_]) :: tail =>
      val nested = filter(l)

      if (nested.isEmpty) filter(tail, acc)
      else
        filter(tail, acc :+ nested)

    case 1 :: tail =>     filter(tail, acc)
    case other :: tail => filter(tail, acc :+ other)
  }
}

Input

filter(List(1, 3, 4, List(1, 4 ,5), List(1)))

Output

res0: List[Any] = List(3, 4, List(4, 5))

My question is: How can I make this tail recursive? My main problem is how to handle the nested lists: val nested = filter(l, Nil)

Thanks

Answer 1

Ok, think you can use Free monad to solve this, available, for example, in Typelevel Cats lib.

"org.typelevel" %% "cats-free" % "1.0.1"

Idea behind is to move recursion which could not be converted to tail-recursion to another part of code, where this is feasible

So, if you want some understanding on how to solve such task you need to read this paper by Runar Oli Bjarnason (and probably watch this presentation). After that you can try to solve your problem with trampolines written by yourself (based on paper and video):

First lets re-write your filter in a way (this step is not necessary, I was just under impression of even/odd example from paper, below I show how to do the same based on your variant):

private def filterList(l: List[Any]): Option[Any] = {
  l.foldLeft(Option.empty[List[Any]]) {
    case (acc, next) if acc.isEmpty =>
      filter1(next).map(_ :: Nil)
    case (acc, next) =>
      acc.map(_ ++ filter1(next))
  }
}

private def filter1(a: Any): Option[Any] = a match {
  case 1 => Option.empty[Any]
  case l: List[Any] => filterList(l)
  case t => Some(t)
}

def filter(l: List[Any]): List[Any] = {
  filterList(l) match {
    case Some(l: List[Any]) => l
    case _ => Nil
  }
}

Now let's assume that Trampoline is already available (and it is available, see below) and rewrite this with trampolines:

def filterList(l: List[Any]): Trampoline[Option[Any]] = {
  l.foldLeft[Trampoline[Option[List[Any]]]](Done(Option.empty[List[Any]])) {
    case (acc, next) =>
      acc.flatMap {
        case None => filter1(next).map(_.map(_ :: Nil))
        case v =>
          filter1(next).map (r => v.map(_ ++ r))
      }
  }
}

def filter1(a: Any): Trampoline[Option[Any]] = a match {
  case 1 => Done(Option.empty[Any])
  case l: List[Any] => More(() => filterList(l))
  case t => Done(Some(t))
}

def filter(l: List[Any]): List[Any] = {
  filterList(l).runT match {
    case Some(l: List[Any]) => l
    case _ => Nil
  }
}

Trampoline implementation (based on paper and video)

sealed trait Trampoline[+A] {
  final def runT: A =
    this match {
      case Done(a) => a
      case More(k) => k().runT
      case t: FlatMap[Any, A] => t match {
        case Done(a) FlatMap f => f(a).runT
        case More(k) FlatMap f => k().flatMap(f).runT
        case FlatMap(xg: FlatMap[Any, A], f) =>
          xg.a.flatMap(a => xg.f(a)).runT
      }
    }

  def map[B](f: A => B): Trampoline[B] =
    flatMap(x => More(() => Done(f(x))))

  def flatMap[B](f: A => Trampoline[B]): Trampoline[B] = this match {
    case FlatMap(a, g: Any) => FlatMap(a, (x: Any) => g(x) flatMap f)
    case x => FlatMap(x, f)
  }
}

case class More[+A](k: () => Trampoline[A])
  extends Trampoline[A]

case class Done[+A](result: A)
  extends Trampoline[A]

case class FlatMap[A, B](a: Trampoline[A], f: A => Trampoline[B])
  extends Trampoline[B]

Now, with all these parts you can check that now you have working stack-less recursive implementation. I used such code for test:

def check(i: Int, output: Boolean) {
  val l = (1 to i).foldLeft(List[Any](1, 2, 3)) {
    case (l, _) =>
      List[Any](1, 2, 3, l, List(1))
  }

  val res = filter(l)
  if (output) {
    println(s"Depth: $i, result: " + res)
  }
}

check(3, true)
check(10000, false)

It fails on non-trampoline approach with StackOverflowError and completes successfully on modified one.

Update with Cats

I tried same approach with cats-free lib. It works perfectly as well:

import cats.free._
import cats.implicits._

def filterList(l: List[Any]): Trampoline[Option[Any]] = {
  l.foldLeft[Trampoline[Option[List[Any]]]](Trampoline.done(Option.empty[List[Any]])) {
    case (acc, next) =>
      acc.flatMap {
        case None => filter1(next).map(_.map(_ :: Nil))
        case v =>
          filter1(next).map (r => v.map(_ ++ r))
      }
  }.map(_.map(identity[Any]))
}

def filter1(a: Any): Trampoline[Option[Any]] = a match {
  case 1 => Trampoline.done(Option.empty[Any])
  case l: List[Any] => Trampoline.defer(filterList(l))
  case t => Trampoline.done(Some(t))
}

def filter(l: List[Any]): List[Any] = {
  filterList(l).run match {
    case Some(l: List[Any]) => l
    case _ => Nil
  }
}

Using Cats Trampoline with original code

I tried to do the same with your original code, here is the result:

import cats.free._
import cats.implicits._

private def filterInner(list : List[Any], acc : List[Any] = Nil) : Trampoline[List[Any]] = {
  list match {
    case Nil => Trampoline.done(acc)
    case (l : List[_]) :: tail =>
      Trampoline.defer(filterInner(l)).flatMap {
        case Nil => Trampoline.defer(filterInner(tail, acc))
        case nested => Trampoline.defer(filterInner(tail, acc :+ nested))
      }
    case 1 :: tail =>     Trampoline.defer(filterInner(tail, acc))
    case other :: tail => Trampoline.defer(filterInner(tail, acc :+ other))
  }
}
def filter(list : List[Any]): List[Any] = filterInner(list).run

Note, that every recursive call to filterInner is wrapped to Trampoline.defer, it is done to eliminate recursion. You can test it by removing wrapping in case (l : List[_]) :: tail part and run on my test example.