{-# LANGUAGE ImportQualifiedPost #-}

-- |
-- Module      : Database.RocksDB.Iterator
-- Copyright   : (c) 2012-2013 The leveldb-haskell Authors
--               (c) 2014-2020 The rocksdb-haskell Authors
-- License     : BSD3
-- Maintainer  : jprupp@protonmail.ch
-- Stability   : experimental
-- Portability : non-portable
--
-- Iterating over key ranges.
module Database.RocksDB.Iterator
  ( Iterator,
    withIter,
    withIterCF,
    iter,
    iterCF,
    iterator,
    createIterator,
    destroyIterator,
    iterEntry,
    iterFirst,
    iterGetError,
    iterKey,
    iterLast,
    iterNext,
    iterPrev,
    iterSeek,
    iterValid,
    iterValue,
  )
where

import Data.ByteString (ByteString)
import Data.ByteString qualified as B
import Data.ByteString.Unsafe qualified as BU
import Database.RocksDB.C
import Database.RocksDB.Internal
import Foreign
import Foreign.C.Error (throwErrnoIfNull)
import Foreign.C.String (CString)
import Foreign.C.Types (CSize)
import UnliftIO
import UnliftIO.Resource

-- | Create 'Iterator' and use it.
--
-- Note that an 'Iterator' creates a snapshot of the database implicitly, so
-- updates written after the iterator was created are not visible. You may,
-- however, specify an older 'Snapshot' in the 'ReadOptions'.
--
-- Iterator should not be used after computation ends.
withIter :: (MonadUnliftIO m) => DB -> (Iterator -> m a) -> m a
withIter :: forall (m :: * -> *) a.
MonadUnliftIO m =>
DB -> (Iterator -> m a) -> m a
withIter DB
db = DB -> Maybe ColumnFamily -> (Iterator -> m a) -> m a
forall (m :: * -> *) a.
MonadUnliftIO m =>
DB -> Maybe ColumnFamily -> (Iterator -> m a) -> m a
withIterCommon DB
db Maybe ColumnFamily
forall a. Maybe a
Nothing

withIterCF :: (MonadUnliftIO m) => DB -> ColumnFamily -> (Iterator -> m a) -> m a
withIterCF :: forall (m :: * -> *) a.
MonadUnliftIO m =>
DB -> ColumnFamily -> (Iterator -> m a) -> m a
withIterCF DB
db ColumnFamily
cf = DB -> Maybe ColumnFamily -> (Iterator -> m a) -> m a
forall (m :: * -> *) a.
MonadUnliftIO m =>
DB -> Maybe ColumnFamily -> (Iterator -> m a) -> m a
withIterCommon DB
db (ColumnFamily -> Maybe ColumnFamily
forall a. a -> Maybe a
Just ColumnFamily
cf)

-- | Variation on 'iterator' below.
iter :: (MonadIO m, MonadResource m) => DB -> m Iterator
iter :: forall (m :: * -> *).
(MonadIO m, MonadResource m) =>
DB -> m Iterator
iter DB
db = DB -> Maybe ColumnFamily -> m Iterator
forall (m :: * -> *).
(MonadIO m, MonadResource m) =>
DB -> Maybe ColumnFamily -> m Iterator
iterator DB
db Maybe ColumnFamily
forall a. Maybe a
Nothing

iterCF :: (MonadIO m, MonadResource m) => DB -> ColumnFamily -> m Iterator
iterCF :: forall (m :: * -> *).
(MonadIO m, MonadResource m) =>
DB -> ColumnFamily -> m Iterator
iterCF DB
db ColumnFamily
cf = DB -> Maybe ColumnFamily -> m Iterator
forall (m :: * -> *).
(MonadIO m, MonadResource m) =>
DB -> Maybe ColumnFamily -> m Iterator
iterator DB
db (ColumnFamily -> Maybe ColumnFamily
forall a. a -> Maybe a
Just ColumnFamily
cf)

withIterCommon ::
  (MonadUnliftIO m) =>
  DB ->
  Maybe ColumnFamily ->
  (Iterator -> m a) ->
  m a
withIterCommon :: forall (m :: * -> *) a.
MonadUnliftIO m =>
DB -> Maybe ColumnFamily -> (Iterator -> m a) -> m a
withIterCommon DB {rocksDB :: DB -> RocksDB
rocksDB = RocksDB
rocks_db, readOpts :: DB -> ReadOpts
readOpts = ReadOpts
read_opts} Maybe ColumnFamily
mcf =
  m Iterator -> (Iterator -> m ()) -> (Iterator -> m a) -> m a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket m Iterator
create_iterator Iterator -> m ()
destroy_iterator
  where
    destroy_iterator :: Iterator -> m ()
destroy_iterator = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (Iterator -> IO ()) -> Iterator -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Iterator -> IO ()
c_rocksdb_iter_destroy
    create_iterator :: m Iterator
create_iterator = IO Iterator -> m Iterator
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Iterator -> m Iterator) -> IO Iterator -> m Iterator
forall a b. (a -> b) -> a -> b
$
      String -> IO Iterator -> IO Iterator
