Wednesday, July 18, 2012

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.

14 comments:

  1. Thanks for the walk through of how to prove something with equational reasoning. The term is thrown around a lot, but there are few examples of the actual steps.

    I have been thinking that it might be good to have a wiki for showing the proofs of instances for common type classes (Monad, Applicative, etc). This would be a good resource to show users how to prove their own instances.

    ReplyDelete
    Replies
    1. If you are interested in equational reasoning I recommend Richard Bird's Pearls of Functional Algorithm Design, where he tackles multiple problems using equational reasoning to arrive at efficient solutions.

      Delete
  2. This comment has been removed by the author.

    ReplyDelete
  3. We can do program derivation also! It's not easy but works nice. In this fashion one can make a program more efficient while it's sure that is correct.

    And there is another good point to equational reasioning. People often think that demonstrations only serve to confirm correctness, but they also serve to deduce answers too. And this is a very good point to have in programming, choosing an answer because it's guided by some properties instead of trying to guess it thinking in a human manner the way that the program should behave.

    ReplyDelete
    Replies
    1. I completely agree. Usually you know you implemented something correctly when it has all sorts of nice succinct equational properties.

      Also, like you said, you can use all the equational properties as rewrite rules so that GHC can optimize your program automatically, as I hinted at in the post.

      Delete
  4. Is the "free monad" the key to adding purity/semantics to IO?

    It seems like what you are doing is splitting IO into to two monads: a pure monad and an impure monad. This lets us swap out the impure "implementation" monad with alternate backends, while still maintaining some semantic guarantees.

    The "free monad" is the *least* you can do, so it's a good starting point.

    But it would be even nicer to use a non-free monad, one that lets the type-system encode *even more* semantics about the expected behaviors of the IO functions.

    When you say that we can reason about `runPure` to show that `runPure echo === take `, that shows you that the impure `run echo` *might* be ismorphic/adjoint/equivalent to `take 1`, but there's no guarantee, because the list-management code is inside `runPure`, not in the common (free) monad.

    Basically, the free monad and `runPure` provides a test case that `run` "isn't definitely wrong"

    Now (and here's where I start hand-waving frantically), if you boost the list-management code from `runPure` up to the pure monad (no longer a free monad), but carrying some additional semantics that says that `putStrLn` consumes content from `getLine` (and any intermediate pure transformations), then you have IO "laws". You still can't statically *guarantee* that impure IO follows those laws (same is with the monad laws in general), but you can encode them, and you can write specific test cases to (quick)check those laws on sample data/programs.

    (Disclosure: I don't really know what I am talking about. I'm not sure that using a more interesting pure monad actually carries tangible improvements)

    ReplyDelete
    Replies
    1. To answer your first question, the free monad is not the only way to add a useful semantics to `IO`. The approach you describe (moving some shared logic up into the free monad and making it non-free) is better.

      The main reason I use the free monad as the running example is because it is the most extreme and most instructive example: if you understand how the free monad trick works then you typically also understand how the non-free derivatives work as well.

      Delete
  5. Thanks for your article!
    However, I wonder how you would do this practically for a program which requires maybe tens or hundreds of IO operations. Wrapping all those functions into a free monad would be quite a burden.
    Is there a good approach to this?

    ReplyDelete
    Replies
    1. Yes, the `free` library recently added a Template Haskell function to autogenerate all of these wrapped functions:

      http://hackage.haskell.org/package/free-4.7.1/docs/Control-Monad-Free-TH.html

      Delete
  6. Excellent post - great examples. Thanks a lot!

    A quick question for you: How you would represent a function of the type `IO a -> IO (Double, a)` (say, System.TimeIt.timeItT) as a constructor of your TeletypeF type?

    Thanks again for the excellent post!

    ReplyDelete
    Replies
    1. I think it's probably easier to do this using the `operational` package's encoding of free monads:

      https://hackage.haskell.org/package/operational

      Delete
    2. Perfect. Thanks for pointing me in the right direction.

      Delete
  7. Instead of quickcheck, is there some automated-reasoning tool that would attempt to actually find the proof of the equational property automatically? I suppose it wouldn't work in the general case, but proving simple examples as the one presented here should be possible somehow, I think. Are there some theorethical issues or is it too difficult to implement or has simply nobody done it yet?

    ReplyDelete
    Replies
    1. I think Coq is what a lot of people reach for when they want automatic theorem proving although I'm not an expert on it myself

      Delete