Memory pool backed by ghc storage manager

Pool.hs

{-# 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