recursion - How to write non-leaking tail-recursive function using Stream.cons in Scala?

Question

When writing a function operating on Stream(s), there are different notions of recursion. The first simple sense is not recursive on the compiler level, since the tail if not evaluated instantly so the function returns immediately, but the returned stream is recursive:

final def simpleRec[A](as: Stream[A]): Stream[B] = 
  if (a.isEmpty) Stream.empty              
  else someB(a.head) #:: simpleRec(a.tail) 

The above notion of recursion doesn't cause any problems. The second one is truly tail-recursive on the compiler level:

@tailrec
final def rec[A](as: Stream[A]): Stream[B] = 
  if (a.isEmpty) Stream.empty              // A) degenerated
  else if (someCond) rec(a.tail)           // B) tail recursion
  else someB(a.head) #:: rec(a.tail)       // C) degenerated

The problem here is that the C) case is detected by the compiler as a non-tailrec call, even if there is no actual call carried out. This can be avoided by factoring out the stream tail into a helper function:

@tailrec
final def rec[A](as: Stream[A]): Stream[B] = 
  if (a.isEmpty) Stream.empty              
  else if (someCond) rec(a.tail)          // B)
  else someB(a.head) #:: recHelp(a.tail)  

@tailrec
final def recHelp[A](as: Stream[A]): Stream[B] = 
  rec(as)

While it compiles, this approach eventually results in a memory leak. Since the tail-recursive rec is eventually called from the recHelp function, the stack frame of the recHelp function holds a reference to the steam head, and doesn't let the stream to be garbage collected until the rec call returns, which can quite long (in terms of recursion steps) depending on the number of calls to B).

Note that even in the helperless case, if the compiler allowed the @tailrec, the memory leak may still be present since the lazy stream tail would in effect create an anonymous object holding reference to the stream head.

Answer 1

A possible workaround is to make the recHelp method not hold reference to the stream head. This can be achieved by passing a wrapped stream to it, and mutating the wrapper to erase the reference from it:

@tailrec
final def rec[A](as: Stream[A]): Stream[B] = 
  if (a.isEmpty) Stream.empty              
  else if (someCond) rec(a.tail)          
  else {
    // don't inline and don't define as def,
    // or anonymous lazy wrapper object would hold reference
    val tailRef = new AtomicReference(a.tail)
    someB(a.head) #:: recHelp(tailRef)  
  }

@tailrec
final def recHelp[A](asRef: AtomicReference[Stream[A]]): Stream[B] = 
  // Note: don't put the content of the holder into a local variable
  rec(asRef.getAndSet(null))

The AtomicReference is just convenience, atomicity is not required in this case, any simple holder object would do.

Also note that since recHelp is wrapped in a stream Cons tail, therefore it will only be evaluated once, and Cons also takes care of synchronization.