{-# OPTIONS_GHC -fglasgow-exts #-}

----------------------------------------------------------
-- Section 1
--
-- Two basic data types we use in the paper. Lists are a
-- family of inductive types while Lam is an inductive 
-- family. 
----------------------------------------------------------

data List a  = Nil | Cons a (List a)

data Lam a   = Var a
              | App (Lam a) (Lam a)
              | Abs (Lam (Maybe a))
              deriving (Eq, Show)

----------------------------------------------------------
-- Section 2
--
-- Basic programming infrastructure for inductive types 
-- consisting of the implementations of generic and specialised 
-- fold and build combinators and their fusion laws. We also have 
-- some examples
----------------------------------------------------------


----------------------------------------------------------
-- Specific implementation of the initial algebra, fold 
-- combinator, build combinator and fold/build rule for 
-- lists. Also an example of foldr/build fusion for lists.
----------------------------------------------------------

build :: (forall b. (a -> b -> b) -> b -> b) -> [a]
build g = g (:) []

-- foldr c n (build g) = g c n

sum :: [Int] -> Int
sum xs = foldr (+) 0 xs

mapL :: (a -> b) -> [a] -> [b]
mapL f xs = build (\c n -> foldr (c . f) n xs)

sumSqs :: [Int] -> Int
sumSqs xs = foldr ((+) . sqr) 0 xs where sqr x = x * x


----------------------------------------------------------
-- Specific implementation of the initial algebra, fold 
-- combinator, build combinator and fold/build rule for 
-- a data type of expressions. Also an example of 
-- fold/build fusion for expressions.
----------------------------------------------------------

data Expr a = EVar a 
            | Lit Int 
            | Op Ops (Expr a) (Expr a)

data Ops = Add | Sub | Mul | Div

foldE :: (a -> b) -> (Int -> b) -> (Ops -> b -> b -> b) -> Expr a -> b
foldE v l o e = case e of 
                 EVar x -> v x
                 Lit i -> l i
                 Op op e1 e2 -> o op (foldE v l o e1) 
                                     (foldE v l o e2)
  
buildE :: (forall b. (a -> b) -> (Int -> b) -> (Ops -> b -> b -> b) -> b) -> Expr a
buildE g = g EVar Lit Op

-- foldE v l o (buildE g) = g v l o 

accum :: Expr a -> [a]
accum = foldE (\x -> [x]) (\i -> []) (\op -> (++))

mapE :: (a -> b) -> Expr a -> Expr b
mapE env e = buildE (\v l o -> foldE (v . env) l o e)

renameAccum :: (a -> b) -> Expr a -> [b]
renameAccum env e = foldE ((\x -> [x]) . env) (\i -> []) (\op -> (++)) e

----------------------------------------------------------
-- Generic implementation of the initial algebra, fold 
-- combinator, build combinator and fold/build rule for an 
-- arbitrary instance of the Functor class
----------------------------------------------------------

newtype M f = Inn {unInn :: f (M f)}

ffold :: Functor f => (f a -> a) -> M f -> a
ffold h (Inn k) =  h (fmap (ffold h) k)

fbuild :: Functor f => (forall b. (f b -> b) -> b) -> M f
fbuild g = g Inn

-- ffold h (fbuild g) = g h

----------------------------------------------------------
-- Specific implementation of the initial algebra, fold 
-- combinator, build combinator and fold/build rule for 
-- a data type of input/output computations. Also an example of 
-- fold/build fusion for expressions.
----------------------------------------------------------

data IntIO i o a = Val a 
                 | Inp (i -> IntIO i o a) 
                 | Outp (o, IntIO i o a)

data K i o a b = Vk a | Ik (i -> b) | Ok (o,b)

instance Functor (K i o a) where
  fmap k (Vk x)     = Vk x
  fmap k (Ik h)     = Ik (k . h)
  fmap k (Ok (y,z)) = Ok (y, k z) 

intIOfold :: (a -> b) -> ((i -> b) -> b) -> ((o,b) -> b) -> IntIO i o a -> b
intIOfold v p q k = case k of 
                        Val x      -> v x
                        Inp h      -> p (intIOfold v p q . h)
                        Outp (y,z) -> q (y, intIOfold v p q z)

intIObuild :: (forall b. (a -> b) -> ((i -> b) -> b) -> ((o,b) -> b) -> b) -> IntIO i o a 
intIObuild g = g Val Inp Outp

-- intIOfold v p q (intIObuild g) = g v p q

----------------------------------------------------------
-- Section 3
--
-- Implementation of the core theory of nested types, higher 
-- order functors and their initial algebras, hfolds, hbuilds 
-- and the hfold/hbuild fusion rule. Also, specialisation of 
-- these combinators and rules to specific nested types. 
-- Finally, an example of fusion based upon bit reversal
----------------------------------------------------------

----------------------------------------------------------
-- Examples of nested types
----------------------------------------------------------

data PTree a = PLeaf a
              | PNode (PTree (a,a)) deriving Show

instance Functor PTree where
   fmap f (PLeaf a)  = PLeaf (f a)
   fmap f (PNode xs) = PNode (fmap (pair f) xs)

-- we will use the functoriality of PTree so we establish it here.

data CList a = Pt a | CNil | CCons Int (CList (Maybe a))

-- The example of lambda terms is in the code for the introduction 

--------------------------------------------------------------
-- Higher order functors, their initial algebras and examples
--------------------------------------------------------------

class HFunctor f where
     ffmap :: Functor g => (a -> b) -> f g a -> f g b
     hfmap :: (Functor g, Functor h) => Nat g h -> Nat (f g) (f h)

-- ffmap ensures that an HFunctor maps functors to functors
-- hfmap maps morphisms between functors to morphisms between 
-- functors, ie natural transformations to natural transformations.

type Nat g h = forall a . g a -> h a

-- natural transformations are just polymorphic functions with 
-- parametricity of forall ensuring the relevant diagram commutes.

idNat :: Nat f f
idNat = id

data Mu f a = In (f (Mu f) a)

data HPTree f a = HPLeaf a | HPNode (f (a,a))

instance HFunctor HPTree where
     ffmap f (HPLeaf a) = HPLeaf (f a)
     ffmap f (HPNode a) = HPNode (fmap (pair f) a)
     hfmap f (HPLeaf a) = HPLeaf a
     hfmap f (HPNode a) = HPNode (f a)

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

data HCList f a = HPt a | HCNil | HCCons Int (f (Maybe a))

data HLam f a = HVar a | HApp (f a) (f a) | HAbs (f (Maybe a))

instance HFunctor HLam where
     hfmap s (HVar a)   = HVar a
     hfmap s (HApp u v) = HApp (s u) (s v)
     hfmap s (HAbs t)   = HAbs (s t)

     ffmap f (HVar a)   = HVar (f a)
     ffmap f (HApp u v) = HApp (fmap f u) (fmap f v)
     ffmap f (HAbs t)   = HAbs (fmap (fmap f) t)

data HBush f a = HBLeaf 
               | HBNode (a, f (f a))

instance HFunctor HBush where
   hfmap s HBLeaf          = HBLeaf
   hfmap s (HBNode (x, t)) = HBNode (x, fmap s (s t))

   ffmap f HBLeaf         = HBLeaf
   ffmap f (HBNode (x,t)) = HBNode (f x, fmap (fmap f) t)

data Bush a = BLeaf 
            | BNode (a, Bush (Bush a))

instance Functor Bush where 
      fmap f BLeaf          = BLeaf
      fmap f (BNode (x, t)) = BNode (f x, fmap (fmap f) t)

------------------------------------------------------------
-- Algebras, hfolds and hbuilds for nested types. We give generic 
-- definitions and then specialise them to the specific cases of 
-- individual nested types
------------------------------------------------------------

type Alg f g = Nat (f g) g

-- an algebra is still a map f g -> g, but this map is in the 
-- category of endofunctors, ie is a natural transformation 



hfold :: (HFunctor f, Functor g) => Alg f g -> Nat (Mu f) g
hfold m (In u) = m (hfmap (hfold m) u)

hfoldPTree :: (forall a. a -> f a) -> (forall a . f (a,a) -> f a) -> PTree a -> f a
hfoldPTree f g (PLeaf x)   = f x
hfoldPTree f g (PNode xs)  = g (hfoldPTree f g xs)


hfoldBush :: (forall a. f a) -> (forall a. (a, f (f a)) -> f a) -> 
              Bush a -> f a

hfoldBush l n BLeaf         = l 
hfoldBush l n (BNode (x,t)) = n (x, hfoldBush l n (fmap (hfoldBush l n) t))


hfoldLam :: (forall a. a -> f a) -> (forall a. f a -> f a -> f a) -> 
                (forall a. f (Maybe a) -> f a) -> Lam a -> f a

hfoldLam v ap ab (Var x)   = v x 
hfoldLam v ap ab (App d e) = ap (hfoldLam v ap ab d) 
                                (hfoldLam v ap ab e)
hfoldLam v ap ab (Abs l)   = ab (hfoldLam v ap ab l)


hbuild :: HFunctor f => (forall x. Alg f x -> Nat c x) -> Nat c (Mu f)
hbuild g = g In

hbuildPTree :: (forall x . 
                    (forall a . a -> x a) ->
                    (forall a . x (a,a) -> x a) ->
                    (forall a . c a -> x a)) -> Nat c PTree  
hbuildPTree g = g PLeaf PNode


hbuildBush :: (forall f . 
                    (forall a . f a) ->
                    (forall a . (a, f (f a)) -> f a) ->
                    (forall a . c a -> f a)) -> Nat c Bush  
hbuildBush g = g BLeaf BNode



hbuildLam :: (forall x. (forall a. a -> x a) ->
                        (forall a. x a -> x a -> x a) ->
                        (forall a. x (Maybe a) -> x a) ->
                        (forall a. c a -> x a)) -> 
                  Nat c Lam
hbuildLam g = g Var App Abs


-- hfold h . hbuild g = g h

-- hfoldPTree l n . (hbuildPTree g) = g l n

-- hfoldClist p n c . (hbuildClist g) = g p n c

-- hfoldLam v ap ab . (hbuildLam g) = g v ap ab


------------------------------------------------------------
-- The bit reversal protocol, optimised by the fold/build rule
-- for PTrees
------------------------------------------------------------

brp1 :: [a] -> [a]
brp1 = shuffle . unshuffle

shuffle :: PTree a -> [a]
shuffle = hfoldPTree (\x -> [x]) (cat . unzip) 

unshuffle :: [a] -> PTree a 
unshuffle = hbuildPTree unsh

cat :: ([a],[a]) -> [a]
cat (xs,ys) = xs ++ ys 

zip' :: ([a],[b]) -> [(a,b)]
zip' (xs,ys) = zip xs ys

unsh :: (forall a. a -> x a) -> 
          (forall a. x (a,a) -> x a) -> 
             (forall a. [a] -> x a)
unsh u z [e] = u e
unsh u z es  = z (unsh u z (zip' (uninter es)))

uninter :: [a] -> ([a],[a])
uninter []     = ([],[])
uninter (x:xs) = (x:as,bs) 
                    where (bs,as) = uninter xs

brp2 :: [a] -> [a]
brp2 xs = if length xs == 1 then xs 
                            else (cat . unzip . brp2 . zip' . uninter) xs



------------------------------------------------------------
-- Section 4
--
-- Generalised folds, and right Kan extensions. Similarly, 
-- generalise builds and left Kan extensions.
------------------------------------------------------------

newtype Comp f g a = Comp {icomp :: f (g a) }

instance (Functor f, Functor g) => Functor (f `Comp` g) where
    fmap k (Comp t) = Comp (fmap (fmap k) t)


psum :: PTree Int -> Int
psum xs = psumAux xs id

psumAux :: PTree a -> (a -> Int) -> Int
psumAux (PLeaf x) e  = e x
psumAux (PNode xs) e = psumAux xs (\(x,y) -> e x + e y)


newtype Ran g h a = Ran { iran :: forall b . (a -> g b) -> h b }

-- This definition of right Kan extension is essentially the exact
-- transcription of the categorical end formula with the end replaced
-- by the parametric universal quantifier. This is similar to the
-- implementation of the end defining natural transformations as a
-- universal quantifier.

instance Functor (Ran f g) where
    fmap f (Ran c) = Ran (\c1 -> c (c1 . f))

newtype Con k a  = Con {icon :: k}

toRan :: Functor k => Nat (k `Comp` g) h -> Nat k (Ran g h)
toRan s t = Ran (\val -> s (Comp (fmap val t)))

fromRan :: Nat k (Ran g h) -> Nat (k `Comp` g) h
fromRan s (Comp t) = iran (s t) id

-- the functions toRan and fromRan implement the isomorphism
-- characterising Ran g h

type Interp y g h = Nat y (Ran g h)

runInterp :: Interp y g h -> y a -> (a -> g b) -> h b
runInterp k y e = iran (k y) e

type InterpT f g h = forall y . Functor y => Interp y g h -> Interp (f y) g h

toAlg :: InterpT f g h -> Alg f (Ran g h)
toAlg interpT = interpT idNat




fromAlg :: HFunctor f => Alg f (Ran g h) -> InterpT f g h
fromAlg h interp = h . hfmap interp

-- toAlg and fromAlg implement the isomorphism showing that
-- interpreter transformers are just algebras for a higher order
-- functor. Interestingly, this is also a proof of the Yoneda lemma
-- in disguise.

gfold :: HFunctor f => InterpT f g h -> Nat (Mu f) (Ran g h)
gfold interpT = hfold (toAlg interpT)

-- a gfold is just i) an hfold for an algebra whose carrier is a right
-- Kan extension; and ii) the algebra is presented in the form of an
-- interpreter transformer.

rungfold :: HFunctor f => InterpT f g h -> Mu f a -> (a -> g b) -> h b
rungfold interpT = iran . gfold interpT


------------------------------------------------------------------
-- Examples of gfolds: summing perfect trees, monadicity of lambda 
-- terms and functoriality of nested types.
------------------------------------------------------------------

type PLeafT g h = forall y . forall a . Nat y (Ran g h) -> a -> Ran g h a
type PNodeT g h = forall y . forall a . Nat y (Ran g h) -> y (a, a) -> Ran g h a

gfoldPTree :: PLeafT g h -> PNodeT g h -> PTree a -> Ran g h a
gfoldPTree l n = hfoldPTree (l idNat) (n idNat)

psumL :: PLeafT (Con Int) (Con Int)
psumL pinterp x = Ran (\e -> e x)

psumN :: PNodeT (Con Int) (Con Int)
psumN pinterp x = Ran (\e -> runInterp pinterp x (update e))

update :: (a -> Con Int b) -> (a, a) -> Con Int b
update e (x,y) = e x `cplus` e y where cplus (Con a) (Con b) = Con (a+b)

sumAuxPTree :: PTree a -> Ran (Con Int) (Con Int) a
sumAuxPTree = gfoldPTree psumL psumN

sumPTree :: PTree Int -> Int
sumPTree = icon . fromRan sumAuxPTree . Comp . fmap Con


subAlg :: InterpT HLam (Mu HLam) (Mu HLam)
subAlg k (HVar x)   = Ran (\e -> e x)
subAlg k (HApp t u) = Ran (\e -> In (HApp (runInterp k t e) (runInterp k u e)))
subAlg k (HAbs t)   = Ran (\e -> In (HAbs (runInterp k t (lift e))))

lift e (Just x) = fmap Just (e x)
lift e Nothing  = In (HVar Nothing)

instance Monad (Mu HLam) where
    return  = In . HVar
    t >>= f = rungfold subAlg t f

mapAlg :: HFunctor f => InterpT f Id (Mu f)
mapAlg k t = let k1 t = runInterp k t Id
               in Ran (\e -> In (hfmap k1 (ffmap (unid . e) t)))

instance HFunctor f => Functor (Mu f) where
       fmap k t = rungfold mapAlg t (Id . k)

newtype Id a = Id {unid :: a}

instance Functor Id where
    fmap f (Id x) = Id (f x)


--------------------------------------------------------------------
-- Left Kan Extensions and the natural isomorphism 
-- inherent in their definition
--------------------------------------------------------------------

data Lan g h a = forall b . Lan (g b -> a, h b)

-- This definition of left Kan extension is essentially the exact
-- transcription of the categorical coend formula with the coend replaced
-- by the parametric existential quantifier. 

gbuild :: HFunctor f => (forall x. Alg f x -> Nat (Lan g h) x) -> Nat (Lan g h) (Mu f)
gbuild = hbuild

toLan :: Functor f => Nat h (f `Comp` g) -> Nat (Lan g h) f
toLan s (Lan (val, v)) =  fmap val (icomp (s v))

fromLan :: Nat (Lan g h) f -> Nat h (f `Comp` g)
fromLan s t = Comp (s (Lan (id, t)))

-- the functions toLan and fromLan implement the isomorphism
-- characterising Ran g h


-------------------------------------------------------------
-- An example of fusion using generalised builds with PTrees
------------------------------------------------------------

gbuildPTree :: (forall x. 
                  (forall a. a -> x a) ->
                  (forall a. x (a,a) -> x a) ->
                  (forall a. Lan g h a -> x a)) -> 
                Lan g h a -> PTree a
gbuildPTree g = g PLeaf PNode

-- gfold k . (gbuild l) = l (toAlg k)

tree :: (forall a. a -> x a) -> 
          (forall a. x (a,a) -> x a) -> 
            Nat (Lan (Con Int) (Con Int)) x
tree l n (Lan (z, Con 0)) = l (z (Con 0))
tree l n x                = n (tree l n (count x))

count :: Lan (Con Int) (Con Int) a -> 
            Lan (Con Int) (Con Int) (a,a)
count (Lan (z, Con n)) = 
            Lan (\i -> (z (s i), z (s i)), Con (n-1))

s :: Con Int a -> Con Int a
s (Con n) = Con (n+1)

mkTree :: Int -> PTree Int
mkTree n = gbuildPTree tree (Lan (icon, Con n))

smTree :: Lan (Con Int) (Con Int) a -> Ran (Con Int) (Con Int) a
smTree (Lan (z, Con n)) = if n == 0 then Ran (\e -> e (z (Con 0)))
                                    else Ran (\e -> iran (smTree (count (Lan (z, Con n))))
                                                 (update e))  

------------------------------------------------------------
-- 5. Conclusions 
------------------------------------------------------------

unfold :: Functor f => (x -> f x) -> x -> M f
unfold k x = Inn (fmap (unfold k) (k x))

destroy :: Functor f => (forall x . (x -> f x) -> x -> c) -> M f -> c
destroy g = g outt

outt (Inn t) = t


type CoAlg f g  = Nat g (f g)


hunfold :: (HFunctor f, Functor g) => CoAlg f g -> Nat g (Mu f) 
hunfold k x = In (hfmap (hunfold k) (k x))

hdestroy :: HFunctor f => (forall g. Functor g => CoAlg f g -> Nat g c) -> Nat (Mu f) c
hdestroy g = g (out)


out :: Nat (Mu f) (f (Mu f))
out (In t) = t

-- hdestroy g . hunfold k = g k



--------------------------------------------------------------
-- This research supported in part by NSF award CCF-0700341 --
--------------------------------------------------------------