Bird-style single pass minimum on trees with `fixIO`

Analysis.md

The trick

The following code implements a Bird-style traversal replaceMin which traverses a tree exactly once to produce a new tree of the same shape with the leaves replaced with the minimum of all leaf values in the original tree.

import System.IO (fixIO)

data Tree a = Leaf a | Branch (Tree a) (Tree a)
  deriving Show

f :: Tree Int -> Int -> IO (Int, Tree Int)
f (Leaf n) m = do
  print n
  return (n, Leaf m)
f (Branch t1 t2) m = do
  (m1, r1) <- f t1 m
  (m2, r2) <- f t2 m
  return (min m1 m2, Branch r1 r2)

replaceMin :: Tree Int -> IO (Int, Tree Int)
replaceMin t = fixIO (\ ~(m, r) -> f t m)

testTree :: Tree Int
testTree = Branch (Leaf 5) (Leaf 3)

main :: IO ()
main = do
  res <- replaceMin testTree
  print res

Running this code will produce:

5
3
(3, Branch (Leaf 3) (Leaf3))

Note that we can observer that the tree is traversed exactly once as the leaf case of f is effectful by virtue of printing the leaf value, and we see that each leaf is printed precisely once.

Analyzing the trick in general

The reason this code functions is due to a few things including:

• MVars
• unsafeDupablePerformIO
• laziness

MVars in GHC are mutable "boxes" that are either empty or full. They also have a blocking read operation. Thus, one can create an empty MVar and put it into a data structure with the expectation that in the future a value will be put into the MVar and then subsequently read in a non-blocking manner.

Next let's look at unsafeDupableIO.

unsafeDupablePerformIO  :: IO a -> a
unsafeDupablePerformIO (IO m) = case runRW# m of (# _, a #) -> lazy a

This definition bypasses the IO monad by running the IO action but pretending it is a pure value. Importantly, the result of the computation is lazy, through the use of the GHC magic lazy function. Thus, the result can be used lazily.

We now consider the definition of fixIO.

fixIO :: (a -> IO a) -> IO a
fixIO k = do
  v <- newEmptyMVar                                           -- 1
  ans <- unsafeDupableInterleaveIO                            -- 2
           (readMVar v `catch` \BlockedIndefinitelyOnMVar ->
                                  throwIO FixIOException)
  result <- k ans                                             -- 3
  putMVar v result                                            -- 4
  return result                                               -- 5

Reading one line at a time:

1. Creeate an empty MVar v, i.e. an empty box to be filled in the future.
2. Using this empty v, create an IO action ans which reads from it, and use unsafeDupableInterleaveIO to pretend it is a pure value. If GHC determines that this read will block, then error.
3. Run the given function k exactly once! The input to this function is ans, which looks like a pure value but is sercretly and IO action which reads from v. The function k should not inspect its argument, i.e. it should be lazy on its argument. If it is not, it will attempt to read from v before a value has been place into it and block forever, triggering a FixIOException.
4. Take the final result and fill the empty v with it.
5. Return the result.

Thus we can see in the implementation of fixIO that its argument k is run exactly once. Every place which the argument of k is used is initially an empty MVar when k is run, and eagerly/strictly attempting to access the argument will cause an error. When the argument is only used lazily however, the function k can succesfully terminate with result, which fixIO then places into v. Thus, when the result is demanded later and is lazily evaluated, there is now a value in v which can be read.

Analyzing replaceMin

Recall the definition of replaceMin:

f :: Tree Int -> Int -> IO (Int, Tree Int)
f (Leaf n) m = do
  print n
  return (n, Leaf m)
f (Branch t1 t2) m = do
  (m1, r1) <- f t1 m
  (m2, r2) <- f t2 m
  return (min m1 m2, Branch r1 r2)

replaceMin :: Tree Int -> IO (Int, Tree Int)
replaceMin t = fixIO (\p -> f t (fst p))

For clarity we have replaced the lazy pattern match ~(m, r) with equivalent code, making it clear that the argument is passed lazily into f. Our test tree is

testTree :: Tree Int
testTree = Branch (Leaf 5) (Leaf 3)

which looks like the following in memory

        _______
      /         \
-------------    ---------    ---------
| B | * | * |    | L | 5 |    | L | 3 |
-------------    ---------    ---------
          \__________________/

Let's now fun replaceMin by hand. We thus run fixIO, which begins by creating the empty MVar v off to the side and starts f consuming testTree at its Branch, which then immediately recursively calls f on Leaf 5. This prints 5 andcreates a new Leaf node (morally) pointing to m = fst v:

Old
        _______
      /         \  * f *
-------------    ---------    ---------
| B | * | * |    | L | 5 |    | L | 3 |
-------------    ---------    ---------
          \__________________/
New
                         ________________
                       /                  \
                 ---------                 -----
                 | L | * |                 | _ | m = fst v
                 ---------                 -----

Next, f is called on Leaf 3 which prints 3 and creates a new Leaf pointing to m = fst v:

Old
        _______
      /         \               * f *
-------------    ---------    ---------
| B | * | * |    | L | 5 |    | L | 3 |
-------------    ---------    ---------
          \__________________/
New
                         ________________
                       /                  \
                 ---------    ---------    -----
                 | L | * |    | L | * |    | _ | m = fst v
                 ---------    ---------    -----
                                    \_____/

Now that both recursive calls have finished, the minimums m1 and m2 are available as 3 and 5 respectively. f then computes the minimum (namely 3) and returns that along with a new Branch pointing to the new leaves:

Old
        _______
      /         \               * f *
-------------    ---------    ---------
| B | * | * |    | L | 5 |    | L | 3 |
-------------    ---------    ---------
          \__________________/
New 
        ________          ________________
      /         \       /                  \
-------------    ---------    ---------    -----
| B | * | * |    | L | * |    | L | * |    | _ | m = fst v
-------------    ---------    ---------    -----
          \__________________/      \_____/

The function f has now finished, and returns the pair (3, <new tree>). Note that as it currently stands the leaves of the new tree are currently pointing to an empty v, and so if we tried to consume the new tree now (by say printing it) we would attempt to read from an empty MVar which would block. Thus, the next line in fixIO takes the pair (3, <new tree>) and fills v with it. The first component of this is 3, and so we can update the new tree to be

        ________          ________________
      /         \       /                  \
-------------    ---------    ---------    -----
| B | * | * |    | L | * |    | L | * |    | 3 |
-------------    ---------    ---------    -----
          \__________________/      \_____/

Now when we want to use this <new tree>, i.e. print it, when we attempt to use the values stored in the leaves we read from the filled MVar v and thus observe the desired result.