RingBuffer In A Byte String

ByteRing.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module ByteRing where
import Foreign.Ptr
import Foreign.C.Types
import Foreign.C.String
import Foreign.ForeignPtr
import System.Posix.Types
import System.Posix.Internals
import Data.ByteString qualified as BS
import Data.ByteString.Unsafe qualified as UnsafeBS
import Data.ByteString.Internal qualified as InternalBS
import Data.Word
import Control.Monad
import Data.Coerce
import Data.Char

foreign import ccall "memfd_create" mem_fd_create :: CString -> CUInt -> IO FD
foreign import ccall "ftruncate" ftruncate :: FD -> COff -> IO CInt
foreign import ccall "mmap" mmap :: Ptr Word8 -> CSize -> CInt -> CInt -> FD -> COff -> IO (Ptr Word8)
foreign import ccall "getpagesize" getpagesize :: IO CInt

-- Create a buffer capable of holding @size@ elements, where @size@
-- must be pagesize aligned. Returns a pointer to the start of an area
-- 2 * @size@ bytes in width, where the second half of the space is
-- mapped back to the same underlying memory region as the first.
ringMap :: CInt -> IO (Ptr Word8)
ringMap size = do
  let
    protReadWrite = 3
    mapPrivateAnonymous = 34
    mapSharedFixed = 17
  fd <- BS.useAsCString "ringMapFD" $ \name -> mem_fd_create name 0
  ftruncate fd (fromIntegral size)
  region <- mmap nullPtr (fromIntegral $ 2 * size) protReadWrite mapPrivateAnonymous (-1) 0
  let region' = plusPtr region (fromIntegral size)
  mmap region  (fromIntegral size) protReadWrite mapSharedFixed fd 0
  mmap region' (fromIntegral size) protReadWrite mapSharedFixed fd 0
  pure region

byteRing :: Int -> IO BS.ByteString
byteRing size = do
  ptr <- ringMap (fromIntegral size)
  UnsafeBS.unsafePackCStringLen (coerce ptr, size)

writeByteRing :: Int -> BS.ByteString -> BS.ByteString -> IO ()
writeByteRing offset inputData outputData = do
  let (inputForeignPtr, _, inputSize) = InternalBS.toForeignPtr inputData
      (outputForeignPtr, _, outputSize) = InternalBS.toForeignPtr outputData
  guard (outputSize >= inputSize)
  withForeignPtr inputForeignPtr $ \inputRawPtr -> do
    withForeignPtr outputForeignPtr $ \outputRawPtr -> do
      InternalBS.memcpy (plusPtr outputRawPtr $ fromIntegral offset) inputRawPtr inputSize


testData :: BS.ByteString
testData = BS.pack (replicate 4096 $ fromIntegral $ ord '1')

testData' :: BS.ByteString
testData' = BS.pack (replicate 100 $ fromIntegral $ ord '2')

testBuffer :: IO BS.ByteString
testBuffer = byteRing 4096

example :: IO ()
example = do
  b <- testBuffer
  writeByteRing 0 testData b
  print b
  writeByteRing 4090 testData' b
  print b