One of the things I love about Scala is how you can write things in so many different ways. Part of the fun is finding the best-looking way to accomplish a task. In this post, I take a big, ugly method with several logical branches and levels of nested conditions, and try to make it as pretty as I can. My quest brings me through call-by-name blocks all the way into some custom functional data structures and ‘for comprehensions’, leading to what I think is a pretty successful ‘beautification’.

The Scenario

Let’s say you’re writing a RESTful web server for… something. In each request you might do things like check the permissions, extract parameters from the request, interpret JSON values from the parameters, or look up values from a database. With each of these operations, there’s a possible failure point; if the permission check fails, you’d want to respond with a 403 Forbidden response. If parameters are missing, you might want to respond with a 400 Bad Request. Handling all of these cases and their possible failures is not pretty… so let’s give it a code-makeover!

Here’s how all of that might look at first:

def respondWithChecks(req: Request): Response = {

    // 1st logical branch: permission check
    if(checkForPermissions(req, someKindOfPermission)){

        // 2nd logical branch: parsing the request body as JSON
        extractJsonBody(req) match {
            case None => Response(400, "Request needs a JSON body")
            case Some(json) =>

                // 3rd logical branch: extracting params from the JSON
                extractJsonField[Int]("userId", json) match {
                    case None => Response(400, "Request body needs a 'userId' field")
                    case Some(userId) =>

                        // 4th logical branch: database lookup
                        lookupUserFromDB(userId) match {
                            case None => Response(404, "User not found")
                            case Some(user) => Response(200, user.toJson)
                        }
                }
        }
    } else {
        Response(403, "Insufficient Permissions")
    }
}

Assume these methods are already implemented somewhere.

def checkForPermissions(req: Request, perms: Permissions): Boolean

def extractJsonBody(req: Request): Option[JSON]

def extractJsonField[T](fieldName: String, json: JSON): Option[T]

def lookupUserFromDB(userId: Int): Option[User]

Call-by-Name Blocks

You don’t even have to read the respondWithChecks method above to see that it gets pretty heavily indented (by 7 tabs at the deepest point)! It also feels pretty Java-ish; the only Scala-ish thing in the method is the use of Options and match blocks.

I can eliminate a few levels of indentation by creating some ‘call-by-name’ style versions of the convenience functions. Here’s the general pattern:

def exampleRequestHandler(someParams: Params)(
    onSuccess: SomeValue => Response): Response = {

    if(someParams.areValid) {
        val someValue = calculateSomeValue(someParams)
        onSuccess(someValue)
    } else {
        someFailureResponse
    }
}

In the example above, the function takes some arbitrary input parameters and performs some logical check. If the check passes, it sends a value to the onSuccess function in order to return a Response. If the check fails, it returns some default failure response. Here’s the lookupUserFromDB method, wrapped to use this new style:

def withDBUser(userId: Int)(onSuccess: User => Response): Response = {
    lookupUserFromDB(userId) match {
        case Some(user) => onSuccess(user)
        case None => Response(404, "User not found")
    }
}

First Attempt

Assuming the other three convenience methods get refactored in a similar way, the original request handler method can be rewritten as follows:

def respondWithChecks(req: Request): Response = {
    withPermissions(req, someKindOfPermission){ () =>
        withJsonBody(req){ json =>
            withJsonField[Int]("userId", json){ userId =>
                withDBUser(userId){ user =>
                    Response(200, user.toJson)
                }
            }
        }
    }
}

This seems a little bit better. The deepest level of indentation is now 5 tabs (an improvement of two). Also, the error-condition handling has been extracted into our convenience methods (e.g. withDBUser giving the 404 response). That makes it easier to reuse the same blocks and not have to repeat that logic for each request handler. Yes, it’s an improvement, but we can do better!

Sequence Comprehensions

I spent some time thinking about this one, and arrived at the conclusion that the most effective way to de-indent this code would be to take advantage of Scala’s “sequence comprehensions” syntax sugar. (Also known as “for comprehensions”. I tend to prefer that way despite the fact that it’s more awkward to write/read/say/hear.)

Scala’s has these comprehensions baked into most of its standard library. Lists, Options, Trys, Futures, and pretty much every collection have the map, flatMap, withFilter, and foreach methods defined on them. These four methods are what Scala’s compiler desugars comprehensions into:

for {
    i <- listOfInts
    j <- anotherList
} yield i * j

// is equivalent to

listOfInts.flatMap { i =>
    anotherList.map { j =>
        i * j
    }
}

If sequence comprehensions are new to you, I encourage you to do a bit of reading on them before continuing down this post. I’m about to dive into the deep end of this concept.

The End Goal

Now here’s the leap I’m making: listOfInts.flatMap { i => ... } looks a lot like withJsonBody(req){ json => ... } in its format; maybe I can refactor withJsonBody to return something that has a flatMap and map method so that I can use it in a sequence comprehension, like this:

for {
    _ <- withPermissions(req, someKindOfPermission)
    json <- withJsonBody(req)
    userId <- withJsonField[Int]("userId", json)
    user <- withDBUser(userId)
} yield Response(200, user.toJson)

I don’t think you could possibly have less indentation without completely inlining the code. It’s extremely concise; there’s almost no syntactical overhead to this style. It feels pretty readable to me; most of the important bits are indented on the same level, so it reads more like regular text. The values being returned by each step are clearly listed on the left (json, userId, and user). There aren’t a lot of extra brackets and parentheses. Seems nice to me, now let’s make it happen!

The Calculation Data Structure

To make the comprehensions-based approach work, we need to design a data structure that fits its needs:

Each step can either short-circuit to a return value immediately, or succeed by passing along the input to the next step
Must have a flatMap and map method in order to fit a for comprehension.

For our purposes, the end result type should always be the same. All of our convenience methods eventually return a Response. The value they pass along can vary. Generically, let’s refer to the result type as Result, and the passed-along input type as B.

sealed trait Calculation[+Result, +B]

case class Done[Result](value: Result) extends Calculation[Result, Nothing]
case class Next[B](input: B) extends Calculation[Nothing, B]

The Calculation ADT is a lot like Scala’s Either; the calculation can either be Done, containing a result value, or be ready to pass along the next input to another calculation.

The + in Calculation[+Result, +B] marks Result and B as covariant types. Note that Done extends Calculation[Result, Nothing]. Nothing is a magical class that is a subtype of everything. When combined with being covariant, it means that a Next[B] could count as a Calculation[String, B], or a Calculation[Response, B], or a Calculation[PrettyMuchAnything, B]. This is the same concept as Nil for Scala’s Lists, or None for Options.

With the Calculation data structure in mind, let’s refactor the convenience methods a second time:

def withDBUser(userId: Int): Calculation[Response, User] = {
    lookupUserFromDB(userId) match {
        case Some(user) => Next(user)
        case None => Done(Response(404, "User not found"))
    }
}

This version of the function feels pretty similar to the previous one. The arguments to each function get simpler since there is no second argument list, and no onSuccess function in the signature. Rather than calling onSuccess, it simply returns a Next containing the user. If it fails, it returns a Done with the error response, since the calculation is done at that point. The rest of the convenience methods go pretty much the same way:

def withPermissions(req: Request, perms: Permissions): Calculation[Response, Unit]

def withJsonBody(req: Request): Calculation[Response, JSON]

def withJsonField[T](fieldName: String, json: JSON): Calculation[Response, T]

def withDBUser(userId: Int): Calculation[Response, User]

Note that the Result type for each of the four convenience methods is Response, because the end goal of the calculation is a Response. The B type is up to the individual methods.

Calculation Mapping

The final piece of the puzzle is to implement a flatMap and map method for the Calculation class. It turns out to be easier than you might think. The hardest part is figuring out how to deal with the consequences of making Result and B be covariant.

For flatMap, the functionality can be read as “If the calculation is done, it stays done. If it’s not done, I should pass my input to the next part of the calculation.”

def flatMap[R2 >: Result, B2](f: B => Calculation[R2, B2]):
    Calculation[R2, B2] = this match {

    case Done(result) => Done(result)
    case Next(input) => Next(f(input))
}

The “next part of the calculation” is represented as the function f, which takes an input of type B and returns the next state of the calculation. R2 is there to make the compiler happy with covariance. Basically it allows for f to return a Calculation with a less-specific Result type.

The map method is a little bit tricky. Since it will generally be called from inside a flatMap function, it needs to return a Calculation[Result, B] instead of a Result, even though what we want is the Result. For now, we can craft the type signature for map so that it forces the B type to be equal to the Result type. Later on we can define a method to help get the Result value out.

def map[R2 >: Result](f: B => R2): Calculation[R2, R2] = this match {
    case Done(result) => Done(result)
    case Next(input) => Done(f(input))
}

Again, R2 is there because the Result type is covariant. You can think of R2 as equivalent to Result if it helps. Now, to get the actual result value out…

def result[R2 >: Result](implicit eq: B <:< R2): R2 = this match {
    case Done(r) => r
    case Next(in) => eq(in)
}

Here, I used a cool trick where an implicit “subtype equality” checker method ensures at compile time that B is either equal to, or a subtype of R2, which itself is either Result, or a supertype of Result. (confused yet?) What this boils down to is that if you try calling result on a Calculation[Response, Unit], you’ll get a compile error, but if you try calling result on a Calculation[Response, Response], you’ll get a Response. Check out the scaladoc for <:<

The Final Result

Now that all of the pieces have been made, we can finally put them together into the new-and-improved respondWithChecks method!

def respondWithChecks(req: Request): Response = {
    val calculation = for {
        _ <- withPermissions(req, someKindOfPermission)
        json <- withJsonBody(req)
        userId <- withJsonField[Int]("userId", json)
        user <- withDBUser(userId)
    } yield Response(200, user.toJson)

    calculation.result
}

dylemma.io

Functional Beautification