Purify code using free monads

Experienced Haskell programmers commonly advise newcomers to keep as much of their program as pure as possible. Purity confers many practical advantages:

• You can formally prove things about your code
• Barring that, you can reason easily about your code
• If all else fails, you can QuickCheck your code

To demonstrate this, I'll use the following simple echo program:

import System.Exit

main = do x <- getLine
          putStrLn x
          exitSuccess
          putStrLn "Finished"

The above program only makes one mistake: It mixes business logic with side-effects. Now, there's nothing wrong with that and I program like that all the time for simple programs that I can fit in my head. However, I hope to interest you in all the cool things you can do by factoring out the side-effects from the program logic.

Free monads

First, we must purify our program, and we can always factor out impure parts from any code by using free monads. Free monads let you decompose any impure program into a pure representation of its behavior and a minimal impure interpreter:

import Control.Monad.Free
import System.Exit hiding (ExitSuccess)

data TeletypeF x
  = PutStrLn String x
  | GetLine (String -> x)
  | ExitSuccess

instance Functor TeletypeF where
    fmap f (PutStrLn str x) = PutStrLn str (f x)
    fmap f (GetLine      k) = GetLine (f . k)
    fmap f  ExitSuccess     = ExitSuccess

type Teletype = Free TeletypeF

putStrLn' :: String -> Teletype ()
putStrLn' str = liftF $ PutStrLn str ()

getLine' :: Teletype String
getLine' = liftF $ GetLine id

exitSuccess' :: Teletype r
exitSuccess' = liftF ExitSuccess

run :: Teletype r -> IO r
run (Pure r) = return r
run (Free (PutStrLn str t)) = putStrLn str >>  run t
run (Free (GetLine  f    )) = getLine      >>= run . f
run (Free  ExitSuccess    ) = exitSuccess

echo :: Teletype ()
echo = do str <- getLine'
          putStrLn' str
          exitSuccess'
          putStrLn' "Finished"

main = run echo

The above rewrite concentrates all of the impure code within the run function. I like to call this process "purifying the code" because we distill out all of the program's logic into pure code leaving behind only the bare minimum impure residue in our run interpreter.

Proofs

Now, it seems like we didn't gain much from purifying our code and we paid a price in code size. However, now we can prove things about our code using equational reasoning.

For example, everybody reading this probably noticed the obvious bug in the original echo program:

import System.Exit

main = do x <- getLine
          putStrLn x
          exitSuccess
          putStrLn "Finished" <-- oops!

The last command never executes ... or does it? How would we prove that?

Actually, we can't prove that because it's not true. The author of the System.Exit module could simply change the definition of exitSuccess to:

exitSuccess :: IO ()
exitSuccess = return ()

The above program would still type-check, but now it would also successfully print Finished to the console.

However, for our purified version, we can prove that any command after exitSuccess' never executes:

exitSuccess' >> m = exitSuccess'

We do so using "equational reasoning", which is the most important benefit of purity. "Equational reasoning" means that we can take any expression and just substitute in the function definitions of the components. Each step of the following proof uses a comment to indicate the specific function definition used to reach the next step:

exitSuccess' >> m

-- exitSuccess' = liftF ExitSuccess
= liftF ExitSuccess >> m

-- m >> m' = m >>= \_ -> m'
= liftF ExitSuccess >>= \_ -> m

-- liftF f = Free (fmap Pure f)
= Free (fmap Pure ExitSuccess) >>= \_ -> m

-- fmap f ExitSuccess = ExitSuccess
= Free ExitSuccess >>= \_ -> m

-- Free m >>= f = Free (fmap (>>= f) m)
= Free (fmap (>>= \_ -> m) ExitSuccess)

-- fmap f ExitSuccess = ExitSuccess
= Free ExitSuccess

-- fmap f ExitSuccess = ExitSuccess
= Free (fmap Pure ExitSuccess)

-- liftF f = Free (fmap Pure f)
= liftF ExitSuccess

