-- Code from Section 12.4

module Traversal where
import PhantomTypes

type Name = String
type Age = Int
data Person = Person Name Age
  deriving Show

data Type t =              RInt    (Int    :=: t)
            |              RChar   (Char   :=: t)
            | forall a   . RList   ([a]    :=: t) (Type a)
            | forall a b . RPair   ((a,b)  :=: t) (Type a) (Type b)
            |              RPerson (Person :=: t)


tick :: Name -> Traversal
tick s (RPerson p) person | s==n = from p $ Person n (a+1)
   where Person n a = to p person
tick s rt t = t

rInt :: Type Int
rInt = RInt refl

rChar :: Type Char
rChar = RChar refl

rList :: forall a . Type a -> Type [a]
rList t = RList refl t

rPair :: forall a b . Type a -> Type b -> Type (a,b)
rPair t1 t2 = RPair refl t1 t2

rPerson :: Type Person
rPerson = RPerson refl

rString :: Type String
rString = rList rChar


type Traversal = forall t . Type t -> t -> t

copy :: Traversal
copy rt = id

o :: Traversal -> Traversal -> Traversal
(f `o` g) rt = f rt . g rt

imap :: Traversal -> Traversal
imap f (RInt p) i = i
imap f (RChar p) c = c
imap f (RList p ra) as 
 = case to p as of 
     [] -> from p []
     (a:as) -> from p $ f ra a : to p (f (RList p ra) (from p as))
imap f (RPair p ra rb) pair = let (a,b) = to p pair in from p (f ra a,f rb b)
imap f (RPerson p) person 
 = let Person n a = to p person in from p $ Person (f rString n) (f rInt a)

everywhere, everywhere' :: Traversal -> Traversal
everywhere f = f `o` imap (everywhere f)
everywhere' f = imap (everywhere' f) `o` f

type Query th = forall t . Type t -> t -> th

isum :: Query Int -> Query Int
isum f (RInt p) i = 0
isum f (RChar p) c = 0
isum f (RList p ra) as = case to p as of
                           [] -> 0
                           (a:as) -> f ra a + f (RList p ra) (from p as)
isum f (RPair p ra rb) pair = let (a,b) = to p pair in f ra a + f rb b
isum f (RPerson p) person 
 = let Person n a = to p person in f rString n + f rInt a

total :: Query Int -> Query Int
total f rt t = f rt t + isum (total f) rt t

age :: Query Age
age (RPerson p) person = let Person n a = to p person in a
age _ _ = 0

sizeof :: Query Int
sizeof (RInt p) _ = 2
sizeof (RChar p) _ = 2
sizeof (RList p ra) as = case to p as of 
                           [] -> 0
                           (_:_) -> 3
sizeof (RPair p ra rb) _ = 3
sizeof (RPerson p) _ = 3