osa1 github about atom

Non-local returns in functional programs

July 21, 2013 - Tagged as: en, haskell, ocaml.

Let’s say we want to get nth visited element in depth-first traversal of a tree. Doing this is almost too easy in a language with statements(all imperative languages, some functional ones): Just run the depth-first traversal algorithm with explicit stack, and use return when you visit nth node.

In an expression language(Haskell, some Lisp languages) this is somewhat tricky.

Since I almost always prefer simplest possible solution of a problem, this would be my first attempt in a real-world situation:

data Tree a
    = Branch [Tree a]
    | Leaf a

toList :: Tree a -> [a]
toList (Branch bs) = concatMap toList bs
toList (Leaf a)    = [a]

dfsNth :: Tree a -> Int -> Maybe a
dfsNth tree n = listToMaybe . drop n $ toList tree

One concern about this function may be that the complexity of list generation. It’s hard to predict complexity of this function, but traversing the whole tree just to get first element of it would be costly anyway.

But thanks to lazy evaluation, this function still not very bad. Because only required parts of the intermediate list will be generated. To see why you can do two things: 1) Just place some Debug.Trace.trace calls in toList function and see how many times a leaf node is visited and 2) evaluate this function by hand and observe unevaluated thunks.

Let’s just do the first one, since it’s easier:

toList :: Tree a -> [a]
toList (Branch bs) = concatMap toList bs
toList (Leaf a)    = trace "leaf node visited" [a]

...

testTree = Branch [ Branch [ Leaf 1, Leaf 2 ], Branch [ Leaf 3 ], Branch [ Branch [ Branch [ Leaf 4 ] ] ] ]
ghci> dfsNth testTree 0
leaf node visited
Just 1

ghci> dfsNth testTree 1
leaf node visited
leaf node visited
Just 2

ghci> dfsNth testTree 10
leaf node visited
leaf node visited
leaf node visited
leaf node visited
Nothing

Other solutions are still worth exploring. When I think of “returning in the middle of a function” in Haskell, I always think Either. It’s monad definition is a great fit for this:

instance Monad (Either e) where
    return        = Right
    Left  l >>= _ = Left l
    Right r >>= k = k r

So when Left data used in monadic bind(>>=), second parameter just ignored and Left data is returned. Just like returning in the middle of a function in imperative setting, by ignoring rest of statements.

Using monad instance of Either, we can easily implement our function:

dfsNth' :: Tree a -> Int -> Maybe a
dfsNth' tree n =
    case iter tree n of
      Left a  -> Just a
      Right i -> Nothing
  where
    iter :: Tree a -> Int -> Either a Int
    iter (Branch []) n     = return n
    iter (Branch (b:bs)) n = do
      n' <- iter b n
      iter (Branch bs) n'
    iter (Leaf a) 0 = Left a
    iter Leaf{}   n = return (n - 1)

It works exactly like our first implementation, but without generating an intermediate list.

Exceptions and Either monad

If I were using OCaml, I’d probably implement this function using an exception.

exception NonLocal of int

type 'a tree =
  | Branch of ('a tree) list
  | Leaf   of 'a

let dfs_nth tree n =
  let rec iter tree n =
    match tree with
    | Branch [] ->
        n
    | Branch (b :: bs) ->
        iter (Branch bs) (iter b n)
    | Leaf a ->
        if n = 0 then raise (NonLocal a) else n - 1
  in
  try
    iter tree n;
    None
  with
    | NonLocal a -> Some a

An interesting thing to realize here is that this solution is very similar to our Haskell solution. In Haskell, Either is an instance of MonadError:

instance Error e => MonadError e (Either e) where
    throwError             = Left
    Left  l `catchError` h = h l
    Right r `catchError` _ = Right r

This means if you replace Lefts with throwError (just like raise in OCaml code), you have a similar solution with OCaml.

This doesn’t mean exceptions are same thing as Either types in functional programming. There are just too many differences that I won’t delve into in this post. With an exception, you can return from arbitrary deep contexts(ie. function calls), which is not easily possible with Either types. This is why exceptions sometimes referred as non-local returns.

We discussed this stuff over OCaml IRC channel, and smart people over there gave me some really good insights about non-local returns and exceptions. I’ll probably delve into details in another blog post. I’m especially interested in functional solutions that we can have in Haskell.

For the curious, for now I’ll just leave these two links here: (I haven’t read that links yet, but they’re probably related)

Several other ideas also discussed at IRC channel, some of them were using delimited continuations, or passing a handler function as parameter and just calling it instead of raising an exception. I’ll continue investigating this stuff later.

I also came across this StackOverflow post that explains how Scala’s non-local returns implemented as exceptions internally. Interesting stuff.