Recovering sharing in a GADT

soreifygadt.hs

{-# LANGUAGE GADTs, TypeFamilies, FlexibleInstances, RankNTypes, StandaloneDeriving, DeriveDataTypeable #-}

import Data.Reify
import Control.Applicative
import Data.Map
import Data.Typeable


data Ast e where
  IntLit :: Int -> Ast Int
  Add :: Ast Int -> Ast Int -> Ast Int
  BoolLit :: Bool -> Ast Bool
  IfThenElse :: Typeable e => 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 :: Typeable e => Ast2 Bool s -> Ast2 e s -> Ast2 e s -> Ast2 e s
  Var :: Typeable e => s -> Ast2 e s
  
deriving instance Show s => Show (Ast2 e s)
deriving instance Typeable2 Ast2


data WrappedAst s where
  Wrap :: Typeable e => Ast2 e s -> WrappedAst s

deriving instance Show (WrappedAst Unique)


instance Typeable e => MuRef (Ast e) where
  type DeRef (Ast e) = WrappedAst
  mapDeRef f e = Wrap <$> 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 <$> f a) <*> (Var <$> f b)
      mapDeRef' f (BoolLit b) = pure $ BoolLit2 b
      mapDeRef' f (IfThenElse b t e) = IfThenElse2 <$> (Var <$> f b) <*> (Var <$> f t) <*> (Var <$> f e)


getVar :: Typeable e => Map Name (WrappedAst Name) -> Name -> Maybe (Ast2 e Name)
getVar m n = case m ! n of Wrap e -> cast e


conv :: Typeable e => 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 n)


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