Porting Haskell inversion counter to Scala

Question

Here's relevant code from Haskell implementation of a routine that counts inversion in a list

mergeAndCount :: Ord a => [a] -> [a] -> (Int,[a])
mergeAndCount l@(x:xs) r@(y:ys) | x < y = let (inv, s) = mergeAndCount xs r in (inv, x:s)
                                | otherwise = let (inv, s) = mergeAndCount l ys in (inv + rest, y:s)
                                                where rest = length l
mergeAndCount l [] = (0, l)
mergeAndCount [] r = (0, r)

I've tried to write similar routine in Scala but it crashes with stack overflow (lazy sorting works, though). Here's the the non-working version:

  def mergeAndCount(l: Stream[Int], r: Stream[Int]) : (Long, Stream[Int]) = {
    (l, r) match {
      case (x#::xs, Empty) => (0, l)
      case (Empty, y#::ys) => (0, r)
      case (x#::xs, y#::ys) => if(x < y) {
        lazy val (i, s) = mergeAndCount(xs, r)
        (i, x#::s)
      } else {
        lazy val (i, s) = mergeAndCount(l, ys)
        (i + l.length, y#::s)
      }
    }
  }

Any ideas on how to make Scala version to behave like the Haskell one?

Answer 1

You can trade heap for stack in a situation like this (where turning the recursive call into a tail call might be complicated) by using trampolines :

import Stream.Empty
import scalaz.std.tuple._
import scalaz.syntax.bifunctor._
import scalaz.Free.Trampoline, scalaz.Trampoline._

def mergeAndCount(
  l: Stream[Int],
  r: Stream[Int]
): Trampoline[(Long, Stream[Int])] = (l, r) match {
  case (_ #:: _, Empty) => done((0, l))
  case (Empty, _ #:: _) => done((0, r))
  case (x #:: xs, y #:: _) if x < y => suspend(
    mergeAndCount(xs, r).map(_.rightMap(x #:: _))
  )
  case (_, y #:: ys) => suspend(
    mergeAndCount(l, ys).map(_.bimap(_ + l.length, y #:: _))
  )
}

Note that I'm using Scalaz 's implementation here, since the standard library's is currently missing some pieces (although that will change soon ).

Now you can write eg the following:

mergeAndCount((1 to 20000).toStream, (2 to 20001).toStream).run

Where this would almost certainly blow the stack if we weren't trampolining the tail calls.

Answer 2

I'd leave this as a comment but unfortunately I don't have the reputation yet to do so...

Anyway, you're likely to run into stack overflow errors if you recurse before the last call in your function - in Scala only tail recursion is optimised not to use the stack. If you can move the recursive call to the last line of each case (which would mean abandoning the laziness) then you might get better results.