Binary.hs


-----------------------------------------------------------------------------
-- |
-- Module: Data.Enumerator.Binary
-- Copyright: 2010 John Millikin
-- License: MIT
--
-- Maintainer: jmillikin@gmail.com
-- Portability: portable
--
-- This module is intended to be imported qualified:
--
-- @
-- import qualified Data.Enumerator.Binary as EB
-- @
--
-- Since: 0.4.5
--
-----------------------------------------------------------------------------

module Data.Enumerator.Binary (

	  -- * Binary IO
	  enumHandle
	, enumFile
	, iterHandle

	-- * List analogues
	, Data.Enumerator.Binary.head
	, Data.Enumerator.Binary.drop
	, Data.Enumerator.Binary.dropWhile
	, Data.Enumerator.Binary.take
	, Data.Enumerator.Binary.takeWhile
	, Data.Enumerator.Binary.consume
	, require
	, isolate

	) where
import Prelude hiding (head, drop, takeWhile)
import Data.Enumerator hiding (head, drop)
import qualified Data.ByteString as B

import Data.Enumerator.Util (tryStep)
import qualified Control.Exception as Exc
import Control.Monad.IO.Class (MonadIO)
import qualified System.IO as IO
import System.IO.Error (isEOFError)

import Data.Word (Word8)
import qualified Data.ByteString.Lazy as BL


-- | Read bytes (in chunks of the given buffer size) from the handle, and
-- stream them to an 'Iteratee'. If an exception occurs during file IO,
-- enumeration will stop and 'Error' will be returned. Exceptions from the
-- iteratee are not caught.
--
-- This enumerator blocks until at least one byte is available from the
-- handle, and might read less than the maximum buffer size in some
-- cases.
--
-- The handle should be opened with no encoding, and in 'IO.ReadMode' or
-- 'IO.ReadWriteMode'.

enumHandle :: MonadIO m
           => Integer -- ^ Buffer size
           -> IO.Handle
           -> Enumerator B.ByteString m b
enumHandle bufferSize h = loop where
	loop (Continue k) = withBytes $ \bytes ->
		if B.null bytes
			then continue k
			else k (Chunks [bytes]) >>== loop
	
	loop step = returnI step
	
	intSize = fromInteger bufferSize
	withBytes = tryStep $ do
		hasInput <- Exc.catch
			(IO.hWaitForInput h (-1))
			(\err -> if isEOFError err
				then return False
				else Exc.throwIO err)
		if hasInput
			then B.hGetNonBlocking h intSize
			else return B.empty


-- | Opens a file path in binary mode, and passes the handle to 'enumHandle'.
-- The file will be closed when the 'Iteratee' finishes.

enumFile :: FilePath -> Enumerator B.ByteString IO b
enumFile path = enum where
	withHandle = tryStep (IO.openBinaryFile path IO.ReadMode)
	enum step = withHandle $ \h -> do
		Iteratee $ Exc.finally
			(runIteratee (enumHandle 4096 h step))
			(IO.hClose h)


-- | Read bytes from a stream and write them to a handle. If an exception
-- occurs during file IO, enumeration will stop and 'Error' will be
-- returned.
--
-- The handle should be opened with no encoding, and in 'IO.WriteMode' or
-- 'IO.ReadWriteMode'.

iterHandle :: MonadIO m => IO.Handle
           -> Iteratee B.ByteString m ()
iterHandle h = continue step where
	step EOF = yield () EOF
	step (Chunks []) = continue step
	step (Chunks bytes) = let
		put = mapM_ (B.hPut h) bytes
		in tryStep put (\_ -> continue step)

toChunks :: BL.ByteString -> Stream B.ByteString
toChunks = Chunks . BL.toChunks


-- | Get the next byte from the stream, or 'Nothing' if the stream has
-- ended.
--
-- Since: 0.4.5

head :: Monad m => Iteratee B.ByteString m (Maybe Word8)
head = continue loop where
	loop (Chunks xs) = case BL.uncons (BL.fromChunks xs) of
		Just (char, extra) -> yield (Just char) (toChunks extra)
		Nothing -> head
	loop EOF = yield Nothing EOF


-- | @drop n@ ignores /n/ bytes of input from the stream.
--
-- Since: 0.4.5

