Recovering sharing in a GADT

soreifygadt.hs

{-# LANGUAGE GADTs, TypeFamilies, FlexibleInstances, RankNTypes #-}
import Data.Reify
import Control.Applicative
import Data.Map

data Ast e where
  IntLit :: Int -> Ast Int
  Add :: Ast Int -> Ast Int -> Ast Int
  BoolLit :: Bool -> Ast Bool
  IfThenElse :: Ast Bool -> Ast e -> Ast e -> Ast e

type Name = Unique

data Ast2 e s where
  IntLit2 :: Int -> Ast2 Int s
  Add2 :: Ast2 Int s -> Ast2 Int s -> Ast2 Int s
  BoolLit2 :: Bool -> Ast2 Bool s
  IfThenElse2 :: Ast2 Bool s -> Ast2 e s -> Ast2 e s -> Ast2 e s
  Var :: Type e -> s -> Ast2 e s
  
data Type a where
  Bool :: Type Bool
  Int :: Type Int

data Refl a b where Refl :: Refl a a

typeEq :: Type a -> Type b -> Maybe (Refl a b)
typeEq Bool Bool = Just Refl
typeEq Int Int = Just Refl
typeEq _ _ = Nothing

getType :: Ast e -> Type e
getType (IntLit _) = Int
getType (Add _ _) = Int
getType (BoolLit _) = Bool
getType (IfThenElse _ _ e) = getType e

data WrappedAst s where
  Wrap :: Type e -> Ast2 e s -> WrappedAst s


instance MuRef (Ast e) where
  type DeRef (Ast e) = WrappedAst
  mapDeRef f e = Wrap (getType e) <$> mapDeRef' f e
    where
      mapDeRef' :: Applicative f => (forall b. (MuRef b, WrappedAst ~ DeRef b) => b -> f u) -> Ast e -> f (Ast2 e u)
      mapDeRef' f (IntLit i) = pure $ IntLit2 i
      mapDeRef' f (Add a b) = Add2 <$> (Var Int <$> f a) <*> (Var Int <$> f b)
      mapDeRef' f (BoolLit b) = pure $ BoolLit2 b
      mapDeRef' f (IfThenElse b t e) = IfThenElse2 <$> (Var Bool <$> f b) <*> (Var (getType t) <$> f t) <*> (Var (getType e) <$> f e)


getVar :: Map Name (WrappedAst Name) -> Type e -> Name -> Maybe (Ast2 e Name)
getVar m t n = case m ! n of Wrap t' e -> (\Refl -> e) <$> typeEq t t'


conv :: Ast e -> IO (Map Name (WrappedAst Name), Maybe (Ast2 e Name))
conv e = do
  Graph l n <- reifyGraph e
  let m = fromList l
  return (m, getVar m (getType e) n)


test = IfThenElse (IfThenElse true true true) (Add one one) one
  where
    one = IntLit 1
    true = BoolLit True


instance Show (Type a) where
  show Bool = "Bool"
  show Int = "Int"

instance Show (WrappedAst Unique) where
  show (Wrap t a) = "(" ++ show a ++ " :: " ++ show t ++ ")"

instance Show s => Show (Ast2 e s) where
  show (IntLit2 i) = "(IntLit2 " ++ show i ++ ")"
  show (Add2 l r) = "(Add " ++ show l ++ " " ++ show r ++ ")"
  show (BoolLit2 b) = "(BoolLit2 " ++ show b ++ ")"
  show (IfThenElse2 b t e) = "(IfThenElse2 " ++ show b ++ " " ++ show t ++ " " ++ show e ++ ")"
  show (Var t n) = "(Var " ++ show t ++ " " ++ show n ++ ")"