Dealing with external effects building a pure functional api using Scala is an interesting topic in the context of pure functional programs design. Today there are great frameworks four developing your external interactions apis like Cats Effect or ZIO, to name a few. However, it is recommendable for developers that want to work with pure IO to become familiar with the concepts that I will explain in this post.

Case

Imagine a groupedWithin operator that pushes downstream a mini-batch of data in a List[A] in every interval of t seconds. For each element of this collection it is necessary to execute an IO operation that will produce a runtime value, that could be considered as a kind of context for the next operation.

Because this functionality can be reused for other modules we may think in building a new library in a pure functional style: declaring a new IO parameterized type and some basic functions that will act as the primitives that will be composed to create more complex ones. For these basic units of execution to be composable they must be pure functions.

Building an api to offer this functionality through a generic IO type can be good option because it provides the mechanisms of abstraction over different implementations: Database, Queue, TCP socket, etc …

The first problem that comes in to my mind is that if I use pure functions, how am I going to avoid the Stack Overflow if I execute the effect over a collection of several thousands of elements?. Remember that the api design does not include any side effect, only functional composition over the IO type.

Note: An IO effect represents an interaction with the external world, so it is in essence a side effect. You can find different points of view of this: “Using IO makes your input/output functions pure” or the “it just represents that your code is obvious impure”.

Antecedents

There are a lot of explanations about this (and a little old) topic and I would suggest reading the red book “Functional Programming in Scala” by Chiusano and Bjarnason, chapter 13 “External Effects and IO”.

Scala supports tail call optimization if the last step of the function is the recursive call. In this case the compiler is able to optimize the recursion process, so it is executed in a constant stack space, thereby avoiding the Stack Overflow exception. There is a good explanation in the Alvin Alexander´s blog.

Example 1: Not tail recursive, the sum call is not in the tail position. It will blow the stack if the list is several thousands of elements long.

def sum(list: List[Int]): Int = list match {
case Nil => 0
case x :: xs => x + sum(xs)
}

Example2: Tail recursive function. the compiler is able to optimize the function to use the same stack frame.

def sumWithAccumulator(list: List[Int], accumulator: Int): Int = {
list match {
case Nil => accumulator
case x :: xs => sumWithAccumulator(xs, accumulator + x)
}
}

Abstracting over the IO Type

We can think that an IO type represents an interaction with the outside world, there is no need in thinking about any specific implementation yet. The only thing that it is necessary for the moment is that the effect execution creates a runtime value for the successive one.

Having in mind that, the first primitive functions to be implemented are flatMap, map and execute. The two first are mandatory to use a Scala for-comprehension to compose functions that produce IO effects, and the latter is the one that runs the effect producing a value. A very simple example would be like the following:

trait IO[A] { self =>
def execute : A
def flatMap[B](f: A => IO[B]): IO[B] = {
new IO[B] {
def execute = f(self.execute).execute
}
}

def map[B](f: A => B): IO[B] =
new IO[B] { def execute = f(self.execute) }
}

The IO´s type parameter corresponds to the value that the execution yields. For example, consider that you need to print messages in the console, you need to create a new type PrintConsole as a IO implementation that only writes a message in the console and does not give any value:

case class PrintToConsole(msg: String) extends IO[Unit] {
override def execute: Unit = println(msg)
}

Now, you can compose several effects using a for comprehension structure:

for {
_ <- PrintToConsole("Hello")
_ <- PrintToConsole("World")
} yield()

The result, as expected is:

Hello
World

Now, we think to repeat this effect but for a set of elements, as we said before, those that come from a stream as a collection of variable size. To process the elements of this collection we think in creating a fold function:

def fold[A, B](l: List[A])(z: B)(f: (B, A) => IO[B]): IO[B] = {
l match {
case h :: tail => f(z, h) flatMap (z2 => fold(tail)(z2)(f))
case _ => new IO[B] {
override def execute: B = z
}
}
}

This function is part of the api and uses the flatMap method of the IO. We test with a collection of five messages:

val list = List.fill(5)("Hello")
val result2 = IO.fold(list)() { (_, b) =>
PrintToConsole(b)
}
result2.execute

This function looks elegant and pure and it actually works as expected:

Hello
Hello
Hello
Hello
Hello

However, if the collection is of size several thousands elements it will crash with the following message:

Exception in thread “main” java.lang.StackOverflowError
at java.base/java.io.FileOutputStream.write(FileOutputStream.java:354)

If we examine the problem in detail we can see that we can not leverage the flatMap of our IO type if it is implemented in that way because each iteration adds a stack frame and the compiler is not able to detect that problem so there is not any kind of optimization.

