Memory.hsc

{-# LANGUAGE CApiFFI #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE NoFieldSelectors #-}

{-| Raw bindings to @roaring/memory.h@.
It sets the allocator that CRoaring uses for every bitmap, container, iterator and bitset.

The hook is process-wide.
Install it before CRoaring allocates anything.
Never change it while any CRoaring structure is alive.
Otherwise a different allocator frees memory that the old one allocated.

Every hook must be a C function, not a Haskell function wrapped with @foreign import ccall "wrapper"@.
CRoaring allocates inside @unsafe@ foreign calls, which must not re-enter Haskell.
-}
module Data.Bitmap.Roaring.FFI.Memory
  ( Hooks (..)
  , Malloc
  , Realloc
  , Calloc
  , Free
  , AlignedMalloc
  , AlignedFree
  , initMemoryHook
  ) where

#include <roaring/memory.h>

import Foreign.C.Types (CSize)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (FunPtr, Ptr)
import Foreign.Storable (Storable (..))

-- | @roaring_malloc_p@
type Malloc = CSize -> IO (Ptr ())

-- | @roaring_realloc_p@
type Realloc = Ptr () -> CSize -> IO (Ptr ())

-- | @roaring_calloc_p@
type Calloc = CSize -> CSize -> IO (Ptr ())

-- | @roaring_free_p@: must accept and ignore NULL.
type Free = Ptr () -> IO ()

-- | @roaring_aligned_malloc_p@: alignment first, size second.
type AlignedMalloc = CSize -> CSize -> IO (Ptr ())

-- | @roaring_aligned_free_p@: frees blocks from 'AlignedMalloc'. It must accept and ignore NULL.
type AlignedFree = Ptr () -> IO ()

-- | @roaring_memory_t@
data Hooks = Hooks
  { malloc :: FunPtr Malloc
  , realloc :: FunPtr Realloc
  , calloc :: FunPtr Calloc
  , free :: FunPtr Free
  , alignedMalloc :: FunPtr AlignedMalloc
  , alignedFree :: FunPtr AlignedFree
  }
  deriving (Eq, Show)

instance Storable Hooks where
  sizeOf _ = #{size roaring_memory_t}
  {-# INLINE sizeOf #-}
  alignment _ = #{alignment roaring_memory_t}
  {-# INLINE alignment #-}
  peek ptr =
    Hooks
      <$> #{peek roaring_memory_t, malloc} ptr
      <*> #{peek roaring_memory_t, realloc} ptr
      <*> #{peek roaring_memory_t, calloc} ptr
      <*> #{peek roaring_memory_t, free} ptr
      <*> #{peek roaring_memory_t, aligned_malloc} ptr
      <*> #{peek roaring_memory_t, aligned_free} ptr
  {-# INLINE peek #-}
  poke ptr hooks = do
    #{poke roaring_memory_t, malloc} ptr hooks.malloc
    #{poke roaring_memory_t, realloc} ptr hooks.realloc
    #{poke roaring_memory_t, calloc} ptr hooks.calloc
    #{poke roaring_memory_t, free} ptr hooks.free
    #{poke roaring_memory_t, aligned_malloc} ptr hooks.alignedMalloc
    #{poke roaring_memory_t, aligned_free} ptr hooks.alignedFree
  {-# INLINE poke #-}

-- | @roaring_init_memory_hook@: see the module header for when it is safe to call.
initMemoryHook :: Hooks -> IO ()
initMemoryHook hooks =
  alloca $ \ptr -> do
    poke ptr hooks
    initMemoryHookFrom ptr

foreign import capi unsafe "croaring-ffi.h croaring_ffi_init_memory_hook"
  initMemoryHookFrom :: Ptr Hooks -> IO ()