-- exitSuccess' = liftF ExitSuccess
= exitSuccess'

Notice how in the last steps we reversed the equality and worked backwards from the function definition to the defined expression. This is perfectly legitimate because thanks to purity the equals sign in Haskell doesn't mean assignment: it actually means equality! This means we can translate both ways across an equals sign when equationally reasoning about code. That's pretty amazing!

Equally important, the above proof is true no matter how the free monad is interpreted. We could swap out the run function for any other interpreter, including a pure interpreter, and the equation still holds:

exitSuccess' >> m = exitSuccess'

This means that we could safely use this as a GHC rewrite rule for an optimization pass over our program:

{-# RULES
  "exit" forall m. exitSuccess' >> m = exitSuccess'
    #-}

...and we can guarantee that the rule is always safe to apply and will never be broken.

Reasoning

We'd like to prove that our program always outputs the same string it received as input. Unfortunately, we can't prove this because it's not true. The maintainer of the putStrLn function could always have a change of heart and redefine it to:

putStrLn str = return () -- worst maintainer, ever

In fact, we can't even prove this is true for our free monad version, either. run uses the same putStrLn, so it would suffer from the same bug. However, we can still prove things about the free monad itself by instead studying it through a pure interpreter:

runPure :: Teletype r -> [String] -> [String]
runPure (Pure r)                  xs  = []
runPure (Free (PutStrLn y  t))    xs  = y:runPure t xs
runPure (Free (GetLine     k))    []  = []
runPure (Free (GetLine     k)) (x:xs) = runPure (k x) xs
runPure (Free  ExitSuccess   )    xs  = []

Now, we can prove that:

runPure echo = take 1

... based off of the Haskell98 Prelude's take. I leave this one as an exercise for the reader, because this post is already pretty long.

Notice that by examining echo through a pure lens, we caught a small corner case we didn't originally anticipate: The user might not enter any input! In that scenario our interpreter must return an empty list, just like take. Equational reasoning forces us to be rigorous in a way that simply saying "our program always outputs the same string it received" can never be. The more you work through these kinds of equations, the more you improve your ability to reason about your own code and catch mistakes in your own assumptions.

Equally important, equational transformations let you view your program in a whole new light. We saw our program was a glorified take 1 when filtered through runPure, meaning we could leverage our intuition about take for understanding our program and catching that corner case.

Once you factor your code through the Free monad and prove its soundness it acts like a trusted kernel and then you need only maintain the interpreter from that point forward. So while we couldn't quite prove the entire program was correct, we were able to prove that everything except the interpreter is correct. Equally important, we reduced the interpreter to the absolute minimal attack surface possible so that we can fit it in our head and maintain it by hand.

Testing

We can't prove anything about code in the IO monad. How would we even do that? We could say something like: "If you compile this and run the program and type some string into the prompt, you will get the same string back", but that's not really an equation, so there's nothing rigorous about it. However, we could write it as a test for our program.

Unfortunately, tests for impure code don't really scale to large and complicated programs and in test-dependent programming languages writing these tests consumes a significant fraction of the programmer's time.

Fortunately, though, we can easily exercise pure code with the QuickCheck library, which pathologically scours pure assertions for a violation in a completely automated fashion.

For example, let's assume that you were too lazy to prove that runPure echo = take 1. We can instead ask QuickCheck to test that proposition for us:

>>> import Test.QuickCheck
>>> quickCheck (\xs -> runPure echo xs == take 1 xs)
+++ OK, passed 100 tests.

How cool is that! You can test your code much more aggressively if you keep it as pure as possible.

Conclusions

Equational reasoning only works for pure code, so the pure component of your code base serves as a trusted kernel for which you can:

• prove soundness,
• reason about behavior, and
• aggressively test.

This is why we always strive to maximize the pure portions of our code bases and minimize the impure parts.

Fortunately, the Free monad guarantees that you can always easily achieve the absolute maximal level of purity possible and the absolute minimal amount of impure code. This is why every experienced Haskell programmer should be fluent in free monads.