Map.hsc

module Database.QDBM.Cabin.Map
    ( Map
    , CBMAP

    , wrapMap
    , unsafePeekMap
    , withMapPtr

    , newMap
    , put
    , get

    , toList
    , fromList
    )
    where


import qualified Data.ByteString as Strict (ByteString)
import qualified Data.ByteString.Char8 as C8 hiding (ByteString)
import           Foreign.C.Types
import           Foreign.ForeignPtr
import           Foreign.Ptr
import           Foreign.Storable
import           Foreign.Marshal.Alloc


newtype Map = Map (ForeignPtr CBMAP)
data CBMAP


foreign import ccall unsafe "cabin.h cbmapopen"
        _open :: IO (Ptr CBMAP)

foreign import ccall unsafe "cabin.h &cbmapclose"
        _close :: FunPtr (Ptr CBMAP -> IO ())

foreign import ccall unsafe "cabin.h cbmapput"
        _put :: Ptr CBMAP -> Ptr CChar -> CInt -> Ptr CChar -> CInt -> CInt -> IO CInt

foreign import ccall unsafe "cabin.h cbmapget"
        _get :: Ptr CBMAP -> Ptr CChar -> CInt -> Ptr CInt -> IO (Ptr CChar)

foreign import ccall unsafe "cabin.h cbmapiterinit"
        _iterinit :: Ptr CBMAP -> IO ()

foreign import ccall unsafe "cabin.h cbmapiternext"
        _iternext :: Ptr CBMAP -> Ptr CInt -> IO (Ptr CChar)

foreign import ccall unsafe "cabin.h cbmapiterval"
        _iterval :: Ptr CChar -> Ptr CInt -> IO (Ptr CChar)


wrapMap :: Ptr CBMAP -> IO Map
wrapMap = fmap Map . newForeignPtr _close


unsafePeekMap :: Ptr CBMAP -> IO Map
unsafePeekMap = fmap Map . newForeignPtr_


withMapPtr :: Map -> (Ptr CBMAP -> IO a) -> IO a
withMapPtr (Map m) = withForeignPtr m


newMap :: IO Map
newMap = _open >>= wrapMap


put :: Map -> Strict.ByteString -> Strict.ByteString -> Bool -> IO Bool
put m key value overwrite
    = withMapPtr         m     $ \ mapPtr               ->
      C8.useAsCStringLen key   $ \ (keyPtr  , keyLen  ) ->
      C8.useAsCStringLen value $ \ (valuePtr, valueLen) ->
      fmap (/= 0) (_put mapPtr
                        keyPtr   (fromIntegral keyLen  )
                        valuePtr (fromIntegral valueLen)
                        (fromIntegral $ fromEnum overwrite))


get :: Map -> Strict.ByteString -> IO (Maybe Strict.ByteString)
get m key
    = withMapPtr         m   $ \ mapPtr           ->
      C8.useAsCStringLen key $ \ (keyPtr, keyLen) ->
      alloca                 $ \ valLenPtr        ->
      do valPtr <- _get mapPtr keyPtr (fromIntegral keyLen) valLenPtr
         if valPtr == nullPtr then
             return Nothing
           else
             do valLen <- peek valLenPtr
                value  <- C8.packCStringLen (valPtr, fromIntegral valLen)
                return $ Just value


initIterator :: Map -> IO ()
initIterator m
    = withMapPtr m _iterinit

iterateNext :: Map -> IO (Maybe (Strict.ByteString, Strict.ByteString))
iterateNext m
    = withMapPtr m $ \ mapPtr    ->
      alloca       $ \ keyLenPtr ->
      alloca       $ \ valLenPtr ->
      do keyPtr <- _iternext mapPtr keyLenPtr
         if keyPtr == nullPtr then
             return Nothing
           else
             do keyLen <- peek keyLenPtr
                key    <- C8.packCStringLen (keyPtr, fromIntegral keyLen)
                valPtr <- _iterval keyPtr valLenPtr
                valLen <- peek valLenPtr
                value  <- C8.packCStringLen (valPtr, fromIntegral valLen)
                return $ Just (key, value)

-- Internal state of the iterator is stored in the Map itself. That's
-- not thread-safe. So we can't iterate it lazily.
toList :: Map -> IO [(Strict.ByteString, Strict.ByteString)]
toList m = initIterator m >> loop
    where
      loop :: IO [(Strict.ByteString, Strict.ByteString)]
      loop = do next <- iterateNext m
                case next of
                  Nothing   -> return []
                  Just pair -> do -- We want to do unsafeInterleaveIO
                                  -- here, but we can't.
                                  rest <- loop
                                  return $ pair : rest


fromList :: [(Strict.ByteString, Strict.ByteString)] -> IO Map
fromList pairs
    = do m <- newMap
         mapM_ (putPair m) pairs
         return m
    where
      putPair :: Map -> (Strict.ByteString, Strict.ByteString) -> IO ()
      putPair m (key, value) = put m key value True >> return ()