Index.hs

{-| An off-heap set of entity ids in a CRoaring 64-bit bitmap, freed by the garbage collector.

Stores that keep one expose it through 'HasIndex', so queries can combine them without touching the Haskell heap.
-}
module Apecs.Stores.Index
  ( Index
  , HasIndex (..)
  , new
  , insert
  , delete
  , member
  , size
  , clear
  , toVector
  , with
  , bitmapToVector
  , checked
  ) where

import Control.Exception (throwIO)
import Data.Bitmap.Roaring.FFI.Bitmap64 (Bitmap64 (..))
import Data.Bitmap.Roaring.FFI.Bitmap64 qualified as Bitmap64
import Data.Bitmap.Roaring.FFI.NullHandle qualified as NullHandle
import Data.Primitive.ByteArray (newPinnedByteArray, unsafeFreezeByteArray, withMutableByteArrayContents)
import Data.Vector.Primitive qualified as Primitive
import Data.Vector.Unboxed qualified as Unboxed
import Data.Vector.Unboxed.Base qualified as Unboxed.Base
import Foreign.ForeignPtr (ForeignPtr, newForeignPtr, withForeignPtr)
import Foreign.Ptr (castPtr, nullPtr)
import Foreign.Storable (sizeOf)
import GHC.ForeignPtr (unsafeWithForeignPtr)

newtype Index = Index (ForeignPtr Bitmap64)

class HasIndex s where
  getIndex :: s -> Index

new :: IO Index
new = do
  Bitmap64 handle <- Bitmap64.create >>= checked "Apecs.Stores.Index.new"
  Index <$> newForeignPtr Bitmap64.finalizerFree handle

checked :: String -> Bitmap64 -> IO Bitmap64
checked function bitmap@(Bitmap64 handle)
  | handle == nullPtr = throwIO $ NullHandle.NullHandle function
  | otherwise = pure bitmap

with :: Index -> (Bitmap64 -> IO a) -> IO a
with (Index foreignPtr) action = withForeignPtr foreignPtr $ action . Bitmap64
{-# INLINE with #-}

withCall :: Index -> (Bitmap64 -> IO a) -> IO a
withCall (Index foreignPtr) action = unsafeWithForeignPtr foreignPtr $ action . Bitmap64
{-# INLINE withCall #-}

insert :: Index -> Int -> IO ()
insert index ety = withCall index $ \bitmap -> Bitmap64.add bitmap $ fromIntegral ety
{-# INLINE insert #-}

delete :: Index -> Int -> IO ()
delete index ety = withCall index $ \bitmap -> Bitmap64.remove bitmap $ fromIntegral ety
{-# INLINE delete #-}

member :: Index -> Int -> IO Bool
member index ety = withCall index $ \bitmap -> Bitmap64.contains bitmap $ fromIntegral ety
{-# INLINE member #-}

size :: Index -> IO Int
size index = withCall index $ fmap fromIntegral . Bitmap64.getCardinality
{-# INLINE size #-}

clear :: Index -> IO ()
clear index = withCall index Bitmap64.clear
{-# INLINE clear #-}

toVector :: Index -> IO (Unboxed.Vector Int)
toVector index = with index bitmapToVector
{-# INLINE toVector #-}

bitmapToVector :: Bitmap64 -> IO (Unboxed.Vector Int)
bitmapToVector bitmap = do
  count <- fromIntegral <$> Bitmap64.getCardinality bitmap
  buffer <- newPinnedByteArray $ count * sizeOf (0 :: Int)
  withMutableByteArrayContents buffer $ Bitmap64.toUint64Array bitmap . castPtr
  Unboxed.Base.V_Int . Primitive.Vector 0 count <$> unsafeFreezeByteArray buffer