ceph/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/IOBuffer.hs

   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
  20 module Thrift.Transport.IOBuffer
  21        ( WriteBuffer
  22        , newWriteBuffer
  23        , writeBuf
  24        , flushBuf
  25        , ReadBuffer
  26        , newReadBuffer
  27        , fillBuf
  28        , readBuf
  29        , peekBuf
  30        ) where
  31
  32 import Data.ByteString.Lazy.Builder
  33 import Data.Functor
  34 import Data.IORef
  35 import Data.Monoid
  36 import Data.Word
  37
  38 import qualified Data.ByteString.Lazy as LBS
  39
  40 type WriteBuffer = IORef Builder
  41 type ReadBuffer = IORef LBS.ByteString
  42
  43 newWriteBuffer :: IO WriteBuffer
  44 newWriteBuffer = newIORef mempty
  45
  46 writeBuf :: WriteBuffer -> LBS.ByteString -> IO ()
  47 writeBuf w s = modifyIORef w ( <> lazyByteString s)
  48
  49 flushBuf :: WriteBuffer -> IO LBS.ByteString
  50 flushBuf w = do
  51   buf <- readIORef w
  52   writeIORef w mempty
  53   return $ toLazyByteString buf
  54
  55 newReadBuffer :: IO ReadBuffer
  56 newReadBuffer = newIORef mempty
  57
  58 fillBuf :: ReadBuffer -> LBS.ByteString -> IO ()
  59 fillBuf = writeIORef
  60
  61 readBuf :: ReadBuffer -> Int -> IO LBS.ByteString
  62 readBuf r n = do
  63   bs <- readIORef r
  64   let (hd, tl) = LBS.splitAt (fromIntegral n) bs
  65   writeIORef r tl
  66   return hd
  67
  68 peekBuf :: ReadBuffer -> IO (Maybe Word8)
  69 peekBuf r = (fmap fst . LBS.uncons) <$> readIORef r