Folding Two Things at Once

There’s a whole family of Haskell brainteasers surrounding one function: foldr. The general idea is to convert some function on lists which uses recursion into one that uses foldr. map, for instance:

map :: (a -> b) -> [a] -> [b]
map f = foldr (\e a -> f e : a) []

Some can get a little trickier. dropWhile, for instance. (See here and here for interesting articles on that one in particular.)

dropWhile :: (a -> Bool) -> [a] -> [a]
dropWhile p = fst . foldr f ([],[]) where
  f e ~(xs,ys) = (if p e then xs else zs, zs) where zs = e : ys

Zip

One function which was a little harder to convert than it first seemed was zip.

Here’s the first (non) solution:

zip :: [a] -> [b] -> [(a,b)]
zip = foldr f (const []) where
  f x xs (y:ys) = (x,y) : xs ys
  f _ _  [] = []

The problem with the above isn’t that it doesn’t work: it does. The problem is that it’s not really using foldr. It’s only using it on the first list: there’s still a manual uncons being performed on the second. Ideally, I would want the function to look something like this:

zip :: [a] -> [b] -> [(a,b)]
zip xs ys = foldr f (\_ _ -> []) xs (foldr g (const []) ys)

The best solution I found online only dealt with Folds, not Foldables. You can read it here.

Recursive Types

Reworking the solution online for Foldables, the initial intuition is to have the foldr on the ys produce a function which takes an element of the xs, and returns a function which takes an element of the xs, and so on, finally returning the created list. The problem with that approach is the types involved:

zip :: [a] -> [b] -> [(a,b)]
zip xs = foldr f (const []) xs . foldr g (\_ _ -> []) where
  g e2 r2 e1 r1 = (e1,e2) : (r1 r2)
  f e r x = x e r

You get the error: Occurs check: cannot construct the infinite type: t0 ~ a -> (t0 -> [(a, b)]) -> [(a, b)]. Haskell’s typechecker doesn’t allow for infinitely recursive types.

You’ll be familiar with this problem if you’ve ever tried to encode the Y-combinator, or if you’ve fiddled around with the recursion-schemes package. You might also be familiar with the solution: a newtype, encapsulating the recursion. In this case, the newtype looks very similar to the signature for foldr:

newtype RecFold a b = 
  RecFold { runRecFold :: a -> (RecFold a b -> b) -> b }

Now you can insert and remove the RecFold wrapper, helping the typechecker to understand the recursive types as it goes:

zip :: [a] -> [b] -> [(a,b)]
zip xs =
  foldr f (const []) xs . RecFold . foldr g (\_ _ -> []) where
    g e2 r2 e1 r1 = (e1,e2) : (r1 (RecFold r2))
    f e r x = runRecFold x e r

As an aside, the performance characteristics of the newtype wrapper are totally opaque to me. There may be significant improvements by using coerce from Data.Coerce, but I haven’t looked into it.

Generalised Zips

The immediate temptation from the function above is to generalise it. First to zipWith, obviously:

