content: {-# LANGUAGE GADTs, RankNTypes #-}
  
  module Main where
  
  import Control.Applicative (Applicative(..))
  import Control.Monad (liftM, ap)
  
  -- abstract syntax -------------------------------------------------------------
  data Ty t where
    Bool :: Ty Bool
    Arr  :: Ty a -> Ty b -> Ty (a -> b)
  
  data Exp g t where
    Var 	 :: Var g t -> Exp g t
    Lam 	 :: Ty a -> Exp (g,a) b -> Exp g (a->b) 
    App 	 :: Exp g (s -> t) -> Exp g s -> Exp g t
    If  	 :: Exp g Bool -> Exp g t -> Exp g t -> Exp g t
    ETrue  :: Exp g Bool
    EFalse :: Exp g Bool
  
  data Var g t where
    ZVar :: Var (h,t) t
    SVar :: Var h t -> Var (h,s) t
  
  -- smart constructors ----------------------------------------------------------
  lamE :: Ty s -> (Exp (g,s) s -> Exp (g,s) t) -> Exp g (s -> t)
  lamE s f = Lam s (f (Var ZVar))
  
  ifE :: Exp g Bool -> Exp g t -> Exp g t -> Exp g t
  ifE t ETrue EFalse = t
  ifE t e e' = if eqE e e' then e else If t e e'
  
  -- boring equality tests -------------------------------------------------------
  eqB :: BoxExp t -> BoxExp s -> Bool
  eqB (Box e) (Box e_) = eqE e e_
  
  eqE :: Exp g t -> Exp h s -> Bool
  eqE (Var x) (Var y) = eqV x y
  eqE (Lam s e) (Lam s_ e_) = eqT s s_ && eqE e e_
  eqE (App e1 e2) (App e1_ e2_) = eqE e1 e1_ && eqE e2 e2_
  eqE (If e1 e2 e3) (If e1_ e2_ e3_) = eqE e1 e1_ && (eqE e2 e2_ && eqE e3 e3_)
  eqE (ETrue) (ETrue) = True
  eqE (EFalse) (EFalse) = True
  eqE _ _ = False
  
  eqT :: Ty t -> Ty s -> Bool
  eqT (Arr s t) (Arr s_ t_) = eqT s s_ && eqT t t_
  eqT Bool Bool = True
  eqT _ _ = False
  
  eqV :: Var g t -> Var h s -> Bool
  eqV (SVar x) (SVar y) = eqV x y
  eqV ZVar ZVar = True
  eqV _ _ = False
  
  -- evaluation ------------------------------------------------------------------
  var :: Var g t -> g -> t
  var ZVar     (_,t) = t
  var (SVar x) (h,s) = var x h
  
  eval :: Exp g t -> g -> t
  eval (Var x)    g = var x g
  eval (Lam _ e)  g = \a -> eval e (g,a)
  eval (App e e') g = eval e g (eval e' g)
  eval (ETrue)    g = True
  eval (EFalse)   g = False
  eval (If c t e) g = if eval c g then eval t g else eval e g
  
  -- type inference --------------------------------------------------------------
  data TyEnv g where
    Nil :: TyEnv g
    Cons :: Ty t -> TyEnv h -> TyEnv (h,t)
  
  infer :: TyEnv g -> Exp g t -> Ty t
  infer g (Var x)		= inferVar g x
  infer g (Lam t e)	= Arr t (infer (Cons t g) e)
  infer g (App e e')	= case infer g e of Arr _ t -> t
  infer g (ETrue)		= Bool
  infer g (EFalse)	= Bool
  infer g (If _ e _)	= infer g e
  
  inferVar :: TyEnv g -> Var g t -> Ty t
  inferVar (Cons t h) (SVar x) = inferVar h x
  inferVar (Cons t h) (ZVar)   = t
  
  -- tree monad ------------------------------------------------------------------
  
  data Tree a = Val a | Choice (Tree a) (Tree a)
  -- doesn't yet force trees to be fully balanced:
  -- 	Val :: a -> Tree a Z
  -- 	Choice :: Tree a n -> Tree a n -> Tree a (S n)
  
  instance Functor Tree where
      fmap = liftM
  
  instance Applicative Tree where
      pure = return
      (<*>) = ap
  
  instance Monad Tree where
    return x = Val x
    (Val a) >>= f = f a
    (Choice l r) >>= f = Choice (l >>= f) (r >>= f)
  
  tmap :: Monad m => (a->b) -> m a -> m b
  tmap f x = do { a <- x; return (f a) }
  
  flatten t = flatten_ t []
   where
     flatten_ (Val a)      k = a:k
     flatten_ (Choice l r) k = flatten_ l (flatten_ r k)
  
  
  -- quote & friends -------------------------------------------------------------
  
  -- for values --------------------------
  enumV		:: Ty t -> Tree t
  questionsV	:: Ty t -> [t -> Bool]
  
  
  enumV Bool      = Choice (Val True) (Val False)
  enumV (Arr s t) = mkEnum (questionsV s) (enumV t)
   where
     mkEnum [] t = tmap const t
     mkEnum (q:qs) es = do
     		   f1 <- mkEnum qs es
  		   f2 <- mkEnum qs es
  		   return (\d -> if q d then f1 d else f2 d)
  
  questionsV Bool		= return (\x -> x)
  questionsV (Arr s t)	= do
  			  d <- flatten (enumV s)
  			  q <- questionsV t
  			  return (\f -> q (f d))
  
  -- for expressions ---------------------
  enumE		:: Ty t -> Tree (Exp g t)
  questionsE	:: Ty t -> [Exp g t -> Exp g Bool]
  
  enumE Bool      = Choice (Val ETrue) (Val EFalse)
  enumE (Arr s t) = tmap (lamE s) (mkEnumE (questionsE s) (enumE t))
   where
     mkEnumE [] t = tmap const t
     mkEnumE (q:qs) es = do
     			f1 <- mkEnumE qs es
  			f2 <- mkEnumE qs es
  			return (\d -> ifE (q d) (f1 d) (f2 d))
  
  questionsE Bool		= return (\x -> x)
  questionsE (Arr s t)	= do
  			  d <- flatten (enumE s)
  			  q <- questionsE t
  			  return (\f -> q (App f d))
  
  -- should be 
  -- 	find (List (Exp g Bool) n) -> Tree (Exp g a) n -> Exp g a
  find :: [Exp g Bool] -> Tree (Exp g a) -> Exp g a
  find []		(Val a)		= a
  find (b:bs)	(Choice l r)	= ifE b (find bs l) (find bs r)
  find _		_		= error "bad arguments to find"
  
  quote :: Ty t -> t -> Exp g t
  quote Bool      t = case t of True -> ETrue; False -> EFalse
  quote (Arr s t) f = lamE s (\e -> find (do q <- questionsE s; return (q e))
  					(tmap (quote t . f) (enumV s)))
  
  -- normalization (by evaluation) -----------------------------------------------
  data BoxExp t = Box (forall g. Exp g t)
  
  normalize :: Ty t -> BoxExp t -> BoxExp t
  normalize s (Box e) = Box (quote s (eval e ()))
  
  -- examples --------------------------------------------------------------------
  b2b = Arr Bool Bool
  b22b = Arr b2b b2b
  zero = Var ZVar
  one = Var (SVar ZVar)
  once   = Box (Lam b2b (Lam Bool (App one zero)))
  twice  = Box (Lam b2b (Lam Bool (App one (App one zero))))
  thrice = Box (Lam b2b (Lam Bool (App one (App one (App one zero)))))
  
  test = [ eqB (nf b22b thrice) (nf b22b once)
         , eqB (nf b22b twice)  (nf b22b once)]
    where nf = normalize
  
  main = print test