Transducers with explicit local state

T.hs

{-# LANGUAGE GADTs, RankNTypes #-}

import Control.Applicative
import Control.Category
import Data.List (foldl')
import Data.Profunctor
import Prelude hiding ((.), id)

-- | Explicit state-passing Moore machine
data Moore i o where
  Moore :: (i -> x -> x) -> x -> (x -> o) -> Moore i o

instance Profunctor Moore where
  dimap f g (Moore ixx x xo) = Moore (ixx . f) x (g . xo)

instance Functor (Moore i) where fmap = rmap

-- | Strict pair since we'll be doing strict left folds
data Two a b = Two !a !b

instance Applicative (Moore i) where
  pure = Moore (const id) () . const 
  Moore ixx x xf <*> Moore iyy y ya =
    let izz i (Two x y) = Two (ixx i x) (iyy i y)
        z               = Two x y
        zb    (Two x y) = (xf x) (ya y)
    in Moore izz z zb

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

-- I'm ignoring early termination for now

newtype T a b = T (forall r . Moore b r -> Moore a r)

instance Profunctor T where
  dimap f g (T t) = T (lmap f . t . lmap g)

instance Functor (T i) where fmap = rmap

instance Category T where
  id = T id
  T f . T g = T (g . f) -- composes the "right way" despite CPS

-- Is there an Arrow instance?

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

tmap :: (a -> b) -> T a b
tmap f = T (lmap f)

tfilter :: (a -> Bool) -> T a a
tfilter p = T $ \(Moore axx x xr) ->
  Moore (\a -> if p a then axx a else id) x xr

tflatMap :: (a -> [b]) -> T a b
tflatMap f = T $ \(Moore bxx x xr) ->
  Moore (\a x -> foldl' (flip bxx) x (f a)) x xr

ttake :: Int -> T a a
ttake n = T $ \(Moore axx x xr) ->
  let ayy a y@(Two x n)
        | n > 0     = Two (axx a x) (n-1)
        | otherwise = y
      y               = Two x n
      yr    (Two x _) = xr x
  in Moore ayy y yr

-- Note that this generalizes
tseq :: T a b -> ([a] -> [b])
tseq (T t) as =
  case t m of
    Moore cons x xr -> 
      -- a way to feed the input
      xr (foldl' (flip cons) x as)
  where
    -- a machine for executing the output
    -- could easily be effectful
    -- 
    -- here we build lists "from the left" using
    -- difference lists. This is probably the most
    -- complex thing in this entire module
    m = Moore (\a f -> f . (a:)) id ($ [])

data Three a b c = Three !a !b !c

tpartition :: Int -> T a [a]
tpartition n = T $ \(Moore asxx x xr) ->
  -- the difference list trick would work here too
  -- but I demonstrate the usual accumulate/reverse
  -- mechanism instead
  let ayy a (Three x m acc)
        | m == 0    = Three (asxx (reverse acc) x) n []
        | otherwise = Three x (m-1) (a:acc)
      y                     = Three x n []
      yr    (Three x m acc) = xr x
  in Moore ayy y yr