drop :: Monad m => Integer -> Iteratee B.ByteString m ()
drop n | n <= 0 = return ()
drop n = continue (loop n) where
	loop n' (Chunks xs) = iter where
		lazy = BL.fromChunks xs
		len = toInteger (BL.length lazy)
		iter = if len < n'
			then drop (n' - len)
			else yield () (toChunks (BL.drop (fromInteger n') lazy))
	loop _ EOF = yield () EOF


-- | @dropWhile p@ ignores input from the stream until the first byte which
-- does not match the predicate.
--
-- Since: 0.4.5

dropWhile :: Monad m => (Word8 -> Bool) -> Iteratee B.ByteString m ()
dropWhile p = continue loop where
	loop (Chunks xs) = iter where
		lazy = BL.dropWhile p (BL.fromChunks xs)
		iter = if BL.null lazy
			then continue loop
			else yield () (toChunks lazy)
	loop EOF = yield () EOF


-- | @take n@ extracts the next /n/ bytes from the stream, as a lazy
-- ByteString.
--
-- Since: 0.4.5

take :: Monad m => Integer -> Iteratee B.ByteString m BL.ByteString
take n | n <= 0 = return BL.empty
take n = continue (loop id n) where
	loop acc n' (Chunks xs) = iter where
		lazy = BL.fromChunks xs
		len = toInteger (BL.length lazy)
		
		iter = if len < n'
			then continue (loop (acc . (BL.append lazy)) (n' - len))
			else let
				(xs', extra) = BL.splitAt (fromInteger n') lazy
				in yield (acc xs') (toChunks extra)
	loop acc _ EOF = yield (acc BL.empty) EOF


-- | @takeWhile p@ extracts input from the stream until the first byte which
-- does not match the predicate.
--
-- Since: 0.4.5

takeWhile :: Monad m => (Word8 -> Bool) -> Iteratee B.ByteString m BL.ByteString
takeWhile p = continue (loop id) where
	loop acc (Chunks []) = continue (loop acc)
	loop acc (Chunks xs) = iter where
		lazy = BL.fromChunks xs
		(xs', extra) = BL.span p lazy
		iter = if BL.null extra
			then continue (loop (acc . (BL.append lazy)))
			else yield (acc xs') (toChunks extra)
	loop acc EOF = yield (acc BL.empty) EOF


-- | Read all remaining input from the stream, and return as a lazy
-- ByteString.
--
-- Since: 0.4.5

consume :: Monad m => Iteratee B.ByteString m BL.ByteString
consume = continue (loop id) where
	loop acc (Chunks []) = continue (loop acc)
	loop acc (Chunks xs) = iter where
		lazy = BL.fromChunks xs
		iter = continue (loop (acc . (BL.append lazy)))
	loop acc EOF = yield (acc BL.empty) EOF


-- | @require n@ buffers input until at least /n/ bytes are available, or
-- throws an error if the stream ends early.
--
-- Since: 0.4.5

require :: Monad m => Integer -> Iteratee B.ByteString m ()
require n | n <= 0 = return ()
require n = continue (loop id n) where
	loop acc n' (Chunks xs) = iter where
		lazy = BL.fromChunks xs
		len = toInteger (BL.length lazy)
		iter = if len < n'
			then continue (loop (acc . (BL.append lazy)) (n' - len))
			else yield () (toChunks (acc lazy))
	loop _ _ EOF = throwError (Exc.ErrorCall "require: Unexpected EOF")


-- | @isolate n@ reads at most /n/ bytes from the stream, and passes them
-- to its iteratee. If the iteratee finishes early, bytes continue to be
-- consumed from the outer stream until /n/ have been consumed.
--
-- Since: 0.4.5

isolate :: Monad m => Integer -> Enumeratee B.ByteString B.ByteString m b
isolate n step | n <= 0 = return step
isolate n (Continue k) = continue loop where
	loop (Chunks []) = continue loop
	loop (Chunks xs) = iter where
		lazy = BL.fromChunks xs
		len = toInteger (BL.length lazy)
		
		iter = if len <= n
			then k (Chunks xs) >>== isolate (n - len)
			else let
				(s1, s2) = BL.splitAt (fromInteger n) lazy
				in k (toChunks s1) >>== (\step -> yield step (toChunks s2))
	loop EOF = k EOF >>== (\step -> yield step EOF)
isolate n step = drop n >> return step