Convenient embedded abstract syntax for Haskell.

ReifiedEmbeddedAbstractSyntax.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# OPTIONS_GHC -Wall #-}

import Data.Reify ( reifyGraph, MuRef(..), Graph(..))
import Data.Function (fix)
import Data.IntMap.Strict qualified as Map
import Data.Foldable ( Foldable(toList) )
import Data.Bifunctor ( Bifunctor(second) )

-- Public data types:

data Lam = Var Lam | Abs Lam | App Lam Lam deriving Show
data LamN = VarN Int | AbsN Int LamN | AppN LamN LamN deriving Show
data LetRec = LetRec [(Int, LamN)] Int deriving Show

-- Private data types

data Free f a = Pure a | Free (f (Free f a)) deriving Functor
deriving instance (forall x. Show x => Show (f x), Show a) => Show (Free f a)
data LamF f = VarF f | AbsF f | AppF f f deriving (Eq, Ord, Foldable, Traversable, Functor, Show)
data LamF' f = AbsF' Int f | AppF' f f deriving (Eq, Ord, Foldable, Traversable, Functor, Show)

-- Black magic:

instance MuRef Lam where
  type DeRef Lam = LamF
  mapDeRef f (Var x) = VarF <$> f x
  mapDeRef f (App x y) = AppF <$> f x <*> f y
  mapDeRef f (Abs x) = AbsF <$> f x

inline :: Graph LamF -> Graph (Free LamF')
inline (Graph g z) = Graph (Map.toList (Map.filterWithKey (\i _ -> not (isUnique Map.! i)) res)) z where
  isUnique = Map.unionsWith (\_ _ -> False)
    (Map.singleton z False : concatMap ((\x -> if not (isVar x) then map (`Map.singleton` True) (toList x) else []) . snd) g)
  res = Map.fromList [(i, go i x) | (i, x) <- g]
  go _ (VarF x) = Pure x
  go i (AbsF x) = Free (AbsF' i (helper x))
  go _ (AppF x y) = Free (AppF' (helper x) (helper y))

  helper :: Int -> Free LamF' Int
  helper i 
    | isUnique Map.! i = res Map.! i
    | otherwise = Pure i

isVar :: LamF Map.Key -> Bool
isVar (VarF _) = True
isVar _ = False

-- Public syntax interface:

lam :: (Lam -> Lam) -> Lam
lam f = fix $ Abs . f . Var

infixl 9 $$
($$) :: Lam -> Lam -> Lam
f $$ x = App f x

myConst :: Lam
myConst = lam \x -> lam \_ -> x

myId :: Lam
myId = lam \x -> x

myFlip :: Lam
myFlip = lam \f -> lam \x -> lam \y -> f $$ y $$ x

-- Public transformation interface:

reify :: Lam -> IO LetRec
reify x0 = do
  Graph g z <- inline <$> reifyGraph x0
  let
    go (Pure x) = VarN x
    go (Free (AbsF' i f)) = AbsN i (go f)
    go (Free (AppF' f1 f2)) = AppN (go f1) (go f2)
  pure (LetRec (map (second go) g) z)

main :: IO ()
main = do
  print =<< reify myId
  print =<< reify myConst
  print =<< reify myFlip
  print =<< reify (myId $$ myId)

-- Results:
--
-- LetRec [(1,AbsN 1 (VarN 1))] 1
-- LetRec [(1,AbsN 1 (AbsN 2 (VarN 1)))] 1
-- LetRec [(1,AbsN 1 (AbsN 2 (AbsN 3 (AppN (AppN (VarN 1) (VarN 3)) (VarN 2)))))] 1
-- LetRec [(1,AppN (VarN 2) (VarN 2)),(2,AbsN 2 (VarN 2))] 1