SequenceId.hs

{-# LANGUAGE BangPatterns #-}

module System.IO.Streams.SequenceId
       ( sequenceIdInputStream
       , sequenceIdOutputStream
       ) where

import           Control.Applicative ((<$>))
import           Data.IORef          (newIORef, readIORef, writeIORef)
import           Data.SequenceId     (SequenceId, SequenceIdError, checkSeqId,
                                      incrementSeqId)
import           System.IO.Streams   (InputStream, OutputStream)
import qualified System.IO.Streams   as Streams


------------------------------------------------------------------------------
-- | Wrap an 'System.IO.Streams.InputStream' and check for dropped or duplicated sequence IDs.
--
-- Example:
--
-- @
-- ghci> 'System.IO.Streams.fromList' [1..10 :: 'Data.SequenceId.SequenceId'] >>= 'sequenceIdInputStream' 0 'id' ('fail' . 'show') >>= 'System.IO.Streams.toList'
-- [1,2,3,4,5,6,7,8,9,10]
--
-- ghci> 'System.IO.Streams.fromList' [5..10 :: 'Data.SequenceId.SequenceId'] >>= 'sequenceIdInputStream' 0 'id' ('fail' . 'show') >>= 'System.IO.Streams.toList'
-- *** Exception: user error ('Data.SequenceId.SequenceIdError' {errType = 'Data.SequenceId.SequenceIdDropped', lastSeqId = 0, currSeqId = 5})
--
-- ghci> 'System.IO.Streams.fromList' [1..10 :: 'Data.SequenceId.SequenceId'] >>= 'sequenceIdInputStream' 5 'id' ('fail' . 'show') >>= 'System.IO.Streams.toList'
-- *** Exception: user error ('Data.SequenceId.SequenceIdError' {errType = 'Data.SequenceId.SequenceIdDuplicated', lastSeqId = 5, currSeqId = 1})
-- @
sequenceIdInputStream :: SequenceId                 -- ^ Initial sequence ID
                      -> (a -> SequenceId)          -- ^ Function applied to each element of the stream to get the sequence ID
                      -> (SequenceIdError -> IO ()) -- ^ Error handler
                      -> InputStream a              -- ^ 'System.IO.Streams.InputStream' to check the sequence of
                      -> IO (InputStream a)         -- ^ Pass-through of the given stream
sequenceIdInputStream initSeqId getSeqId seqIdFaultHandler inStream =
    fst <$> Streams.inputFoldM f initSeqId inStream
  where
    f lastSeqId x = do
        let currSeqId = getSeqId x
        maybe (return ()) seqIdFaultHandler $ checkSeqId lastSeqId currSeqId
        return $ max currSeqId lastSeqId


------------------------------------------------------------------------------
-- | Wrap an 'System.IO.Streams.OutputStream' to give a sequence ID for each element written.
--
-- Example:
--
-- @
-- (outStream', getSeqId) <- 'sequenceIdOutputStream' 1 outStream
-- return $ 'System.IO.Streams.Combinators.contramapM' (addSeqId getSeqId) outStream'
-- @
sequenceIdOutputStream :: SequenceId                         -- ^ Initial sequence ID
                       -> OutputStream a                     -- ^ 'System.IO.Streams.OutputStream' to count the elements of
                       -> IO (OutputStream a, IO SequenceId) -- ^ ('IO' 'SequenceId') is the action to run to get the current sequence ID
sequenceIdOutputStream = outputFoldM count
  where count a _ = return $ incrementSeqId a


outputFoldM :: (a -> b -> IO a)           -- ^ fold function
            -> a                          -- ^ initial seed
            -> OutputStream b             -- ^ output stream
            -> IO (OutputStream b, IO a)  -- ^ returns a new stream as well as
                                          -- an IO action to fetch the updated
                                          -- seed value.
outputFoldM f initial stream = do
    ref <- newIORef initial
    os  <- Streams.makeOutputStream (wr ref)
    return (os, readIORef ref)

  where
    wr _ Nothing       = Streams.write Nothing stream
    wr ref mb@(Just x) = do
        !z  <- readIORef ref
        !z' <- f z x
        writeIORef ref z'
        Streams.write mb stream