Specialising Arrowized FRP

test.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE GADTs #-}

module Lib where

import Control.Arrow
import Control.Category

data Step s a = Yield !s a

data SF a b where
  SF :: !(a -> s -> Step s b) -> !s -> SF a b

runSF :: SF a b -> a -> (b, SF a b)
runSF (SF f s) a = case f a s of
  Yield s' b -> (b, (SF f s'))

instance Category SF where
  id = SF (flip Yield) ()
  {-# INLINE id #-}
  SF f2 s2 . SF f1 s1 = SF g (s1, s2)
    where
      g a (s1, s2)
        | Yield s1' b <- f1 a s1
        , Yield s2' c <- f2 b s2
        = Yield (s1', s2') c
  {-# INLINE (Control.Category..) #-}

instance Arrow SF where
  arr f = SF (\a _ -> Yield () (f a)) ()
  {-# INLINE arr #-}
  first (SF f s) = SF g s
    where
      g (b, d) s = case f b s of
        Yield s' c -> Yield s' (c, d)
  {-# INLINE first #-}
  second (SF f s) = SF g s
    where
      g (d, b) s = case f b s of
        Yield s' c -> Yield s' (d, c)
  {-# INLINE second #-}
  SF f1 s1 *** SF f2 s2 = SF g (s1, s2)
    where
      g (b, b') (s1, s2)
        | Yield s1' c  <- f1 b  s1
        , Yield s2' c' <- f2 b' s2
        = Yield (s1', s2') (c, c')
  {-# INLINE (***) #-}
  SF f1 s1 &&& SF f2 s2 = SF g (s1, s2)
    where
      g b (s1, s2)
        | Yield s1' c  <- f1 b s1
        , Yield s2' c' <- f2 b s2
        = Yield (s1', s2') (c, c')
  {-# INLINE (&&&) #-}

instance ArrowChoice SF where
  left (SF f s) = SF g s
    where
      g (Right d) s = Yield s (Right d)
      g (Left b)  s = case f b s of
        Yield s' c -> Yield s' (Left c)
  {-# INLINE left #-}
  right (SF f s) = SF g s
    where
      g (Left d)  s = Yield s (Left d)
      g (Right b) s = case f b s of
        Yield s' c -> Yield s' (Right c)
  {-# INLINE right #-}
  SF f1 s1 +++ SF f2 s2 = SF g (s1, s2)
    where
      g (Left b) (s1, s2)
        | Yield s1' c  <- f1 b  s1
        = Yield (s1', s2) (Left c)
      g (Right b') (s1, s2)
        | Yield s2' c' <- f2 b' s2
        = Yield (s1, s2') (Right c')
  {-# INLINE (+++) #-}
  SF f1 s1 ||| SF f2 s2 = SF g (s1, s2)
    where
      g (Left b) (s1, s2)
        | Yield s1' c  <- f1 b  s1
        = Yield (s1', s2) c
      g (Right b') (s1, s2)
        | Yield s2' c  <- f2 b' s2
        = Yield (s1, s2') c
  {-# INLINE (|||) #-}

inc :: SF Int Int
inc = SF go ()
  where
    go a _ = let !b = a+1 in Yield () b

double :: SF Int Int
double = arr (*2)

test :: SF a b -> _
test a b c d e f = first (a ||| b) >>> c *** second (d ||| e) >>> f

test2 :: SF Int Int
test2 = (double &&& inc) >>> arr (uncurry (+))