Breadth first tree relabelling, or: the power of lazyness

[haanss]

Back in 2018 or so, I attended the Advanced Functional Programming summer school at Utrecht University. Doaitse Swierstra was one of the lecturers there and showed this very nifty example of laziness that totally expanded my mind at the time. Last week I came back to the AFP summer school, but as a TA this time, and because, as far as I could tell, none of the lecturers showed this example, I took it upon myself to show it. Googling the specific example, I can find no real reference to it online, so I’ve also taken it upon myself to share it here [0].

With the preamble done, here is the actual problem: Traverse a tree data structure in a breadth-first manner, and replace all the values with their “index”.

data Tree a
  = Leaf a 
  | Node (Tree a) a (Tree a)
  deriving (Eq, Show)

bfLabel :: Tree a -> Tree Int 
bfLabel t  = ???

But what is the correct LHS of the function here[1]? It’s not trivial at all to implement, and indeed, only one of the students in this year’s batch managed to fill it in at all. One thing we can look at to try and figure it out is the same problem, but depth first.

dfLabel :: Tree a -> Tree Int 
dfLabel = fst . go 0 where 
  go n (Leaf _)          = (Leaf n, n+1) 
  go n (Node l _ r)    =  let 
    (l', n')  = go (n+1) l
    (r', n'') = go n'   r
    in (Node l' n r', n'')

We need some counter to keep track of the highest label we’ve applied so far, and we pass this counter down to a node’s children as a function argument and pass it up to a node’s (or leaf’s) parent as the return argument. This works fine because at each node I really only need 1 value, and I really only need to send 1 value.

Visually, if I have the following tree (bad ms paint incoming):

I have the following counter-path:

but if I tried to use just one counter in a BF case I’d have the following counter path:

generalizing from this specific tree to a more general tree I have, I need to receive arbitrary counters at the top level (from the other leg of the tree) and receive one less counter as I descend down the tree. The same pattern holds for counters I need to send up from the function call. We do have a data structure for arbitrary amounts of values, lists! Except getting the list to use as an argument is… tricky. You might notice that the list that we need to send is exactly the return list of the function, with the seed value 0 attached as its head.

bfLabel :: Tree a -> Tree Int 
bfLabel t = let (result, future) = go (0:future) t in result where
  go (n:ns) (Leaf a)     = (Leaf n, n+1:ns) 
  go (n:ns) (Node l a r) = let 
    (l', ns')  = go ns  l 
    (r', ns'') = go ns' r
    in (Node l' n r', n+1 : ns'')

And through the magic of laziness, this all functions as we’d want . This is one example of tying the knot. I like it because it’s elegant (look how it mirrors the df solution function! ), not trivial, and at the same time not so hard that it becomes completely impossible to understand (for me).

TLDR: Laziness is cool.

[0] I’m certain most “advanced” haskellers will be aware of the concept, though perhaps not the specific example.
[1] There are actually multiple solutions for the lhs of the function. It can be done using a queue or by splitting the tree into a list of lists, where each list contains the values of a specific ‘level’, and then constructing the correct tree back with that list. The one student that managed to fill it in did it using a queue approach; the other TA besides me did it with the list of list approach.

[jaror]

I was there too in 2017: https://afp-2017.github.io/. Sadly it was one of Doaitse’s last summer schools.

[jackdk]

Extremely cool. I find that I don’t explicitly use laziness especially often in my practical work, with the exception of nixpkgs-style module systems. (Of course, there’s all the implicit laziness in libraries, where-clauses, etc., that we all rely upon.)

@haanss

dfLabel :: Tree a -> Tree Int 
dfLabel = fst . go 0 where 
  go n (Leaf _)          = (Leaf n, n+1) 
  go n (Node l _ r)    =  let 
    (l', n')  = go (n+1) l
    (r', n'') = go n''   r
    in (Node l' n r', n'')

Should that second go n'' r be go n' r?

[haanss]

Should that second go n'' r be go n' r?

Yep. I spend all my time making sure the bf version was correct and neglected to double check the df version . Silly mistake.

[jaror]

@haanss

I can find no real reference to it online

I think the first to publish it were Jones and Gibbons: https://researchspace.auckland.ac.nz/handle/2292/3470. That’s still from the Squiggol days, so prepare for symbol soup.

There’s a later paper by Okasaki that discusses several different approaches that don’t need laziness, but it does include Jones’ and Gibbons’ version in an appendix: https://dl.acm.org/doi/10.1145/357766.351253

[jaror]

As a tribute to Doaitse, here’s a version of the algorithm implemented using UUAGC, his attribute grammar system:

-- Main.ag

data Root
  | Root t :: Bin

data Bin
  | Leaf x :: Int
  | Bin  l :: Bin  r :: Bin
deriving Bin : Show

attr Root
  syn res :: Bin

attr Bin
  syn res :: Bin
  chn xss :: { [[Int]] }

sem Root
  | Root t.xss = { [1..] : @t.xss }

sem Bin
  | Leaf lhs.res = { Leaf @loc.x }
         (loc.x, lhs.xss) = { let ((x:xs):xss) = @lhs.xss in (x, (x+1:xs):xss) }
  | Bin l.xss = { tail @lhs.xss }
        lhs.xss = { head @lhs.xss : @r.xss }
        lhs.res = { Bin @l.res @r.res }

{
bfl :: Bin -> Bin
bfl expr = res_Syn_Root (wrap_Root (sem_Root (Root expr)) Inh_Root)

main :: IO ()
main = print (bfl (Bin (Bin (Leaf 0) (Bin (Leaf 0) (Leaf 0))) (Bin (Leaf 0) (Leaf 0))))
}

You can run it like this (needs GHC 9.6 or earlier):

$ cabal install uuagc
$ uuagc --haskellsyntax --module Main --catas --semfuns --signatures --wrappers --data Main.ag
$ ghc Main.hs
$ ./Main
Bin (Bin (Leaf 1) (Bin (Leaf 4) (Leaf 5))) (Bin (Leaf 2) (Leaf 3))

If you want to explore more of UUAGC, I have an example repo and check out the uuagc manual in the web archive.

[Iceland_jack]

Phases in Software Architecture https://www.cs.ox.ac.uk/jeremy.gibbons/publications/phases.pdf

[rhendric]

That paper also contains a breathtakingly beautiful breadth-first traversal for rose trees. Here it is adapted to the binary tree used in the OP:

import Control.Applicative (liftA2)
import Control.Arrow ((&&&))
import Control.Monad.State (MonadState, evalState, state)

-- This is generic infrastructure; Phases is isomorphic to the free
-- applicative, but the Applicative (Phases f) instance is *not* the free
-- applicative instance. If you squint, you can see that this resembles a
-- list, with Pure playing the role of [] and Link serving as (:). Much
-- like the OP used lists to track the next label at each level of the
-- tree, this structure is used to build an applicative effect at each
-- level of the tree, keeping them separate from each other.
data Phases f a where
  Pure :: a -> Phases f a
  Link :: (a -> b -> c) -> f a -> Phases f b -> Phases f c

deriving instance Functor (Phases f)

-- This instance is like zipping two lists. The interesting case is the
-- combining of two Links, which pairwise combines the effects in their
-- ‘heads’ and ‘tails’ without letting them interfere with each other. This
-- is the essence of the breadth-first traversal: first do all the effects
-- at a given level, and wait to do the effects at the level below that.
--
-- (Edited to add lazy patterns in h; this allows for fully lazy
-- traversals of infinite or partial trees.)
instance Applicative f => Applicative (Phases f) where
  pure = Pure
  Pure f <*> xs = fmap f xs
  fs <*> Pure x = fmap ($ x) fs
  Link f xs ys <*> Link g zs ws =
    Link h (liftA2 (,) xs zs) (liftA2 (,) ys ws)
    where h ~(x, z) ~(y, w) = f x y (g z w)

-- now places an effect in the ‘head’ of a Link.
now :: Applicative f => f a -> Phases f a
now xs = Link const xs (Pure ())

-- later places an effect in the ‘tail’ of a Link.
later :: Applicative f => Phases f a -> Phases f a
later xs = Link (const id) (pure ()) xs

-- runPhases is just the retraction of the free applicative structure into
-- a concrete Applicative. In practical terms, it runs the effects in the
-- ‘list’ in order, top to bottom.
runPhases :: Applicative f => Phases f a -> f a
runPhases (Pure x) = pure x
runPhases (Link f xs ys) = liftA2 f xs (runPhases ys)


-- Now, for a given tree type, we can define its breadth-first traversal
-- using the above primitives.

data Tree a
  = Leaf a
  | Node (Tree a) a (Tree a)
  deriving (Eq, Show)

-- Just look at this beauty.
bfTraverse :: Applicative f => (a -> f b) -> Tree a -> f (Tree b)
bfTraverse f = runPhases . go
  where
  go (Leaf a)     = Leaf <$> now (f a)
  go (Node l a r) = Node <$> later (go l) <*> now (f a) <*> later (go r)

-- To use the traversal to relabel a tree, we'll use the State monad.

-- fresh gets a fresh label from the state and increments it.
fresh :: MonadState Int m => m Int
fresh = state (id &&& succ)

-- And all that remains is to use it.
bfLabel :: Tree a -> Tree Int
bfLabel = flip evalState 0 . bfTraverse (const fresh)

[olf]

The way that bfTraverse combines the Phases effects is like the ordinary Traversable instance (that yields depth-first), only that one distinguishes between the base and recursive fields of constructors. I think this can be generalized to fixed points of arbitrary bitraversable bifunctors.

data BiTree now later = BiLeaf now | BiNode later now later
-- Tree a  ~ Fix (BiTree a)

[rhendric]

Interesting; I wonder how that would compare to the TreeLike abstraction used in the library that was cited in the paper?

[jeremygibbons]

There’s also a longer version, Breadth First Traversal Via Staging

[jaror]

For an application of that technique to type checkers, see: Casper Bach Poulsen - A Monadic Framework for Name Resolution in Multi-Phased Type Checkers

I believe we could avoid .hs-boot files if we type checked Haskell using that approach.

[meooow]

This is undoubtedly elegant, but it is also less efficient: O(n log n) compared to OP’s O(n).
A bit surprising that this is never mentioned in the paper.

Edit: I was mistaken, I believe it is O(n).

[jaror]

Can you elaborate? The running time is really not obvious to me. Where does the log n factor come from?

Edit: The extended paper does say:

Anyway, we make no claims that these transformations improve running time
or space usage. More broadly, we have not concerned ourselves with making these
traversals take linear time; for example, bf in Section 3 is not linear, because of
repeated concatenations of lists. This issue can be addressed by using difference
lists, and more generally by Cayley representations [12], but is orthogonal to our
main argument.

[meooow]

I corrected my statement, it seems to be O(n). For every node, the bulk of the work is done in <*> for Phases, which takes time proportional to the length of the smaller Phases. The length of a Phases structure returned from a node is equal to the height of the subtree rooted at that node. So for a node, the work done in <*> is proportional to the minimum of the heights of its two children. The total work is the sum of this minimum-child-height over all nodes. The sum of heights over all nodes can be shown to be O(n) (such as here) (and isn’t O(n log n) as I initially thought), so our complexity is also O(n).

The extended paper does say:…

Ah thanks, they did address it!

[rhendric]

One of the authors has written about a modified O(n) implementation that wraps the Phases structure in a difference-list-like thing; empirically, this does seem to improve performance, even if just Phases is already O(n) (I’m not sure who to believe on this point!). It complicates the presentation slightly, but the modified now and later combinators are used in exactly the same way in the resulting traversal function, so it ends up looking just as pretty at the surface level.

(The wrapper type is called Day, though I don’t understand the connection, if any, between that and the Day convolution type I was previously familiar with. Something something Cayley transforms—that part went right over my head, sadly.)

[meooow]

Very interesting! I don’t think it helps the complexity, but it does reassociate the <*>s for f which could improve things depending on f.

To make it easier to analyze, suppose we calculate the sum of node values at each level of the tree. This is like setting f to Const (Sum Int).

-- Structure when using Phases
levelSum :: Tree Int -> [Int]
levelSum (Leaf x) = [x]
levelSum (Node l x r) = x : lzw (+) (levelSum l) (levelSum r)
  where
    -- "long zip with" (from the paper)
    lzw :: (a -> a -> a) -> [a] -> [a] -> [a]
    lzw _ [] ys = ys
    lzw _ xs [] = xs
    lzw f (x:xs) (y:ys) = f x y : lzw f xs ys

-- Structure when using Day of Phases
-- This happens to be quite similar to OP!
levelSum2 :: Tree Int -> [Int]
levelSum2 t = go t []
  where
    go (Leaf x) = \case
      [] -> [x]
      y:ys -> x+y : ys
    go (Node l x r) = \case
      [] -> x : go l (go r [])
      y:ys -> x+y : go l (go r ys)

Both are O(n).

However, if we want levels :: Tree a -> [[a]], that means f ~ Const [a] and the second option is better because the <*>s for Const [a] and consequently the <>s for [a] would be associated to the right.

[olf]

@rhendric

I wonder how that would compare to the TreeLike abstraction used in the library that was cited in the paper?

That is a strange type class. At a glance, its definition is so circular that is seems independent of actual tree-like structure. For example, any Traversable could be considered TreeLike by defining treeTraverse f _ = traverse f. (Edit: The module authors acknowledge this in their Flat newtype.) More to the intention of the class, we can write

newtype BiFix b a = Fix {unFix :: b a (BiFix b a)}
instance Bifunctor b => Functor (BiFix b) where
    fmap f (Fix b) = Fix (bimap f (fmap f) b)
instance Bitraversable b => TreeLike (BiFix b) where
    treeTraverse f go = fmap Fix . bitraverse f go . unFix

This suggests that the TreeLike class captures the essence of a tying-the-knot bitraversal:

traverseFix :: (Bitraversable b, Applicative f) => (x -> f y) -> BiFix b x -> f (BiFix b y)
traverseFix f = go where
    go = fmap Fix . bitraverse f go . unFix

[Iceland_jack]

I also realize Phrases is more akin to your suggestion of PadList (from the excellent Constructing Applicative Functors paper): Using a “singleton” for pure rather than a repeated list like ZipList, and then padding the shorter list to match the longer one.

[jeremygibbons]

As @meooow says, folding over a tree with longZipWith f, where f takes constant time and longZipWith takes time proportional to the length of the shorter list, takes linear time overall. You can think of it colouring in the gaps in the picture of the tree. (Indeed, I used this property, and this intuition, in my paper on drawing trees). But that argument doesn’t apply when f is not constant time - for example, list append. I think this is where Oisin’s Cayley transformation comes in.