-----------------------------------------------------------------------------
-- |
-- Module: Data.Enumerator.IO
-- Copyright: 2010 John Millikin
-- License: MIT
--
-- Maintainer: jmillikin@gmail.com
-- Portability: portable
--
-- Enumerator-based IO
--
-----------------------------------------------------------------------------
module Data.Enumerator.IO
	( enumHandle
	, enumFile
	, iterHandle
	) where
import Data.Enumerator
import Data.Enumerator.Util
import qualified Control.Exception as E
import qualified Data.ByteString as B
import qualified System.IO as IO
-- | 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 :: Integer -- ^ Buffer size
           -> IO.Handle
           -> Enumerator E.SomeException B.ByteString IO b
enumHandle bufferSize h = Iteratee . loop where
	loop (Continue k) = withBytes $ \bytes -> if B.null bytes
		then return $ Continue k
		else runIteratee (k (Chunks [bytes])) >>= loop
	
	loop step = return step
	
	intSize = fromInteger bufferSize
	withBytes = tryStep $ do
		hasInput <- E.catch
			(IO.hWaitForInput h (-1))
			(\(E.SomeException _) -> return False) 
		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 E.SomeException B.ByteString IO b
enumFile path s = Iteratee io where
	withHandle = tryStep (IO.openBinaryFile path IO.ReadMode)
	io = withHandle $ \h -> E.finally
		(runIteratee (enumHandle 4096 h s))
		(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 :: IO.Handle -> Iteratee E.SomeException B.ByteString IO ()
iterHandle h = continue step where
	step EOF = yield () EOF
	step (Chunks []) = continue step
	step (Chunks bytes) = Iteratee io where
		put = mapM_ (B.hPut h) bytes
		io = tryStep put (\_ -> return $ Continue step)