Free Arrow

FreeArrow.hs

{-# LANGUAGE KindSignatures, GADTs, RankNTypes, ScopedTypeVariables #-}
module Main where

import Prelude hiding (id, (.))

import Control.Arrow
import Control.Category
import Data.Kind (Type)
import Data.Function ((&))
import Data.Profunctor


data FreeArrow (g :: Type -> Type -> Type) a b where
  Nil
    :: (a -> b)
    -> FreeArrow g a b
  Cons
    :: (a -> (b, rest))
    -> g b c
    -> FreeArrow g (c, rest) d
    -> FreeArrow g a d

instance Profunctor (FreeArrow g) where
  lmap f1 (Nil f2)
    = Nil (f1 >>> f2)
  lmap f1 (Cons f2 g2 gs2)
    = Cons (f1 >>> f2) g2 gs2
  rmap = flip go
    where
      go
        :: FreeArrow g a b
        -> (b -> c)
        -> FreeArrow g a c
      go (Nil f1) f2
        = Nil (f1 >>> f2)
      go (Cons f1 g1 gs1) f2
        = Cons f1 g1 (go gs1 f2)

instance Category (FreeArrow g) where
  id = Nil id
  (.) = flip go
    where
      go
        :: forall a b c
         . FreeArrow g a b
        -> FreeArrow g b c
        -> FreeArrow g a c
      go (Nil f1) gs2
        = lmap f1 gs2
      go (Cons f1 g gs1) gs2
        = Cons f1 g (go gs1 gs2)

assocL
  :: (a, (b, c))
  -> ((a, b), c)
assocL (a,(b,c))
  = ((a,b),c)

assocR
  :: ((a, b), c)
  -> (a, (b, c))
assocR ((a,b),c)
  = (a,(b,c))

instance Arrow (FreeArrow g) where
  arr f
    = Nil f

  first (Nil f)
    = Nil (first f)
  first (Cons f0 g0 gs0)
    = let (f0', gs0') = go (f0, gs0)
      in Cons f0' g0 gs0'
    where
      go
        :: forall a b c d rest x
         . ( (a -> (b, rest))
           , FreeArrow g (c, rest) d
           )
        -> ( ((a,x) -> (b, (rest, x)))
           , FreeArrow g (c, (rest, x)) (d, x)
           )
      go (f, gs)
        = (f', gs')
        where
          f' :: ((a,x) -> (b, (rest, x)))
              = -- (a, x)
                first f
            >>> -- ((b, rest), x)
                assocR
                -- (b, (rest, x))

          gs' :: FreeArrow g (c, (rest, x)) (d, x)
             = gs
             & -- FreeArrow g (c, rest) d
               first
               -- FreeArrow g ((c, rest), x) (d, x)
             & lmap assocL
               -- FreeArrow g (c, (rest, x)) (d, x)

runFreeArrow
  :: Arrow g
  => FreeArrow g a b
  -> g a b
runFreeArrow (Nil f)
  = arr f
runFreeArrow (Cons f g gs)
    = -- a
      arr f
  >>> -- (b, rest)
      first g
  >>> -- (c, rest)
      runFreeArrow gs
      -- d