How to enumerate a recursive datatype in Haskell?

Question

This blog post has an interesting explanation of how to use the Omega monad to enumerate an arbitrary grammar diagonally. He offers an example of how to do so, resulting in

Answer 1

Below is a terrible solution, but perhaps an interesting one.

---

We might consider the idea of adding "one more layer"

grow :: T -> Omega T
grow t = each [A, B t, C t t]

which is close to correct but has a flaw—in particular, in the C branch, we end up having both of the arguments take the exact same values instead of being able to vary independently. We can fix this by computing the "base functor" of T which looks like this

data T    = A  | B  T | C  T T
data Tf x = Af | Bf x | Cf x x deriving Functor

In particular, Tf is just a copy of T where the recursive calls are functor "holes" instead of direct recursive calls. Now we can write:

grow :: Omega T -> Omega (Tf (Omega T))
grow ot = each [ Af, Bf ot, Cf ot ot ]

which has a whole computation of a new set of T in each hole. If we could somehow "flatten" the Omega (Tf (Omega T)) into Omega T then we'd have a computation which adds "one new layer" to our Omega computation correctly.

flatten :: Omega (Tf (Omega T)) -> Omega T
flatten = ...

and we could take the fixed point of this layering with fix

fix :: (a -> a) -> a

every :: Omega T
every = fix (flatten . grow)

So the only trick is to figure out flatten. To do this we need to notice two features of Tf. First, it's Traversable so we can use sequenceA to "flip" the order of Tf and Omega

flatten = ?f . fmap (?g . sequenceA)

where ?f :: Omega (Omega T) -> Omega T is just join. The final tricky bit is figuring out ?g :: Omega (Tf T) -> Omega T. Obviously, we don't care about the Omega layer so we should just fmap a function of type Tf T -> T.

And this function is very close to the defining notion for the relationship between Tf and T: we can always compress a layer of Tf on the top of T.

compress :: Tf T -> T
compress Af         = A
compress (Bf t)     = B t
compress (Cf t1 t2) = C t1 t2

All together we have

flatten :: Omega (Tf (Omega T)) -> Omega T
flatten = join . fmap (fmap compress . sequenceA)

Ugly, but all together functional.