In short, although our fold function looks pure, it is not. Depending on the number of elements in the collection it behaves quite different.

Trampolines

If you analyze the Cats IO documentation you can see something that people that are not familiar with these concepts may overlook: “IO is trampolined in its flatMap valuation”. Take a look to the following example:

import cats.effect.IOdef fib(n: Int, a: Long, b: Long): IO[Long] =
IO.suspend {
if (n > 0) fib(n - 1, b, a + b) else IO.pure(a)
}

The documentation tries to clarify this concept and says this about the suspend function:

“So it is useful for suspending effects, but that defers the completion of the returnedIO to some other reference. It’s also useful for modeling stack safe, tail recursive loops”

As I a said before, if you are not familiar with this terminology and more specifically with the concept of Trampoline it sounds weird.

I will try to explain that concept redesigning the fold function using trampolines.

Folding a List without blowing the Stack

The first thing that we need is to take a look to the meaning of Trampoline. There are a lot of definitions and blogs about this topic, I still recommend the red book because besides explaining the concept and proposing some interesting exercises, it gets deeper and describes the connection between Trampolines and the Free Monad very well. There is a very interesting paper of Runar Bjarnason about these concepts that I recommend reading carefully.

We create a new ADT. If you are not familiar with this denomination, the Alvin Alexander´s post regarding algebras is a good point to start. Basically, it is a sealed trait with the map and flatMap functions and three case classes that we need to create a structure in memory that will represent a recursive process:

sealed trait TailRec[A] {
def map[B](f: A => B): TailRec[B] = flatMap(f andThen (Return(_)))
def flatMap[B](f: A => TailRec[B]): TailRec[B] = FlatMap(this, f)
}
final case class Return[A](a: A) extends TailRec[A]
final case class Suspend[A](resume: () => TailRec[A])
extends TailRec[A]
final case class FlatMap[A, B](sub: TailRec[A], k: A => TailRec[B]) extends TailRec[B]

If you take a look to the flatMap function of the TailRec type, we can see that it works quite different than the previous implementation. I will explain in more detail after the new fold implementation. We will call the new function foldSafe:

def foldSafe[A, B]
(l: List[A])(z: B)
(f: (B, A) => TailRec[IO[B]]): TailRec[IO[B]] = {

l match {
case h :: tail =>
f(z, h)
.flatMap (z2 => Suspend(() =>
foldSafe(tail)(z2.execute)(f)))
case Nil =>
Return(new IO[B] {
override def execute: B = z
})
}
}

Now, we use the TailRec functions to implement the fold operation, and we use the three different TailRec data types in the process. The flatMap operation has the same semantic, but it is implemented quite different, it does not compose directly two functions as the IO flatMap. Instead, a new FlatMap data type instance is created which holds the “self” effect and the suspension of the recursive step. Note that the Suspend instance holds the recursive call, it does not execute it. It is very important to understand that TailRec is a kind of wrapper that describes a program in a memory structure, therefore it needs to be interpreted.

The Interpreter

To explain how the interpreter works I will borrow the factorial calculation example using the Trapoline:

def fac(n: Int): TailRec[Int] = {
if (n == 0)
Return(1)
else
FlatMap[Int, Int](Suspend(() => fac(n - 1)), x => Return(n * x))
}

this function creates a in memory representation like this:

FlatMap(
Return(1),
x => Flatmap(
Return(1 * x), x => Flatmap(
Return(2 * x), x => Flatmap(
Return(3 * x), x => Flatmap(
Return(4 * x), x => Return(5 * x)
)
)
)
)
)

To interpret this memory structure we can use the following function:

def run[A](tr: TailRec[A]): A = tr match {
case Return(a) => a
case Suspend(r) => run(r())
case FlatMap(x, f) => x match {
case Return(a) => run(f(a))
case Suspend(r) => run(FlatMap(r(), f))
case FlatMap(y, g) => run(y.flatMap(g(_) flatMap f))
}
}

Looking at the implementation, we now can see that this actually is a recursive function that is tail recursive thus the compiler can optimize it.

Testing the API

For testing our new foldSafe function, we create the a new collection and pass the function that represent the side effect as parameter, and this must be a (A, B) => TailRec[IO[B]] function:

val list = List.fill(50000)("Hello")

val result = IO.foldSafe(list)() { (_, b) =>
Return(PrintToConsole(b))
}

IO.run(result)

It works as expected and the Stack Overflow error does not happen for much bigger lists.

Conclusions

As I mentioned at the beginning of this post, different IO monad implementations solve this problem and developers, normally, don´t have to be worried about the stack although sometimes some words related to trampolines can appear and it is good to know where they come from.