Main.hs

module Main where

import Control.Monad

data P i o a
  = Input (i -> P i o a)
  | Output o (P i o a)
  | Return a

instance Functor (P i o) where
  fmap = liftM

instance Applicative (P i o) where
  pure = return
  (<*>) = ap

instance Monad (P i o) where
  return = Return
  (Input h)     >>= f = Input (\i -> (h i) >>= f)
  (Output o r)  >>= f = Output o (r >>= f)
  (Return a)    >>= f = f a   

{-# INLINE[1] mergeP #-}
mergeP :: P i m a -> P m o a -> P i o a
mergeP p q = merge_PL q p where
  merge_PL :: P m o a -> P i m a -> P i o a
  merge_PL (Input h)     p = merge_PR p h
  merge_PL (Output o r)  p = Output o (merge_PL r p)
  merge_PL (Return a)    p = Return a
  merge_PR :: P i m a -> (m -> P m o a) -> P i o a
  merge_PR (Input f)        h  = Input (\v -> merge_PR (f v) h)
  merge_PR (Output o r)     h  = merge_PL (h o) r
  merge_PR (Return a)       h  = Return a
  
{-# INLINE[1] yieldP #-}
yieldP :: o -> P i o ()
yieldP o = Output o (Return ())

{-# INLINE[1] forP #-}
forP :: P i m a -> (m -> P i o ()) -> P i o a
forP p0 f = go p0 where
  go (Input h) = Input (\x -> go (h x))
  go (Output o r) = f o >> go r
  go (Return a) = Return a

{-# INLINE[1] mapP #-}
mapP :: (i -> o) -> P i o a
mapP f =
  Input $ \i ->
  Output (f i) $
  mapP f

upfromP :: Int -> P i Int a
upfromP n = do
  yieldP n
  upfromP (n + 1)

testP :: P i Int ()
testP = upfromP 0 `mergeP` mapP (+1) `mergeP` mapP (+1)

{-# RULES
    "r1" forall p f .
      p `mergeP` mapP f = forP p (\y -> yieldP (f y))
  ; "r2" forall x f .
      forP (yieldP x) f = f x
  ; "r3" forall p f g .
      forP (forP p f) g = forP p (\x -> forP (f x) g)
  #-}

newtype ContPipe i o a = MakePipe { runPipe :: (a -> ResultCP i o) -> ResultCP i o }
type ResultCP i o = InCont i -> OutCont o -> IO ()
newtype InCont i = MakeInCont { resume_I :: OutCont i -> IO () }
newtype OutCont o = MakeOutCont { resume_O :: o -> InCont o -> IO () }

instance Functor (ContPipe i o) where
  fmap = liftM

instance Applicative (ContPipe i o) where
  pure = return
  (<*>) = ap

instance Monad (ContPipe i o) where
  return a  = MakePipe (\k -> k a)
  p >>= f   = MakePipe (\k -> runPipe p (\x -> runPipe (f x) k))

return_CP :: a -> ContPipe i o a
return_CP a = MakePipe (\k k_i k_o -> k a k_i k_o)

input_CP :: ContPipe i o i
input_CP = MakePipe (\k k_i k_o -> resume_I k_i (MakeOutCont (\i k'_i -> k i k'_i k_o)))

{-# INLINE[1] yield_CP #-}
yield_CP :: o -> ContPipe i o ()
yield_CP o = MakePipe (\k k_i k_o -> resume_O k_o o (MakeInCont (\k'_o -> k () k_i k'_o)))

{-# INLINE[1] merge_CP #-}
merge_CP :: ContPipe i m a -> ContPipe m o a -> ContPipe i o a
merge_CP p q = MakePipe (\k k_i k_o ->
  runPipe q err (MakeInCont (\k'_o -> runPipe p err k_i k'_o)) k_o)
  where err = error "terminated"

{-# INLINE[1] map_CP #-}
map_CP :: (i -> o) -> ContPipe i o a
map_CP f = do
  i <- input_CP
  yield_CP (f i)
  map_CP f

{-# INLINE[1] for_CP #-}
for_CP :: ContPipe i m a -> (m -> ContPipe m o ()) -> ContPipe i o a
for_CP p0 f = p0 `merge_CP` go where
  go = do
    a <- input_CP
    f a
    go

upfrom_CP :: Int -> ContPipe i Int a
upfrom_CP n = do
  yield_CP n
  upfrom_CP (n + 1)

test_CP :: ContPipe i Int ()
test_CP = upfrom_CP 0 `merge_CP` map_CP (+1) `merge_CP` map_CP (+1)

{-# RULES
    "rCP1" forall p f .
      p `merge_CP` map_CP f = for_CP p (\y -> yield_CP (f y))
  ; "rCP2" forall x f .
      for_CP (yield_CP x) f = f x
  ; "rCP3" forall p f g .
      for_CP (for_CP p f) g = for_CP p (\x -> for_CP (f x) g)
  #-}