This seems like a good use-case for biplates. I'm relying on the slower Data.Data
method, but it makes this code pretty trivial.
{-# LANGUAGE DeriveDataTypeable #-}
import Data.Data
import Data.Typeable
import Data.Generics.Uniplate.Data
data Statement
= Return Expression
| If Expression Expression
deriving(Data, Typeable)
data Expression
= Literal Int
| Variable String
| Binary Expression Int Expression
deriving(Data, Typeable)
vars :: Statement -> [String]
vars stmt = [ s | Variable s <- universeBi stmt]
Basically biplates are a generalized notion of uniplates where the target type isn't necessarily the same as the source, eg
biplate :: from -> (Str to, Str to -> from)