{-# LANGUAGE LambdaCase #-}
-- |
-- 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           Control.Monad             (when)
import           Data.ByteString           (ByteString)
import qualified Data.ByteString           as BS
import qualified Data.ByteString.Char8     as BC
import qualified Data.ByteString.Unsafe    as BU
import           Database.RocksDB.C
import           Database.RocksDB.Internal
import           Foreign
import           Foreign.C.Error           (throwErrnoIfNull)
import           Foreign.C.String          (CString, peekCString)
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 :: 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 :: 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 :: 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 :: 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 :: 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 (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 (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 :: 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 :: 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 (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 :: Iterator -> m ()
destroyIterator = IO () -> m ()
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 :: Iterator -> m Bool
iterValid Iterator
iter_ptr = IO Bool -> m Bool
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 (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 :: Iterator -> ByteString -> m ()
iterSeek Iterator
iter_ptr ByteString
key = IO () -> m ()
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
$ \(Ptr CChar
key_ptr, Int
klen) ->
    Iterator -> Ptr CChar -> CSize -> IO ()
c_rocksdb_iter_seek Iterator
iter_ptr Ptr CChar
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 :: Iterator -> m ()
iterFirst = IO () -> m ()
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 :: Iterator -> m ()
iterLast = IO () -> m ()
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 :: Iterator -> m ()
iterNext Iterator
iter_ptr = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    CUChar
valid <- Iterator -> IO CUChar
c_rocksdb_iter_valid Iterator
iter_ptr
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (CUChar
valid CUChar -> CUChar -> Bool
forall a. Eq a => a -> a -> Bool
/= CUChar
0) (IO () -> IO ()) -> IO () -> IO ()
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 :: Iterator -> m ()
iterPrev Iterator
iter_ptr = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    CUChar
valid <- Iterator -> IO CUChar
c_rocksdb_iter_valid Iterator
iter_ptr
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (CUChar
valid CUChar -> CUChar -> Bool
forall a. Eq a => a -> a -> Bool
/= CUChar
0) (IO () -> IO ()) -> IO () -> IO ()
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 :: Iterator -> m (Maybe ByteString)
iterKey = IO (Maybe ByteString) -> m (Maybe ByteString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ByteString) -> m (Maybe ByteString))
-> (Iterator -> IO (Maybe ByteString))
-> Iterator
-> m (Maybe ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Iterator
 -> (Iterator -> Ptr CSize -> IO (Ptr CChar))
 -> IO (Maybe ByteString))
-> (Iterator -> Ptr CSize -> IO (Ptr CChar))
-> Iterator
-> IO (Maybe ByteString)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Iterator
-> (Iterator -> Ptr CSize -> IO (Ptr CChar))
-> IO (Maybe ByteString)
iterString Iterator -> Ptr CSize -> IO (Ptr CChar)
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 :: Iterator -> m (Maybe ByteString)
iterValue = IO (Maybe ByteString) -> m (Maybe ByteString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ByteString) -> m (Maybe ByteString))
-> (Iterator -> IO (Maybe ByteString))
-> Iterator
-> m (Maybe ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Iterator
 -> (Iterator -> Ptr CSize -> IO (Ptr CChar))
 -> IO (Maybe ByteString))
-> (Iterator -> Ptr CSize -> IO (Ptr CChar))
-> Iterator
-> IO (Maybe ByteString)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Iterator
-> (Iterator -> Ptr CSize -> IO (Ptr CChar))
-> IO (Maybe ByteString)
iterString Iterator -> Ptr CSize -> IO (Ptr CChar)
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 :: Iterator -> m (Maybe (ByteString, ByteString))
iterEntry Iterator
it = IO (Maybe (ByteString, ByteString))
-> m (Maybe (ByteString, ByteString))
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 (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 (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 :: Iterator -> m (Maybe ByteString)
iterGetError Iterator
iter_ptr = IO (Maybe ByteString) -> m (Maybe ByteString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ByteString) -> m (Maybe ByteString))
-> ((Ptr (Ptr CChar) -> IO (Maybe ByteString))
    -> IO (Maybe ByteString))
-> (Ptr (Ptr CChar) -> IO (Maybe ByteString))
-> m (Maybe ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Ptr (Ptr CChar) -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr (Ptr CChar) -> IO (Maybe ByteString))
 -> m (Maybe ByteString))
-> (Ptr (Ptr CChar) -> IO (Maybe ByteString))
-> m (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr CChar)
err_ptr -> do
    Ptr (Ptr CChar) -> Ptr CChar -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr (Ptr CChar)
err_ptr Ptr CChar
forall a. Ptr a
nullPtr
    Iterator -> Ptr (Ptr CChar) -> IO ()
c_rocksdb_iter_get_error Iterator
iter_ptr Ptr (Ptr CChar)
err_ptr
    Ptr CChar
erra <- Ptr (Ptr CChar) -> IO (Ptr CChar)
forall a. Storable a => Ptr a -> IO a
peek Ptr (Ptr CChar)
err_ptr
    if Ptr CChar
erra Ptr CChar -> Ptr CChar -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr CChar
forall a. Ptr a
nullPtr
        then Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
        else do
            String
err <- Ptr CChar -> IO String
peekCString Ptr CChar
erra
            Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ByteString -> IO (Maybe ByteString))
-> Maybe ByteString -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (String -> ByteString
BC.pack String
err)

--
-- Internal
--

iterString :: Iterator
           -> (Iterator -> Ptr CSize -> IO CString)
           -> IO (Maybe ByteString)
iterString :: Iterator
-> (Iterator -> Ptr CSize -> IO (Ptr CChar))
-> IO (Maybe ByteString)
iterString Iterator
iter_ptr Iterator -> Ptr CSize -> IO (Ptr CChar)
f =
    Iterator -> IO CUChar
c_rocksdb_iter_valid Iterator
iter_ptr IO CUChar
-> (CUChar -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        CUChar
0 -> Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
        CUChar
_ -> (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 ->
            Iterator -> Ptr CSize -> IO (Ptr CChar)
f Iterator
iter_ptr Ptr CSize
len_ptr IO (Ptr CChar)
-> (Ptr CChar -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Ptr CChar
ptr ->
            if Ptr CChar
ptr Ptr CChar -> Ptr CChar -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr CChar
forall a. Ptr a
nullPtr
                then Maybe ByteString -> IO (Maybe ByteString)
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
BS.packCStringLen (Ptr CChar
ptr, CSize -> Int
cSizeToInt CSize
len)