Pool.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Pool
  ( Pool,
    newPool,
    freePool,
    pooledMalloc,
    pooledMallocBytes,
    pooledMallocArray,
    pooledNew,
    pooledNewArray,
    pooledNewByteString0,
  )
where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BS
import Data.Word
import Foreign.C
import Foreign.Marshal.Array
import Foreign.Marshal.Utils
import Foreign.Ptr
import Foreign.Storable

newtype Pool = Pool (Ptr Pool)

pooledMalloc :: forall a. Storable a => Pool -> IO (Ptr a)
pooledMalloc p = pooledMallocBytes p (sizeOf (undefined :: a))

pooledMallocBytes :: Pool -> Int -> IO (Ptr a)
pooledMallocBytes p l
  | l > 0 = c_pooledMallocBytes p (fromIntegral l)
  | otherwise = fail ("pooledMallocBytes: invalid size " <> show l)

pooledMallocArray :: forall a. Storable a => Pool -> Int -> IO (Ptr a)
pooledMallocArray p l = pooledMallocBytes p (sizeOf (undefined :: a) * l)

pooledNew :: Storable a => Pool -> a -> IO (Ptr a)
pooledNew p x = do
  buf <- pooledMalloc p
  poke buf x
  pure buf

pooledNewArray :: Storable a => Pool -> [a] -> IO (Ptr a)
pooledNewArray p xs
  | null xs = pure nullPtr
  | otherwise = do
    buf <- pooledMallocArray p (length xs)
    pokeArray buf xs
    pure buf

pooledNewByteString0 :: Pool -> BS.ByteString -> IO CString
pooledNewByteString0 p bs
  | BS.null bs = pure nullPtr
  | otherwise =
    BS.unsafeUseAsCStringLen
      bs
      ( \(src, len) -> do
          dst <- pooledMallocBytes p (succ len)
          copyBytes dst src len
          pokeByteOff dst len (0 :: Word8)
          pure dst
      )

foreign import ccall unsafe "newArena" newPool :: IO Pool

foreign import ccall unsafe "arenaAlloc" c_pooledMallocBytes :: Pool -> CSize -> IO (Ptr a)

foreign import ccall unsafe "arenaFree" freePool :: Pool -> IO ()