forall a. String -> IO (Ptr a) -> IO (Ptr a)
throwErrnoIfNull String
"create_iterator" (IO Iterator -> IO Iterator) -> IO Iterator -> IO Iterator
forall a b. (a -> b) -> a -> b
$ case Maybe ColumnFamily
mcf of
        Just ColumnFamily
cf -> RocksDB -> ReadOpts -> ColumnFamily -> IO Iterator
c_rocksdb_create_iterator_cf RocksDB
rocks_db ReadOpts
read_opts ColumnFamily
cf
        Maybe ColumnFamily
Nothing -> RocksDB -> ReadOpts -> IO Iterator
c_rocksdb_create_iterator RocksDB
rocks_db ReadOpts
read_opts

-- | Iterator is not valid outside of 'ResourceT' context.
iterator ::
  (MonadIO m, MonadResource m) =>
  DB ->
  Maybe ColumnFamily ->
  m Iterator
iterator :: forall (m :: * -> *).
(MonadIO m, MonadResource m) =>
DB -> Maybe ColumnFamily -> m Iterator
iterator DB
db Maybe ColumnFamily
mcf =
  (ReleaseKey, Iterator) -> Iterator
forall a b. (a, b) -> b
snd ((ReleaseKey, Iterator) -> Iterator)
-> m (ReleaseKey, Iterator) -> m Iterator
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Iterator -> (Iterator -> IO ()) -> m (ReleaseKey, Iterator)
forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
allocate (DB -> Maybe ColumnFamily -> IO Iterator
forall (m :: * -> *).
MonadIO m =>
DB -> Maybe ColumnFamily -> m Iterator
createIterator DB
db Maybe ColumnFamily
mcf) Iterator -> IO ()
forall (m :: * -> *). MonadIO m => Iterator -> m ()
destroyIterator

-- | Manually create unmanaged iterator.
createIterator :: (MonadIO m) => DB -> Maybe ColumnFamily -> m Iterator
createIterator :: forall (m :: * -> *).
MonadIO m =>
DB -> Maybe ColumnFamily -> m Iterator
createIterator DB {rocksDB :: DB -> RocksDB
rocksDB = RocksDB
rocks_db, readOpts :: DB -> ReadOpts
readOpts = ReadOpts
read_opts} Maybe ColumnFamily
mcf = IO Iterator -> m Iterator
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Iterator -> m Iterator) -> IO Iterator -> m Iterator
forall a b. (a -> b) -> a -> b
$
  String -> IO Iterator -> IO Iterator
forall a. String -> IO (Ptr a) -> IO (Ptr a)
throwErrnoIfNull String
"create_iterator" (IO Iterator -> IO Iterator) -> IO Iterator -> IO Iterator
forall a b. (a -> b) -> a -> b
$ case Maybe ColumnFamily
mcf of
    Just ColumnFamily
cf -> RocksDB -> ReadOpts -> ColumnFamily -> IO Iterator
c_rocksdb_create_iterator_cf RocksDB
rocks_db ReadOpts
read_opts ColumnFamily
cf
    Maybe ColumnFamily
Nothing -> RocksDB -> ReadOpts -> IO Iterator
c_rocksdb_create_iterator RocksDB
rocks_db ReadOpts
read_opts

-- | Destroy unmanaged iterator.
destroyIterator :: (MonadIO m) => Iterator -> m ()
destroyIterator :: forall (m :: * -> *). MonadIO m => Iterator -> m ()
destroyIterator = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (Iterator -> IO ()) -> Iterator -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Iterator -> IO ()
c_rocksdb_iter_destroy

-- | An iterator is either positioned at a key/value pair, or not valid. This
-- function returns /true/ iff the iterator is valid.
iterValid :: (MonadIO m) => Iterator -> m Bool
iterValid :: forall (m :: * -> *). MonadIO m => Iterator -> m Bool
iterValid Iterator
iter_ptr = IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> m Bool) -> IO Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ do
  CUChar
