Composing F-algebras

Algebra.hs

{-# LANGUAGE GADTs, RankNTypes, TypeOperators, TupleSections #-}
module Algebra where

import Control.Category
import Control.Arrow
import Control.Applicative
import Label.Simple (get)
import Prelude hiding ((.), id)

import Generics.Combinator

-------------------------------------------------------------------------------

-- | Generic algebra type.

type Alg i -- ^ An additional value used for algebra composition.
         f -- ^ The recursive functor to decompose into a result.
         a -- ^ Additional input for the algebra.
         b -- ^ Result of running the algebra.
  = a -> f (i, b) -> (i, b)

data Algebra f a b where
  Algebra :: Alg i f a b -> Algebra f a b

algebra :: Functor f => (a -> f b -> b) -> Algebra f a b
algebra alg = Algebra (\a -> (,) () . alg a . fmap snd)

type Phi f b = Algebra f () b

phi :: Functor f => (f b -> b) -> Phi f b
phi f = algebra (const f)

-------------------------------------------------------------------------------

instance Functor f => Category (Algebra f) where
  id                    = Algebra (const . (,) ())
  Algebra a . Algebra b = Algebra (compose a b)

instance Functor f => Functor (Algebra f a) where
  fmap f (Algebra g) = Algebra (amap f g)

instance Functor f => Applicative (Algebra f a) where
  pure = Algebra . const . const . (,) ()
  Algebra f <*> Algebra g = Algebra (ap f g)

instance Functor f => Arrow (Algebra f) where
  arr f = Algebra (const . (,) () . f)
  first (Algebra f) = Algebra (\(b, d) -> second (, d) . f b . fmap (second fst))

instance Functor f => ArrowLoop (Algebra f) where
  loop (Algebra f) = Algebra $ \b inp ->
    let (i, (c, d)) = f (b, d) (second (, d) <$> inp) in (i, c)

compose :: Functor f => Alg j f b c -> Alg i f a b -> Alg (j, (i, b)) f a c
compose f g a input =
  let (i, b) = g a (prj_g <$> input)
      (j, c) = f b (prj_f <$> input)
   in ((j, (i, b)), c)
 where prj_f ((j, (_, _)), c) = (j, c)
       prj_g ((_, (i, b)), _) = (i, b)

ap :: Functor f => Alg j f c (a -> b) -> Alg i f c a -> Alg ((i, a), (j, a -> b)) f c b
ap f g c input =
  let (i, a ) = g c (prj_g <$> input)
      (j, ab) = f c (prj_f <$> input)
   in (((i, a), (j, ab)), ab a)
 where prj_f (((_, _), (j, ab)), _) = (j, ab)
       prj_g (((i, a), (_, _ )), _) = (i, a)

amap :: Functor f => (a -> b) -> Alg i f c a -> Alg (i, a) f c b
amap fn f c input =
    let (i, a) = f c (prj <$> input)
     in ((i, a), fn a)
    where prj ((i, a), _) = (i, a)

-------------------------------------------------------------------------------

-- Running algebras.

para :: Functor f => Algebra f (a, Fix f) b -> a -> Fix f -> b
para (Algebra alg) a = snd . recurse
  where recurse g = alg (a, g) (recurse <$> get out g)

cata :: Functor f => Algebra f () b -> Fix f -> b
cata (Algebra alg) = para (Algebra alg .  arr fst) ()

-- Like para but without an additional input.

para1 :: Functor f => Algebra f (Fix f) b -> Fix f -> b
para1 alg = para (alg . arr snd) ()