Composing complex programs with ZIO
Tags: functional, zio, scala
Published: 2021-03-01 22:50:29

---
title: Composing complex programs with ZIO
published: true
description: ZIO is the most powerful monad to write programs, how do I use it correctly?
tags: fp, zio, functional, scala
//cover_image: https://imgflip.com/i/4zzf20
---

There is a fundamental mental shift a procedural programmer must make when embracing functional programming. In this lesson we'll learn to visualize complex ZIO programs to compose them correctly.

## The art of type chasing

Types in Scala give us confidence that the program is doing what we want and that the quality of the data is high enough. With ZIO they also help us catching nasty bugs at compile time instead of runtime.

The canonical ZIO program has this type:
```scala
ZIO[ZEnv, Nothing, ExitCode]
```
This is our end game, every ZIO program we can possibly conceive, if we transform it into (chase) this type, then we can easily run it.

Usually what we start with will be something like

```scala
ZIO[ZEnv with MyDatabase, DomainSpecificErrors, Unit]
```
to make it coincide with our canonical ZIO, we need to act in two different directions.

- The environment type needs to become more generic (up to `ZEnv`)
- The error and return type need to become more specific (error down to `Nothing` and return to `ExitCode`)

By using ZIO provided methods we can creatively chase the canonical types.

### Cat and mouse chasing

As fun example of type chasing we can create a program of a cat chasing a mouse.

We will have a mouse program and a cat program that will run concurrently. Both programs will share a variable to track the mouse distance from the cat.

### Creating a variable
Our variable, tracking the lead the mouse has on the cat, needs to be shared by both the cat and the mouse program.

To create a variable that is safe to use concurrently we can write the following program
```scala
import zio.Ref

val makeMouseLead: UIO[Ref[Int]] = Ref.make(10)
```
This is a `UIO` program which expands to a `ZIO[Any, Nothing, Ref[Int]]`. It can run in any environment, never fails, and returns a `Ref` of an integer.

`Ref` is a reference to an immutable variable that we can check and swap. You can picture a `Ref` like a single place datum cradle. We will use it to read the distance between cat and mouse and update it whenever either of them advances.

### Digital mouse that runs away

To represent the mouse getting away from the cat we can design a simple recursive program
```scala
import zio.clock.sleep
import zio.console.putStrLn
import zio.random.nextIntBounded
import zio.duration._
import zio._

def mouseProgram(mouseLead: Ref[Int]): URIO[ZEnv, Unit] =
  for {
    cm <- nextIntBounded(2)
    _  <- mouseLead.update(_ + cm)
    _  <- putStrLn(s"Mouse advances by $cm cm")
    _  <- sleep(100.milliseconds)
    _  <- mouseProgram(mouseLead)
  } yield ()
```

Let's dissect this automatic mouse:
- We are accepting the reference to an `Int` variable, the mouse lead on the cat.
- A random number, less than 2, is drawn, representing how many centimeters the mouse advances.
- We add the above number to the current lead and put the sum in the `Ref`. This will be the new mouse lead.
- Print on the screen the mouse progress
- We then wait 100 milliseconds
- Recursively call the mouse program

Notice that we are using functions coming from `zio.clock`, `zio.console` and `zio.random`, all contained in `ZEnv` so our final type is `URIO[ZEnv, Unit]`.

I wrote `ZEnv` instead of `Clock with Console with Random`. This is allowed because I can always pick a more specific type as environment. The opposite is true for error type and return type.

This program is recursive, it will never return the `Unit` we promised. To make the return type more expressive it would have to return `Nothing` meaning it will never do. But, as said above we cannot put a more specific type for the return channel.

The problem is that the compiler didn't detect our infinite recursion and assumed at some point we are going to return `Unit`. There is a way we can explicitly say that the mouse will run forever.

```scala
def mouseProgram(mouseLead: Ref[Int]): URIO[ZEnv, Nothing] =
  for {
    cm    <- nextIntBounded(2)
    _     <- mouseLead.update(_ + cm)
    _     <- putStrLn(s"Mouse advances by $cm cm")
    _     <- sleep(100.milliseconds)
    recur <- mouseProgram(mouseLead)
  } yield recur
```

By making the return type the mouse program itself we can force write `Nothing` and the compiler will be happy.

Thanks to the ZIO library we can actually do this by skipping the recursive step entirely:

```scala
def mouseProgram(mouseLead: Ref[Int]): URIO[ZEnv, Nothing] =
  (for {
    cm    <- nextIntBounded(2)
    _     <- mouseLead.update(_ + cm)
    _     <- putStrLn(s"Mouse advances by $cm cm")
    _     <- sleep(100.milliseconds)
  } yield ()).forever
```

In this code, we deleted the recursive call and wrapped the for comprehension in parenthesis. With `forever` we created a new program that repeats the for comprehension forever.

### Digital cat that chases

The third program we are going to write is the dual of the mouse. The cat is very similar, but it runs faster and, for every centimeter it advances, the mouse lead diminishes instead of increasing.

```scala
def catProgram(mouseLead: Ref[Int]): URIO[ZEnv, Unit] =
  for {
    cm   <- nextIntBounded(5)
    lead <- mouseLead.updateAndGet(_ - cm)
    _    <- putStrLn(s"Cat advances by $cm cm, mouse is $lead cm away")
    _    <- sleep(100.milliseconds)
    _    <- catchDetector(lead, mouseLead)
  } yield ()
```

The cat program will actually finish at some point: the last instruction uses a special program called `catchDetector`. It considers two possible scenarios: either the mouse has still some lead or the cat is at or past the mouse.

