uniform-io/src/System/IO/Uniform/ByteString.hs

{-# LANGUAGE OverloadedStrings #-}

module System.IO.Uniform.ByteString (
  ByteStringIO,
  withByteStringIO, withByteStringIO'
  ) where

import System.IO.Uniform

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString.Builder as BSBuild
import System.IO.Error
import Control.Concurrent.MVar

--import Data.Default.Class

--import System.Posix.Types (Fd(..))

-- | Wrapper that does UniformIO that reads and writes on the memory.
data ByteStringIO = ByteStringIO {bsioinput :: MVar (ByteString, Bool), bsiooutput :: MVar BSBuild.Builder}
instance UniformIO ByteStringIO where
  uRead s n = do
    (i, eof) <- takeMVar . bsioinput $ s
    if eof
    then do
      putMVar (bsioinput s) (i, eof)
      ioError $ mkIOError eofErrorType "read past end of input" Nothing Nothing
    else do
      let (r, i') = BS.splitAt n i
      let eof' = (BS.null r && n > 0)
      putMVar (bsioinput s) (i', eof')
      return r
  uPut s t = do
    o <- takeMVar . bsiooutput $ s
    let o' = mappend o $ BSBuild.byteString t
    putMVar (bsiooutput s) o'
  uClose _ = return ()
  startTls _ a = return a
  isSecure _ = True

-- | withByteStringIO' input f
--   Runs f with a ByteStringIO that has the given input, returns f's output and
--   the ByteStringIO output.
withByteStringIO' :: ByteString -> (ByteStringIO -> IO a) -> IO (a, LBS.ByteString)
withByteStringIO' input f = do
  ivar <- newMVar (input, False)
  ovar <- newMVar . BSBuild.byteString $ BS.empty
  let bsio = ByteStringIO ivar ovar
  a <- f bsio
  out <- takeMVar . bsiooutput $ bsio
  return (a, BSBuild.toLazyByteString out)

-- | The same as withByteStringIO', but returns an strict ByteString
withByteStringIO :: ByteString -> (ByteStringIO -> IO a) -> IO (a, ByteString)
withByteStringIO input f = do
  (a, t) <- withByteStringIO' input f
  return (a, LBS.toStrict t)