The free variables plate with the GADT variant of Multiplate.

multiplate-ix-freevars.hs

{-# LANGUAGE RankNTypes, GADTs #-}

import Data.Monoid
import Data.Functor.Constant
import Data.Type.Equality
import Control.Applicative


type Plate fam f = forall x. fam x -> x -> f x

class EqT fam => Multiplate fam where
  multiplate :: Applicative f => Plate fam f -> Plate fam f

appendPlate :: (Multiplate fam, Monoid o) => Plate fam (Constant o) -> Plate fam (Constant o) -> Plate fam (Constant o)
appendPlate f1 f2 w a = f1 w a <* f2 w a

foldFor :: Multiplate fam => fam a -> Plate fam (Constant o) -> a -> o
foldFor w f = getConstant . f w


data Expr = Con Int
          | Add Expr Expr
          | Mul Expr Expr
          | EVar Var
          | Let Decl Expr
          deriving (Eq, Show)

data Decl = Var := Expr
          | Seq Decl Decl
          deriving (Eq, Show)

type Var = String

data AST a where
  Expr :: AST Expr
  Decl :: AST Decl

instance EqT AST where
  eqT Expr Expr = Just Refl
  eqT Decl Decl = Just Refl
  eqT _ _ = Nothing

instance Multiplate AST where
  multiplate child Expr (Add e1 e2) = Add <$> child Expr e1 <*> child Expr e2
  multiplate child Expr (Mul e1 e2) = Mul <$> child Expr e1 <*> child Expr e2
  multiplate child Expr (Let d e) = Let <$> child Decl d <*> child Expr e
  multiplate _     Expr e = pure e
  multiplate child Decl (v := e) = (v :=) <$> child Expr e
  multiplate child Decl (Seq d1 d2) = Seq <$> child Decl d1 <*> child Decl d2


type FreeVarPlate = Plate AST (Constant ([Var] -> ([Var], [Var])))

freeVariablesPlate :: FreeVarPlate
freeVariablesPlate = handleLet (varPlate `appendPlate` multiplate freeVariablesPlate)
  where
    varPlate :: FreeVarPlate
    varPlate Expr (EVar v) = Constant $ \bounded -> (if (v `elem` bounded) then [] else [v], [])
    varPlate Decl (v := _) = Constant $ const ([], [v])
    varPlate _    x        = pure x
    handleLet :: FreeVarPlate -> FreeVarPlate
    handleLet plate Expr (Let d e) = Constant $ \bounded -> 
      let
        (freeD, declD) = foldFor Decl plate d bounded
        (freeE, _)     = foldFor Expr plate e (declD ++ bounded)
      in
        (freeD ++ freeE, [])
    handleLet plate w x = plate w x

freeVars :: Expr -> [Var]
freeVars = fst . ($ []) . foldFor Expr freeVariablesPlate

expr1 :: Expr
expr1 = Let ("x" := Con 42) (Add (EVar "x") (EVar "y"))

test :: [Var]
test = freeVars expr1