Coyoneda lemma & fmap fusion

coyo.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}

import Data.Monoid
import System.Environment

data Coyoneda f a where
  Coyoneda :: (b -> a) -> f b -> Coyoneda f a

instance Functor (Coyoneda f) where
  fmap f (Coyoneda b2a fb) = Coyoneda (f . b2a) fb

coyo :: f a -> Coyoneda f a
coyo = Coyoneda id

uncoyo :: Functor f => Coyoneda f a -> f a
uncoyo (Coyoneda b2a fb) = fmap b2a fb

withCoyo :: Functor f
         => (Coyoneda f a -> Coyoneda f b)
         -> f a
         -> f b
withCoyo f = uncoyo . f . coyo

{-
uncoyo . fmap f . fmap g . coyo
  = uncoyo . fmap f . fmap g . Coyoneda id
  = uncoyo . fmap f . Coyoneda (g . id)
  = uncoyo . fmap f . Coyoneda g
  = uncoyo . Coyoneda (f . g)
  = fmap (f . g)
-}

data Tree a = Bin a (Tree a) (Tree a)
            | Nil
            deriving (Eq, Show)

instance Functor Tree where
  fmap _ Nil         = Nil
  fmap f (Bin a l r) = Bin (f a) (fmap f l) (fmap f r)

instance Foldable Tree where
  foldMap _ Nil = mempty
  foldMap f (Bin a l r) = f a <> foldMap f l <> foldMap f r

sumTree :: Num a => Tree a -> a
sumTree = getSum . foldMap Sum

t :: Tree Integer
t = go 1

  where go r = Bin r (go (2*r)) (go (2*r + 1))

takeDepth :: Int -> Tree a -> Tree a
takeDepth _ Nil = Nil
takeDepth 0 _   = Nil
takeDepth d (Bin r t1 t2) = Bin r (takeDepth (d-1) t1) (takeDepth (d-1) t2)

transform :: (Functor f, Num a) => f a -> f a
transform = fmap (^2) . fmap (+1) . fmap (*2)

main :: IO ()
main = getArgs >>= \args -> case args of
  [k] -> print . sumTree . takeDepth (read k) $ transform t
  [k, "--coyo"] ->
    print . sumTree . takeDepth (read k) $ withCoyo transform t
  _ -> error "wrong arguments"

Derive-pragmas at the top are redundant.