-- Source code for Hugs 1.3c,
-- extracted from "de Bruijn notation as a nested datatype",
-- by Richard S. Bird and Ross Paterson

-- Section 2.  Preliminaries

data BinTree a = Leaf a | Fork (Pair (BinTree a))
type Pair a = (a,a)

mapP :: (a -> b) -> Pair a -> Pair b
mapP f (x, y) = (f x, f y)

mapB :: (a -> b) -> BinTree a -> BinTree b
mapB f (Leaf x) = (Leaf . f) x
mapB f (Fork p) = (Fork . mapP (mapB f)) p

foldB :: (a -> b) -> (Pair b -> b) -> BinTree a -> b
foldB l f (Leaf x) = l x
foldB l f (Fork p) = (f . mapP (foldB l f)) p

size    = foldB (const 1) (uncurry (+))
height  = foldB (const 0) maxp
	  where maxp (x, y) = 1 + max x y
flatten = foldB wrap (uncurry (++))
	  where wrap x = [x]

joinB :: BinTree (BinTree a) -> BinTree a
joinB = foldB id Fork

-- Section 3.  de Bruijn notation

data Term v = Var v | App (Pair (Term v)) | Lam (Term (Incr v))
data Incr v = Zero | Succ v

mapI :: (a -> b) -> Incr a -> Incr b
mapI f Zero = Zero
mapI f (Succ x) = (Succ . f) x

mapT :: (a -> b) -> Term a -> Term b
mapT f (Var x) = (Var . f) x
mapT f (App p) = (App . mapP (mapT f)) p
mapT f (Lam t) = (Lam . mapT (mapI f)) t

foldT :: (forall a. a -> n a) ->
	 (forall a. Pair (n a) -> n a) ->
	 (forall a. n (Incr a) -> n a) ->
	 Term b -> n b
foldT v a l (Var x) = v x
foldT v a l (App p) = (a . mapP (foldT v a l)) p
foldT v a l (Lam t) = (l . foldT v a l) t

gfoldT :: (forall a. m a -> n a) ->
	  (forall a. Pair (n a) -> n a) ->
	  (forall a. n (Incr a) -> n a) ->
	  (forall a. Incr (m a) -> m (Incr a)) ->
	  Term (m b) -> n b
gfoldT v a l k (Var x) = v x
gfoldT v a l k (App p) = (a . mapP (gfoldT v a l k)) p
gfoldT v a l k (Lam t) = (l . gfoldT v a l k . mapT k) t

kfoldT :: (a -> b) -> (Pair b -> b) -> (b -> b) ->
	  (Incr a -> a) ->
	  Term a -> b
kfoldT v a l k (Var x) = v x
kfoldT v a l k (App p) = (a . mapP (kfoldT v a l k)) p
kfoldT v a l k (Lam t) = (l . kfoldT v a l k . mapT k) t

showT :: Term String -> String
showT = kfoldT id showP ('L':) showI
	where showP (x,y) = "(" ++ x ++ " " ++ y ++ ")"
              showI Zero = "0"
              showI (Succ x) = 'S':x

showTC :: Term Char -> String
showTC = showT . mapT wrap
         where wrap x = [x]

joinT :: Term (Term a) -> Term a
joinT = gfoldT id App Lam distT

distT :: Incr (Term a) -> Term (Incr a)
distT Zero = Var Zero
distT (Succ x) = mapT Succ x

-- Section 4.  Abstraction and application

abstract :: Eq a => a -> Term a -> Term a
abstract x = Lam . mapT (match x)

match :: Eq a => a -> a -> Incr a
match x y = if x == y then Zero else Succ y

apply :: Term a -> Term (Incr a) -> Term a
apply t = joinT . mapT (subst t . mapI Var)

subst :: a -> Incr a -> a
subst x Zero = x
subst x (Succ y) = y

-- Section 5.  An extension of de Bruijn's notation

data TermE a = VarE a
	     | AppE (Pair (TermE a))
	     | LamE (TermE (Incr (TermE a)))

mapE :: (a -> b) -> TermE a -> TermE b
mapE f (VarE x) = (VarE . f) x
mapE f (AppE p) = (AppE . mapP (mapE f)) p
mapE f (LamE t) = (LamE . mapE (mapI (mapE f))) t

gfoldE :: (forall a. m a -> n a) ->
	  (forall a. Pair (n a) -> n a) ->
	  (forall a. n (Incr (n a)) -> n a) ->
	  (forall a. Incr a -> m (Incr a)) ->
	  TermE (m b) -> n b
gfoldE v a l k (VarE x) = v x
gfoldE v a l k (AppE p) = (a . mapP (gfoldE v a l k)) p
gfoldE v a l k (LamE t) = (l . gfoldE v a l k .
			   mapE (k . mapI (gfoldE v a l k))) t

joinE :: TermE (TermE a) -> TermE a
joinE = gfoldE id AppE LamE VarE

abstractE :: Eq a => a -> TermE a -> TermE a
abstractE x = LamE . mapE (mapI VarE . match x)

applyE :: TermE a -> TermE (Incr (TermE a)) -> TermE a
applyE t = joinE . mapE (subst t)

-- specialized version of gfoldE for cvtE

gfoldE' :: (forall a. a -> n a) ->
	  (forall a. Pair (n a) -> n a) ->
	  (forall a. n (Incr (n a)) -> n a) ->
	  (forall a. Incr a -> Incr a) ->
	  TermE b -> n b
gfoldE' v a l k (VarE x) = v x
gfoldE' v a l k (AppE p) = (a . mapP (gfoldE' v a l k)) p
gfoldE' v a l k (LamE t) = (l . gfoldE' v a l k .
			   mapE (k . mapI (gfoldE' v a l k))) t

cvtE :: TermE a -> Term a
cvtE = gfoldE' Var App (Lam . joinT . mapT distT) id

cvtBodyE :: TermE (Incr (TermE a)) -> Term (Incr a)
cvtBodyE = joinT . mapT distT . cvtE . mapE (mapI cvtE)