Initial algebras and final coalgebras

algebras.hs

{-# LANGUAGE RankNTypes, DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances, TypeSynonymInstances #-}
module Main where
import Criterion.Main (defaultMain, bench, bgroup, nf)

--------------------------------------------------------------------------------
-- Fixed points of a functor

newtype Mu f = Mu { muF :: f (Mu f) }

--------------------------------------------------------------------------------
-- Initial algebras

type Algebra f a = f a -> a

-- | This is a generic catamorphism.
ffold :: Functor f => Algebra f a -> Mu f -> a
ffold h = go h where go g = g . fmap (go g) . muF
--ffold h = h . fmap (ffold h) . muF     -- This is a slow definition
{-# INLINE[1] ffold #-}

fbuild :: Functor f => (forall b. Algebra f b -> b) -> Mu f
fbuild g = g Mu
{-# INLINE[1] fbuild #-}

--------------------------------------------------------------------------------
-- Final coalgebras

type Coalgebra f a = a -> f a

-- | This is a generic anamorphism.
funfold :: Functor f => Coalgebra f b -> b -> Mu f
funfold h = go h where go g = Mu . fmap (go g) . g
--funfold h = Mu . fmap (funfold h) . h   -- This is a slow definition
{-# INLINE[1] funfold #-}

funbuild :: Functor f => (Coalgebra f (Mu f) -> b) -> b
funbuild g = g muF
{-# INLINE[1] funbuild #-}

--------------------------------------------------------------------------------
-- Meta/hylomorphisms, and RULE optimizations

hylo :: Functor f => Algebra f c -> Coalgebra f a -> a -> c
hylo f h = ffold f . funfold h
{-# INLINE[2] hylo #-}

meta :: Functor f => Coalgebra f b -> Algebra f b -> Mu f -> Mu f
meta f h = funfold f . ffold h
{-# INLINE[2] meta #-}

{-# RULES
-- Builder rule for catamorphisms
"ffold/fbuild"      forall f (g :: forall b. Algebra f b -> b).
                       ffold f (fbuild g) = g f

-- Builder rule for anamorphisms
"funfold/funbuild"  forall f g.
                       funfold f (funbuild g) = g f

-- Hylomorphism rule
"ffold/funfold"     forall f g x.
                       ffold f (funfold g x) = f (g x)

-- Metamorphism rule
"funfold/ffold"     forall f g x.
                       funfold f (ffold g x) = g (f x)
  #-}

--------------------------------------------------------------------------------
-- Naturals

-- μX. 1 + X
data NatF f = Z | S f
  deriving (Eq, Show, Functor)
type Nat = Mu NatF

instance Eq Nat where
  (Mu f) == (Mu g) = f == g

addN :: Nat -> Nat -> Nat
addN (Mu Z) m     = m
addN (Mu (S f)) m = Mu $ S (addN f m)
{-# INLINE addN #-}

toI :: Nat -> Int
toI = ffold g
  where g Z     = 0
        g (S x) = 1 + x
{-# INLINE toI #-}

nat1 :: Nat
nat1 = s (s (s z))
  where z   = Mu Z
        s x = Mu (S x)

nat2 :: Nat
nat2 = s (s z)
  where z   = Mu Z
        s x = Mu (S x)

--------------------------------------------------------------------------------
-- Lists

-- μX. 1 + A * X
data ListF a f = Nil | Cons a f
  deriving (Eq, Show, Functor)
type List a = Mu (ListF a)

instance Eq a => Eq (List a) where
  (Mu f) == (Mu g) = f == g

lengthL :: List a -> Int
lengthL = ffold g
  where g Nil        = 0
        g (Cons _ f) = 1 + f
{-# INLINE lengthL #-}

mapL :: (a -> b) -> List a -> List b
mapL f = ffold g
  where g Nil        = Mu Nil
        g (Cons a x) = Mu (Cons (f a) x)
{-# INLINE mapL #-}

-- | Different implementation via fbuild, which should fuse
-- with other List functions.
mapL2 :: (a -> b) -> List a -> List b
mapL2 f xs = fbuild (\h -> ffold (h . g) xs)
  where g Nil        = Nil
        g (Cons a x) = Cons (f a) x
{-# INLINE mapL2 #-}

-- | Different implementation via hylomorphisms.
mapL3 :: (a -> b) -> List a -> List b
mapL3 f = hylo g muF
  where
    g Nil        = Mu Nil
    g (Cons a x) = Mu $ Cons (f a) x
{-# INLINE mapL3 #-}

list1 :: List Int
list1 = c 1 (c 2 (c 3 (c 4 (c 5 (c 6 (c 7 (c 8 n)))))))
  where n     = Mu Nil
        c a r = Mu (Cons a r)

--------------------------------------------------------------------------------
-- Trees

-- μX. A + A*X*X
data TreeF a f = Leaf a | Tree a f f
  deriving (Eq, Show, Functor)
type Tree a = Mu (TreeF a)

instance Eq a => Eq (Tree a) where
  (Mu f) == (Mu g) = f == g

depthT :: Tree a -> Int
depthT = ffold g
  where g (Leaf _)     = 0
        g (Tree _ l r) = 1 + max l r
{-# INLINE depthT #-}

sumT :: Tree a -> Int
sumT = ffold g
  where g (Leaf _)     = 1
        g (Tree _ l r) = l + r
{-# INLINE sumT #-}

tree1 :: Tree Int
tree1 = tree 7 (tree 5 (leaf 2) (leaf 6))
               (tree 9 (leaf 8)
                       (tree 11 (leaf 10)
                                (leaf 13)))
  where
    tree x l r = Mu (Tree x l r)
    leaf x     = Mu (Leaf x)

--------------------------------------------------------------------------------
-- Driver and benchmarks

main :: IO ()
main = defaultMain
  [ bgroup "nat"
      [ bench "add1" $ nf toI (addN nat1 nat2)
      , bench "add2" $ nf toI (addN nat2 nat1)
      ]
  , bgroup "list"
      [ bench "length" $ nf lengthL list1
      , bench "map"    $ nf lengthL $ mapL (+1) list1
      , bench "map2"   $ nf lengthL $ mapL2 (+1) list1
      , bench "map3"   $ nf lengthL $ mapL3 (+1) list1
      ]
  , bgroup "tree"
      [ bench "depth" $ nf depthT tree1
      , bench "sum"   $ nf sumT   tree1
      ]
  , bgroup "rewrite"
      [ bench "length/map[ffold/ffold]"  $ nf (lengthL . mapL (+1)) list1
        -- The ffold/fbuild rule cuts runtime in half.
      , bench "length/map[ffold/fbuild]" $ nf (lengthL . mapL2 (+1)) list1
      ]
  ]