1--
2-- Licensed to the Apache Software Foundation (ASF) under one
3-- or more contributor license agreements. See the NOTICE file
4-- distributed with this work for additional information
5-- regarding copyright ownership. The ASF licenses this file
6-- to you under the Apache License, Version 2.0 (the
7-- "License"); you may not use this file except in compliance
8-- with the License. You may obtain a copy of the License at
9--
10--   http://www.apache.org/licenses/LICENSE-2.0
11--
12-- Unless required by applicable law or agreed to in writing,
13-- software distributed under the License is distributed on an
14-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15-- KIND, either express or implied. See the License for the
16-- specific language governing permissions and limitations
17-- under the License.
18--
19
20module Thrift.Transport.IOBuffer
21       ( WriteBuffer
22       , newWriteBuffer
23       , writeBuf
24       , flushBuf
25       , ReadBuffer
26       , newReadBuffer
27       , fillBuf
28       , readBuf
29       , peekBuf
30       ) where
31
32import Data.ByteString.Lazy.Builder
33import Data.Functor
34import Data.IORef
35import Data.Monoid
36import Data.Word
37
38import qualified Data.ByteString.Lazy as LBS
39
40type WriteBuffer = IORef Builder
41type ReadBuffer = IORef LBS.ByteString
42
43newWriteBuffer :: IO WriteBuffer
44newWriteBuffer = newIORef mempty
45
46writeBuf :: WriteBuffer -> LBS.ByteString -> IO ()
47writeBuf w s = modifyIORef w ( <> lazyByteString s)
48
49flushBuf :: WriteBuffer -> IO LBS.ByteString
50flushBuf w = do
51  buf <- readIORef w
52  writeIORef w mempty
53  return $ toLazyByteString buf
54
55newReadBuffer :: IO ReadBuffer
56newReadBuffer = newIORef mempty
57
58fillBuf :: ReadBuffer -> LBS.ByteString -> IO ()
59fillBuf = writeIORef
60
61readBuf :: ReadBuffer -> Int -> IO LBS.ByteString
62readBuf r n = do
63  bs <- readIORef r
64  let (hd, tl) = LBS.splitAt (fromIntegral n) bs
65  writeIORef r tl
66  return hd
67
68peekBuf :: ReadBuffer -> IO (Maybe Word8)
69peekBuf r = (fmap fst . LBS.uncons) <$> readIORef r
70