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

{-# LANGUAGE OverloadedStrings #-}

-- | UniformIO on memory
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 Control.Concurrent.MVar

--import Data.Default.Class

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

{- |
Wrapper that does UniformIO that reads and writes on the memory.

Input and output may be queried and replaced during the execution of
the target, with obviously undefined results in case of concurrent
execution.
-}
data ByteStringIO = ByteStringIO {bsioinput :: MVar ByteString, bsiooutput :: MVar BSBuild.Builder}
instance UniformIO ByteStringIO where
  uRead s n = do
    i <- takeMVar . bsioinput $ s
    let (r, i') = BS.splitAt n i
    putMVar (bsioinput s) i'
    return r
  uPut s t = do
    o <- takeMVar . bsiooutput $ s
    let o' = mappend o $ BSBuild.byteString t
    putMVar (bsiooutput s) o'
  uClose _ = return ()
  startTls _ = return
  isSecure _ = True
  isEOF t = withMVar (bsioinput t) $ return . BS.null

-- | 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
  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)