x <- Iterator -> IO CUChar
c_rocksdb_iter_valid Iterator
iter_ptr
  Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CUChar
x CUChar -> CUChar -> Bool
forall a. Eq a => a -> a -> Bool
/= CUChar
0)

-- | Position at the first key in the source that is at or past target. The
-- iterator is /valid/ after this call iff the source contains an entry that
-- comes at or past target.
iterSeek :: (MonadIO m) => Iterator -> ByteString -> m ()
iterSeek :: forall (m :: * -> *). MonadIO m => Iterator -> ByteString -> m ()
iterSeek Iterator
iter_ptr ByteString
key = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
  ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BU.unsafeUseAsCStringLen ByteString
key ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(CString
key_ptr, Int
klen) ->
    Iterator -> CString -> CSize -> IO ()
c_rocksdb_iter_seek Iterator
iter_ptr CString
key_ptr (Int -> CSize
intToCSize Int
klen)

-- | Position at the first key in the source. The iterator is /valid/ after this
-- call iff the source is not empty.
iterFirst :: (MonadIO m) => Iterator -> m ()
iterFirst :: forall (m :: * -> *). MonadIO m => Iterator -> m ()
iterFirst = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (Iterator -> IO ()) -> Iterator -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Iterator -> IO ()
c_rocksdb_iter_seek_to_first

-- | Position at the last key in the source. The iterator is /valid/ after this
-- call iff the source is not empty.
iterLast :: (MonadIO m) => Iterator -> m ()
iterLast :: forall (m :: * -> *). MonadIO m => Iterator -> m ()
iterLast = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (Iterator -> IO ()) -> Iterator -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Iterator -> IO ()
c_rocksdb_iter_seek_to_last

-- | Moves to the next entry in the source. After this call, 'iterValid' is
-- /true/ iff the iterator was not positioned at the last entry in the source.
--
-- If the iterator is not valid, this function does nothing. Note that this is a
-- shortcoming of the C API: an 'iterPrev' might still be possible, but we can't
-- determine if we're at the last or first entry.
iterNext :: (MonadIO m) => Iterator -> m ()
iterNext :: forall (m :: * -> *). MonadIO m => Iterator -> m ()
iterNext Iterator
iter_ptr = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Iterator -> IO ()
c_rocksdb_iter_next Iterator
iter_ptr

-- | Moves to the previous entry in the source. After this call, 'iterValid' is
-- /true/ iff the iterator was not positioned at the first entry in the source.
--
-- If the iterator is not valid, this function does nothing. Note that this is a
-- shortcoming of the C API: an 'iterNext' might still be possible, but we can't
-- determine if we're at the last or first entry.
iterPrev :: (MonadIO m) => Iterator -> m ()
iterPrev :: forall (m :: * -> *). MonadIO m => Iterator -> m ()
iterPrev Iterator
iter_ptr = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Iterator -> IO ()
c_rocksdb_iter_prev Iterator
iter_ptr

-- | Return the key for the current entry if the iterator is currently
-- positioned at an entry, ie. 'iterValid'.
iterKey :: (MonadIO m) => Iterator -> m (Maybe ByteString)
iterKey :: forall (m :: * -> *). MonadIO m => Iterator -> m (Maybe ByteString)
iterKey Iterator
it = IO (Maybe ByteString) -> m (Maybe ByteString)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ByteString) -> m (Maybe ByteString))
-> IO (Maybe ByteString) -> m (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ Iterator
-> (Iterator -> Ptr CSize -> IO CString) -> IO (Maybe ByteString)
iterString Iterator
it Iterator -> Ptr CSize -> IO CString
c_rocksdb_iter_key

-- | Return the value for the current entry if the iterator is currently
-- positioned at an entry, ie. 'iterValid'.
iterValue :: (MonadIO m) => Iterator -> m (Maybe ByteString)
iterValue :: forall (m :: * -> *). MonadIO m => Iterator -> m (Maybe ByteString)
iterValue Iterator
it = IO (Maybe ByteString) -> m (Maybe ByteString)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ByteString) -> m (Maybe ByteString))
-> IO (Maybe ByteString) -> m (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ Iterator
-> (Iterator -> Ptr CSize -> IO CString) -> IO (Maybe ByteString)
iterString Iterator
it Iterator -> Ptr CSize -> IO CString
c_rocksdb_iter_value

-- | Return the current entry as a pair, if the iterator is currently positioned
-- at an entry, ie. 'iterValid'.
iterEntry :: (MonadIO m) => Iterator -> m (Maybe (ByteString, ByteString))
iterEntry :: forall (m :: * -> *).
MonadIO m =>
Iterator -> m (Maybe (ByteString, ByteString))
iterEntry Iterator
it = IO (Maybe (ByteString, ByteString))
-> m (Maybe (ByteString, ByteString))
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (ByteString, ByteString))
 -> m (Maybe (ByteString, ByteString)))