zipWith :: (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c xs =
  foldr f (const []) xs . RecFold . foldr g (\_ _ -> []) where
    g e2 r2 e1 r1 = c e1 e2 : (r1 (RecFold r2))
    f e r x = runRecFold x e r

What’s maybe a little more interesting, though, would be a foldr on two lists. Something which folds through both at once, using a supplied combining function:

foldr2 :: (Foldable f, Foldable g)
       => (a -> b -> c -> c)
       -> c -> f a -> g b -> c
foldr2 c i xs =
  foldr f (const i) xs . RecFold . foldr g (\_ _ -> i) where
    g e2 r2 e1 r1 = c e1 e2 (r1 (RecFold r2))
    f e r x = runRecFold x e r

Of course, once you can do two, you can do three:

foldr3 :: (Foldable f, Foldable g, Foldable h)
       => (a -> b -> c -> d -> d)
       -> d -> f a -> g b -> h c -> d
foldr3 c i xs ys =
  foldr f (const i) xs . RecFold . foldr2 g (\_ _ -> i) ys where
    g e2 e3 r2 e1 r1 = c e1 e2 e3 (r1 (RecFold r2))
    f e r x = runRecFold x e r

And so on.

There’s the added benefit that the above functions work on much more than just lists.

Catamorphisms

Getting a little formal about the above functions, a fold can be described as a catamorphism. This is a name for a pattern of breaking down some recursive structure. There’s a bunch of them in the recursion-schemes package. The question is, then: can you express the above as a kind of catamorphism? Initially, using the same techniques as before, you can:

newtype RecF f a = RecF { unRecF :: Base f (RecF f a -> a) -> a }

zipo :: (Functor.Foldable f, Functor.Foldable g)
     => (Base f (RecF g c) -> Base g (RecF g c -> c) -> c)
     -> f -> g -> c
zipo alg xs ys = cata (flip unRecF) ys (cata (RecF . alg) xs)

Then, coming full circle, you get a quite nice encoding of zip:

zip :: [a] -> [b] -> [(a,b)]
zip = zipo alg where
  alg Nil _ = []
  alg _ Nil = []
  alg (Cons x xs) (Cons y ys) = (x, y) : ys xs

However, the RecF is a little ugly. In fact, it’s possible to write the above without any recursive types, using the RankNTypes extension. (It’s possible that you could do the same with foldr2 as well, but I haven’t figured it out yet)

You can actually use a newtype that’s provided by the recursion-schemes library as-is. It’s Mu. This is required for an encoding of the Y-combinator. It’s usually presented in this form:

newtype Mu a = Roll { unroll :: Mu a -> a }

However, in the recursion-schemes package, its definition looks like this:

newtype Mu f = Mu (forall a. (f a -> a) -> a)

No recursion! The zipo combinator above can be written using Mu like so:

zipo :: (Functor.Foldable f, Functor.Foldable g)
     => (Base f (Mu (Base g) -> c) -> Base g (Mu (Base g)) -> c)
     -> f -> g -> c
zipo alg xs = cata (\x -> alg x . project) xs . refix

And the new version of zip has a slightly more natural order of arguments:

zip :: [a] -> [b] -> [(a,b)]
zip = zipo alg where
  alg Nil _ = []
  alg _ Nil = []
  alg (Cons x xs) (Cons y ys) = (x,y) : xs ys

Zipping Into

There’s one more issue, though, that’s slightly tangential. A lot of the time, the attraction of rewriting functions using folds and catamorphisms is that the function becomes more general: it no longer is restricted to lists. For zip, however, there’s still a pesky list left in the signature:

zip :: (Foldable f, Foldable g) => f a -> g b -> [(a,b)]

It would be a little nicer to be able to zip through something preserving the structure of one of the things being zipped through. For no reason in particular, let’s assume we’ll preserve the structure of the first argument. The function will have to account for the second argument running out before the first, though. A Maybe can account for that:

zipInto :: (Foldable f, Foldable g) 
        => (a -> Maybe b -> c) 
        -> f a -> g b -> f c

If the second argument runs out, Nothing will be passed to the combining function.

It’s clear that this isn’t a fold over the first argument, it’s a traversal. A first go at the function uses the state monad, but restricts the second argument to a list:

zipInto :: Traversable f => (a -> Maybe b -> c) -> f a -> [b] -> f c
zipInto c xs ys = evalState (traverse f xs) ys where
  f x = do
    h <- gets uncons
    case h of 
      Just (y,t) -> do 
        put t
        pure (c x (Just y))
      Nothing -> pure (c x Nothing)

That code can be cleaned up a little:

zipInto :: Traversable f => (a -> Maybe b -> c) -> f a -> [b] -> f c 
zipInto c = evalState . traverse (state . f . c) where
  f x [] = (x Nothing, [])
  f x (y:ys) = (x (Just y), ys)

But really, the uncons needs to go. Another newtype wrapper is needed, and here’s the end result:

newtype RecAccu a b =
  RecAccu { runRecAccu :: a -> (RecAccu a b, b) }
  
zipInto :: (Traversable t, Foldable f)
        => (a -> Maybe b -> c) -> t a -> f b -> t c
zipInto f xs =
  snd . flip (mapAccumL runRecAccu) xs . RecAccu . foldr h i where
    i e = (RecAccu i, f e Nothing)
    h e2 a e1 = (RecAccu a, f e1 (Just e2))
Advertisements

A Trie in Haskell

Basic Ops

A Trie is one of those data structures that I find myself writing very early on in almost every language I try to learn. It’s elegant and interesting, and easy enough to implement.

I usually write a version that is a set-like data structure, rather than a mapping type, for simplicity’s sake. It stores sequences, in a prefix-tree structure. It has a map (dictionary) where the keys are the first element of every sequence it stores, and the values are the Tries which store the rest of the sequence. It also has a boolean tag, representing whether or not the current Trie is a Trie on which a sequence ends. Here’s what the type looks like in Haskell:

import qualified Data.Map.Lazy as M

data Trie a = Trie { endHere :: Bool
                   , getTrie :: M.Map a (Trie a)
                   } deriving (Eq)

Now, inserting into the Trie is easy. You just uncons on a list, and insert the head into the map, with the value being the tail inserted into whatever existed at that key before:

empty :: Trie a
empty = Trie False M.empty
               
insert :: Ord a => [a] -> Trie a -> Trie a
insert [] (Trie _ m)     = Trie True m
insert (x:xs) (Trie e m) = Trie e (M.alter (Just . insert xs . fromMaybe empty) x m)

Searching is simple, also. For the empty list, you just check if the Trie has its endHere tag set to True, otherwise, you uncons, search the map, and query the Trie with the tail if it eas found, or just return False if it was not:

member :: Ord a => [a] -> Trie a -> Bool
member [] (Trie e _)     = e
member (x:xs) (Trie _ m) = fromMaybe False (member xs <$> M.lookup x m)

Here’s my problem. Both of those functions have the same pattern:

f []     = ...
f (x:xs) = ...

Any good Haskeller should be begging for a fold at this stage. But it proved a little trickier than I’d imagined. Take member, for instance. You want to fold over a list, with the base case being the tag on the Trie:

member :: Ord a => [a] -> Trie a -> Bool
member = foldr f base where
  base = ???
  f e a = M.lookup e ???

Where do you get the base case from, though? You have to specify it from the beginning, but the variable you’re looking for is nested deeply into the Trie. How can you look into the Trie, without traversing the list, to find the tag, at the beginning of the function?

That had been my issue for a while. Every time I cam back to writing a Trie, I would see the pattern, try and write insert and member with a fold, and remember again the trouble I had had with it in the past. Recently, though, I saw a different problem, that gave me an idea for a solution.

The Highest Order

“Rewrite dropWhile using foldr

It’s a (semi) well-known puzzle, that’s maybe a little more difficult than it seems at first. Here, for instance, was my first attempt at it:

dropWhile' :: (a -> Bool) -> [a] -> [a]
dropWhile' p = foldr f [] where
  f e a | p e       = a
        | otherwise = e:a

Yeah. That’s filter, not dropWhile:

dropWhile' (<5) [1, 3, 6, 3, 1] -- [6]

Here was my final solution:

dropWhile' :: (a -> Bool) -> [a] -> [a]
dropWhile' p l = drop (foldr f 0 l) l where
  f e a | p e       = a + 1
        | otherwise = 0

After the problem I found this issue of The Monad Reader, which talks about the same problem. In my drop version, I had been counting the number of items to drop as I went, adding one for every element that passed the test. The corresponding version in the article had been building up tail functions, using . to add them together:

dropWhile' :: (a -> Bool) -> [a] -> [a]
dropWhile' p l = (foldr f id l) l where
  f e a | p e       = tail . a
        | otherwise = id

A quick visit to pointfree.io can generate some monadic pointsfree magic:

dropWhile' :: (a -> Bool) -> [a] -> [a]
dropWhile' p = join (foldr f id) where
  f e a | p e       = tail . a
        | otherwise = id

Now, the final version in the article did not use this technique, as it was very inefficient. It used some cleverness beyond the scope of this post. The second-from-last version I quite liked, though:

dropWhile' :: (a -> Bool) -> [a] -> [a]
dropWhile' p l = foldr f l l where
  f e a | p e       = tail a
        | otherwise = l

However, the idea of building up a function in a fold gave me an idea for adapting it to some of the Trie functions.

Folding Inwards

Let’s start with member. It needs to fold over a list, and generate a function which acts on a Trie:

member :: Ord a => [a] -> Trie a -> Bool
member = foldr f base

The base is the function being built up: the final part of the function chain. Each part of the function is generated based on each element of the list, and then chained with the base using .:

member = foldr f base where
  f e a = ??? . a 

The base here is what’s called when the list is empty. Here’s what it looked like in the explicit recursion version:

member [] (Trie e _) = e

We could simplify this by using record syntax, and getTrie:

member [] t = getTrie t

And this has an obvious pointsfree version:

member [] = getTrie

That fits for the base case. It’s just a function:

member = foldr f endHere where
  f e a = ??? . a 

Then, how to combine it. That’s easy enough, actually. It accesses the map, searches it for the key, and calls the accumulating function on it. If it’s not found in the map, just return False. Here it is:

member :: Ord a => [a] -> Trie a -> Bool
member = foldr f endHere where
  f e a = fromMaybe False . fmap a . M.lookup e . getTrie

One of the other standard functions for a Trie is returning the “completions” for a given sequence. It’s a very similar function to member, actually: instead of calling endHere on the final Trie found, though, just return the Trie itself. And the thing to return if any given element of the sequence isn’t found is just an empty Trie:

complete :: Ord a => [a] -> Trie a -> Trie a
complete = foldr f id where
  f e a = fromMaybe empty . fmap a . M.lookup e . getTrie 

In fact, you could abstract out the commonality here:

follow :: Ord a => c -> (Trie a -> c) -> [a] -> Trie a -> c
follow ifMiss onEnd = foldr f onEnd where
  f e a = fromMaybe ifMiss . fmap a . M.lookup e . getTrie 
  
member :: Ord a => [a] -> Trie a -> Bool
member = follow False endHere

complete :: Ord a => [a] -> Trie a -> Trie a
complete = follow empty id

Folding in and out

insert is another deal entirely. In member, the fold was tunneling into a Trie, applying the accumulator function to successively deeper Tries, and returning a result based on the final Trie. insert needs to do the same tunneling – but the Trie returned needs to be the outer Trie.

It turns out it’s not that difficult. Instead of “building up a function” that is then applied to a Trie, here a function is “sent” into the inner Tries. The cool thing here is that the function being sent hasn’t been generated yet.

Here’s some more illustration of what I mean. Start off with the normal foldr:

insert :: Ord a => [a] -> Trie a -> Trie a
insert = foldr f (\(Trie _ m) -> Trie True m)

With the final function to be applied being one that just flips the endHere tag to True. Then f: this is going to act over the map of the Trie that it’s called on. It’s useful to define a function just for that:

overMap :: Ord b => (M.Map a (Trie a) -> M.Map b (Trie b)) -> Trie a -> Trie b
overMap f (Trie e m) = Trie e (f m)

Then, it will look up the next element of the sequence in the Trie, and apply the accumulating function to it. (if it’s not found it will provide an empty Trie instead) Simple!

insert :: Ord a => [a] -> Trie a -> Trie a
insert = foldr f (\(Trie _ m) -> Trie True m) where
  f e a = overMap (M.alter (Just . a . fromMaybe (Trie False M.empty)) e)

I think this is really cool: with just a foldr, you’re burrowing into a Trie, changing it, and burrowing back out again.

Removal

This is always the tricky one with a Trie. You can just follow a given sequence down to its tag, and flip it from on to off. But that doesn’t remove the sequence itself from the Trie. So maybe you just delete the sequence – but that doesn’t work either. How do you know that there are no other sequences stored below the one you were examining?

What you need to do is to send a function into the Trie, and have it report back as to whether or not it stores other sequences below it. So this version of foldr is going to burrow into the Trie, like member; maintain the outer Trie, like insert; but also send messages back up to the outer functions. Cool!

The way to do the “message sending” is with Maybe. If the function you send into the Trie to delete the end of the sequence returns Nothing, then it signifies that you can delete that member. Luckily, the alter function on Data.Map works well with this:

alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a

Its first argument is a function which is given the result of looking up its second argument. If the function returns Nothing, that key-value pair in the map is deleted (if it was there). If it returns Just something, though, that key-value pair is added. In the delete function, we can chain the accumulating function with =<<. This will skip the rest of the accumulation if any part of the sequence isn’t found. The actual function we’re chaining on is nilIfEmpty, which checks if a given Trie is empty, and returns Just the Trie if it’s not, or Nothing otherwise.

Here’s the finished version:

delete :: Ord a => [a] -> Trie a -> Trie a
delete = (fromMaybe empty .) . foldr f i where
  i (Trie _ m) | M.null m  = Nothing
               | otherwise = Just (Trie False m)
  f e a = nilIfEmpty . overMap (M.alter (a =<<) e) 
  
null :: Trie a -> Bool
null (Trie e m) = (not e) && (M.null m)

nilIfEmpty :: Trie a -> Maybe (Trie a)
nilIfEmpty t | null t    = Nothing
             | otherwise = Just t

Folding the Foldable

So how about folding the Trie itself? Same trick: build up a function with a fold. This time, a fold over the map, not a list. And the function being built up is a cons operation. When you hit a True tag, fire off an empty list to the built-up function, allowing it to evaluate:

foldrTrie :: ([a] -> b -> b) -> b -> Trie a -> b
foldrTrie f i (Trie a m) = M.foldrWithKey ff s m where
  s    = if a then f [] i else i
  ff k = flip (foldrTrie $ f . (k :))

Unfortunately, it’s not easy to make the Trie conform to Foldable. It is possible, and it’s what I’m currently trying to figure out, but it’s non-trivial.

Monty Hall

The Monty Hall problem is a great example of how counter-intuitive probability can sometimes be. It goes something like this: say you’re on a gameshow, with the chance to win a car. You’re shown three doors, and the car is behind one, goats behind the other two. You pick a door, say the leftmost, but then the host of the gameshow stops you before it’s opened. He opens one of the two doors you didn’t pick, revealing a goat. He then asks you if you’d like to change your decision. So? Do you?

Read More

Making FlatMap lazy

flatMap() is pretty cool. When I first learned what map() did, it got shoved into every line of code I had, regardless of how applicable it was. (That’s usually a good sign.) Then, along came flatMap(). It’s deceptively simple: its name pretty much explains exactly what it does. Useful, not game-changing, although it does let you do some pretty cool things:

Read More

Swift Protocols: A Strategy

A Misguided, Over-Simplified Strategy

It Makes Sense to Me, so…

So I’ve been drinking the Protocol-Oriented-Programming gatorade for a while now. I’ve taken it to the extreme a little: you won’t find a class in pretty much any of my code these days. So, before I pull it back a little, I thought I’d put down my strategy so far for how to handle these protocol things.

Read More