Skip to content

Instantly share code, notes, and snippets.

@TerrorJack

TerrorJack/Pool.hs

Created Aug 2, 2019
Embed
What would you like to do?
Memory pool backed by ghc storage manager
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
module Pool
( Pool
, newPool
, freePool
, pooledMallocBytes
, pooledReallocBytes
, pooledFree
) where
import Data.IORef
import qualified Data.IntMap.Strict as IM
import Foreign.Ptr
import GHC.Exts
import GHC.ForeignPtr
import GHC.Types
newtype Pool =
Pool (IORef (IM.IntMap (ForeignPtr ())))
newPool :: IO Pool
newPool = Pool <$> newIORef IM.empty
freePool :: Pool -> IO ()
freePool (Pool ref) = atomicWriteIORef ref IM.empty
pooledMallocBytes :: Pool -> Int -> IO (Ptr a)
pooledMallocBytes (Pool ref) size = do
fptr <- mallocPlainForeignPtrBytes size
atomicModifyIORef' ref $ \im ->
case unsafeForeignPtrToPtr fptr of
ptr -> (IM.insert (unsafePtrToInt ptr) fptr im, castPtr ptr)
pooledReallocBytes :: Pool -> Ptr a -> Int -> IO (Ptr a)
pooledReallocBytes (Pool ref) old_ptr new_size = do
new_fptr@(ForeignPtr _ (PlainPtr new_mba)) <-
mallocPlainForeignPtrBytes new_size
PlainPtr old_mba <-
case unsafePtrToInt old_ptr of
old_key ->
atomicModifyIORef' ref $ \im ->
case im IM.! old_key of
(ForeignPtr _ old_c) ->
( IM.insert
(unsafePtrToInt $ unsafeForeignPtrToPtr new_fptr)
new_fptr $
IM.delete old_key im
, old_c)
IO $ \s0 ->
case getSizeofMutableByteArray# old_mba s0 of
(# s1, old_size_raw #) ->
case copyMutableByteArray#
old_mba
0#
new_mba
0#
(unI# $ min new_size $ I# old_size_raw)
s1 of
s2 -> (# s2, castPtr $ unsafeForeignPtrToPtr new_fptr #)
pooledFree :: Pool -> Ptr a -> IO ()
pooledFree (Pool ref) ptr =
atomicModifyIORef' ref $ \im -> (IM.delete (unsafePtrToInt ptr) im, ())
unsafePtrToInt :: Ptr a -> Int
unsafePtrToInt (Ptr addr) = I# (addr2Int# addr)
unI# :: Int -> Int#
unI# (I# x) = x
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.