```scala
def catchDetector(currentLead: Int, mouseLead: Ref[Int]): ZIO[ZEnv, Nothing, Unit] =
  currentLead match {
    case lead if lead > 0  => catProgram(mouseLead)
    case lead if lead <= 0 => putStrLn("Cat catches the mouse!")
  }
```
Here, if the mouse is not reached yet, the cat program is called (recursively), otherwise the `catchDetector` returns with a program that prints "Cat catches the mouse!".

We can make the program even more fun by including a third possibility: if the cat goes past the mouse, instead of declaring victory, we assume the mouse escaped. We can express this as an exception. From the point of view of the cat program, this is certainly not a welcome outcome.

First, let's declare the exception as a case class.
```scala
case class MouseEscapedException(catLead: Int)
```
The variable `catLead` will tell us how much the cat went past the mouse.

Secondly, we need to add the third case in the `catchDetector`.
```scala
def catchDetector(currentLead: Int, mouseLead: Ref[Int]): ZIO[ZEnv, MouseEscapedException, Unit] =
  currentLead match {
    case 0                 => putStrLn("Cat catches the mouse!")
    case lead if lead > 0  => catProgram(mouseLead)
    case lead if lead < 0 => ZIO.fail(MouseEscapedException(lead))
  }
```
Now if the `currentLead` is strictly less than 0 we return a program that fails with `MouseEscapedException`. This makes our program fail in turn.

Since the cat program is calling this, we need to update the signature by adding the error modality:
```scala
def catProgram(mouseLead: Ref[Int]): ZIO[ZEnv, MouseEscapedException, Unit] = ...
```
### Putting it all together

We have created a total of 4 ZIO programs. This is a good time to list them all.

- `UIO[Ref[Int]]` will give us an integer variable
- `URIO[ZEnv, Nothing]` programs the mouse to run
- `ZIO[ZEnv, MouseEscapedException, Unit]` programs the cat to chase
- `ZIO[ZEnv, MouseEscapedException, Unit]` subroutine of the cat that can either succeed with a `Unit` or trigger failure with a `MouseEscapedException`.

The diagram below is another representation of these programs.

![Not connected diagram](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/nwe3v5h6txwjzei279h0.png)

Notice, at the center, our canonical ZIO program `URIO[ZEnv, ExitCode]`. In green I represented the environment channel. The canonical ZIO program will provide you a perfectly good `ZEnv` (filled green dot). The program that creates the mouse and also the cat, conversely, want the `ZEnv` (hence the hollow green dot).

In blue I represented the success channel. The mouse catch detector may return a unit when a capture happens, this gets transmitted upwards). The mouse program runs forever, and returns nothing, that's why it doesn't have any blue line going outwards. The canonical ZIO waits patiently for a `ExitCode` and won't accept much else.

In red we see the failures channel that, from the catch detectors propagate upwards. The cat, in turn, may propagate the error.

We have to run the cat and the mouse program, to do that we have to chase the canonical ZIO types.

#### Environment channel

There is not a lot to do for the green lines. `ZEnv` with `ZEnv`, everything matches. We won't worry about this.

#### Failure channel

The canonical ZIO doesn't want to deal with errors we have to stop them before they reach the boundaries of our program. One way to do this is to turn the error into a program that never fails (a ZIO with `Nothing` as error type); `putStrLn` is one such a program. We could log the error to screen and this would return a success `Unit` type. A message saying that the mouse escaped will be sufficient.

#### Success channel

The only type accepted here by the canonical ZIO is `ExitCode`. We said this before, but it's worth repeating that `ExitCode` is just a puffed up integer we give to the operating system. We can always return `ExitCode.success`, an alias for `0`. This means we consider the cat both catching and missing the mouse a success.

The above three conditions are satisfied by `compositeProgram`.

```scala
val compositeProgram: URIO[zio.ZEnv, ExitCode] =
  for {
    lead <- makeMouseLead
    exitCode <- catProgram(lead)
      .raceFirst(mouseProgram(lead))
      .catchAll(exc => putStrLn(s"Cat was ${-exc.catLead} cm ahead when it lost the mouse."))
      .as(ExitCode.success)
  } yield exitCode
```

Check out [the complete code](https://scastie.scala-lang.org/bMhWYmaVSpGOQePmAKPHaQ
).

Click "Run" a few times to see all possible outcomes.

## Explanation of `compositeProgram`

Our idea to chase types is a very small for comprehension. In the first place we map over the `makeMouseLead` program to create the variable. Secondly we run both the mouse and the cat with `.raceFirst` which will concurrently spin them up and wait for the first to terminate either with a success or a failure. We know that the mouse program doesn't terminate so we are basically waiting the cat program. This however leaves us with a `ZIO[ZEnv, MouseEscapedException, Unit]`.

To turn the error into a `Nothing` we use `.catchAll` and we pass an infallible program: `putStrLn`. To turn the `Unit` into an exit code we could simply `.flatMap` it to an `ExitCode` but when we are not interested in the result we can use `.as` which is a `.flatMap` that ignores the preceding result.

All these transformations are summarized in the below diagram which now looks complete.

![connected diagram](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/6v8z8h374kcrgk43ksi2.png)

## Bonus section

But wait, there is more. Instead of using the `.catchAll` and `.as(ExitCode.success)` functions, ZIO gives us a nice shortcut that does something very similar and that will work in the majority of cases. You can just add `.exitCode` and this will either print an error (and exit with 1) or succeed (and exit with 0). Not precisely our behavior, we lose the nice error message and we exit with an error when the mouse escape, but good enough: our program is simpler and we save one line.
```scala
val compositeProgram: URIO[zio.ZEnv, ExitCode] =
    for {
      lead <- makeMouseLead
      exitCode <- catProgram(lead)
        .raceFirst(mouseProgram(lead))
        .exitCode
    } yield exitCode
```