How to walk a complex AST generically

Setup

Let's say you have a complex AST. For instance, an AST for Haskell code.

Some nodes are annotated with Comments. Some nodes are annotated with HaddockComments. They are different types, because you thought it would be a good idea (and maybe it is!).

Collect all 'Comment's

If you need to collect all Comments, it's simple enough with uniplate. You derive Data for all your types and then do this:

import Data.Data
import qualified Data.Generics.Uniplate.Data as Uniplate

allComments :: Data a => a -> [Comment]
allComments = Uniplate.universeBi

Collect all comments, both Comment and HaddockComment, in the order that they appear

For certain nefarious purposes you want to collect a list of all comments in the AST, both Comment and HaddockComment, in the order that they appear. Oh shit. What do you do?

Uniplate can't save you now. The answer is — you need to use Data.Data, a weird hodge-podge of methods like gunfold and gmapQl and gmapMo.

You always avoided it because you thought gunfold had something to do with guns. Fear not! We don't have to use gunfold, not this time at least. We just need gmapQ.

gmapQ :: Data a => (forall d. Data d => d -> u) -> a -> [u]

It lets you apply some function to all children of a data type, even if those children have different types. Is it polymorphic recursion? I guess it is. Notice how cast saves the day.

import Data.Data as Data
import Data.Typeable as Typeable

reallyAllComments :: Data a => a -> [Either Comment HaddockComment]
reallyAllComments = go
  where
    go :: forall d. Data d => d -> [Either Comment HaddockComment]
    go x
      | Just comment <- Typeable.cast x =
          [Left comment]
      | Just haddockComment <- Typeable.cast x =
          [Right haddockComment]
      | otherwise =
          concat $ Data.gmapQ go x

Change all comments

You don't want to simply get them out of the structure, you want to change them. We'll need one more trick — an equality witness.

changeAllComments :: Data a => a -> a
changeAllComments = go
  where
    go :: forall d. Data d => d -> d
    go x
      | Just Refl <- Typeable.eqT @d @Comment =
          <something with x>
      | Just Refl <- Typeable.eqT @d @HaddockComment =
          <something with x>
      | otherwise =
          Data.gmapT go x