-> IO (Maybe (ByteString, ByteString))
-> m (Maybe (ByteString, ByteString))
forall a b. (a -> b) -> a -> b
$ do
  Maybe ByteString
mkey <- Iterator -> IO (Maybe ByteString)
forall (m :: * -> *). MonadIO m => Iterator -> m (Maybe ByteString)
iterKey Iterator
it
  Maybe ByteString
mval <- Iterator -> IO (Maybe ByteString)
forall (m :: * -> *). MonadIO m => Iterator -> m (Maybe ByteString)
iterValue Iterator
it
  Maybe (ByteString, ByteString)
-> IO (Maybe (ByteString, ByteString))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (ByteString, ByteString)
 -> IO (Maybe (ByteString, ByteString)))
-> Maybe (ByteString, ByteString)
-> IO (Maybe (ByteString, ByteString))
forall a b. (a -> b) -> a -> b
$ (,) (ByteString -> ByteString -> (ByteString, ByteString))
-> Maybe ByteString
-> Maybe (ByteString -> (ByteString, ByteString))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe ByteString
mkey Maybe (ByteString -> (ByteString, ByteString))
-> Maybe ByteString -> Maybe (ByteString, ByteString)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe ByteString
mval

-- | Check for errors
--
-- Note that this captures somewhat severe errors such as a corrupted database.
iterGetError :: (MonadIO m) => Iterator -> m (Maybe ByteString)
iterGetError :: forall (m :: * -> *). MonadIO m => Iterator -> m (Maybe ByteString)
iterGetError Iterator
iter_ptr = IO (Maybe ByteString) -> m (Maybe ByteString)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ByteString) -> m (Maybe ByteString))
-> IO (Maybe ByteString) -> m (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ (Ptr CString -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr CString -> IO (Maybe ByteString)) -> IO (Maybe ByteString))
-> (Ptr CString -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ \Ptr CString
err_ptr -> do
  Ptr CString -> CString -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr CString
err_ptr CString
forall a. Ptr a
nullPtr
  Iterator -> Ptr CString -> IO ()
c_rocksdb_iter_get_error Iterator
iter_ptr Ptr CString
err_ptr
  CString
err_str <- Ptr CString -> IO CString
forall a. Storable a => Ptr a -> IO a
peek Ptr CString
err_ptr
  if CString
err_str CString -> CString -> Bool
forall a. Eq a => a -> a -> Bool
== CString
forall a. Ptr a
nullPtr
    then Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
    else ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString)
-> IO ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CString -> IO ByteString
BU.unsafePackMallocCString CString
err_str

--
-- Internal
--

iterString ::
  Iterator ->
  (Iterator -> Ptr CSize -> IO CString) ->
  IO (Maybe ByteString)
iterString :: Iterator
-> (Iterator -> Ptr CSize -> IO CString) -> IO (Maybe ByteString)
iterString Iterator
it Iterator -> Ptr CSize -> IO CString
f = do
  Bool
valid <- Iterator -> IO Bool
forall (m :: * -> *). MonadIO m => Iterator -> m Bool
iterValid Iterator
it
  if Bool
valid
    then (Ptr CSize -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr CSize -> IO (Maybe ByteString)) -> IO (Maybe ByteString))
-> (Ptr CSize -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ \Ptr CSize
len_ptr -> do
      CString
str_ptr <- Iterator -> Ptr CSize -> IO CString
f Iterator
it Ptr CSize
len_ptr
      if CString
str_ptr CString -> CString -> Bool
forall a. Eq a => a -> a -> Bool
== CString
forall a. Ptr a
nullPtr
        then Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
        else do
          CSize
len <- Ptr CSize -> IO CSize
forall a. Storable a => Ptr a -> IO a
peek Ptr CSize
len_ptr
          ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString)
-> IO ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CStringLen -> IO ByteString
B.packCStringLen (CString
str_ptr, CSize -> Int
cSizeToInt CSize
len)
    else Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing