TeeMachineExample.hs

{-# LANGUAGE RankNTypes #-}
import Control.Monad
import Control.Monad.State
import Control.Monad.Writer
import Data.Char (intToDigit)
import Data.Machine
import Data.Machine.Plan
import Data.Machine.Source
import Data.Monoid

-- | Produces integers from a list.
m1 :: [Int] -> Source Int
m1 = source

-- | Produces characters from a list.
m2 :: [Char] -> Source Char
m2 = source

-- | Reads a number from its left input. Then reads this many
-- characters from its right input. Outputs the resulting string,
-- together with the number of strings produced so far.
m3 :: Tee Int Char (Int, String)
m3 = construct (loop 0)
  where
    -- `loop` keeps internal state - the count of strings
    -- produced so far.
    loop count = do
        -- Read a number from L.
        n <- awaits L
        -- Read this many characters from L.
        s <- replicateM n (awaits R)
        let count' = count + 1
        -- Output the result.
        yield (count', s)
        loop count'

m3State :: TeeT (State Int) Int Char (Int, String)
m3State = repeatedly $ do
        n <- awaits L
        s <- replicateM n (awaits R)
        lift (modify (+ 1))
        count <- lift get
        yield (count, s)

m3Writer :: TeeT (Writer (Sum Int)) Int Char String
m3Writer = repeatedly $ do
        n <- awaits L
        s <- replicateM n (awaits R)
        tell (Sum 1)
        yield s

input :: (Monad m) => TeeT m Int Char c -> TeeT m a b c
input = tee (m1 [2,3,4,5]) (m2 "Lorem ipsum dolor sit amet")

main = do
    print . run $ input m3
    print . flip evalState 0 . runT $ input m3State
    print . runWriter . runT $ input m3Writer