typed_store/rocks/
keys.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::marker::PhantomData;
6
7use bincode::Options;
8use serde::{Serialize, de::DeserializeOwned};
9
10use super::{RocksDBRawIter, TypedStoreError, be_fix_int_ser};
11
12/// An iterator over the keys of a prefix.
13pub struct Keys<'a, K> {
14    db_iter: RocksDBRawIter<'a>,
15    _phantom: PhantomData<K>,
16}
17
18impl<'a, K: DeserializeOwned> Keys<'a, K> {
19    pub(crate) fn new(db_iter: RocksDBRawIter<'a>) -> Self {
20        Self {
21            db_iter,
22            _phantom: PhantomData,
23        }
24    }
25}
26
27impl<K: DeserializeOwned> Iterator for Keys<'_, K> {
28    type Item = Result<K, TypedStoreError>;
29
30    fn next(&mut self) -> Option<Self::Item> {
31        if self.db_iter.valid() {
32            let config = bincode::DefaultOptions::new()
33                .with_big_endian()
34                .with_fixint_encoding();
35            let key = self.db_iter.key().and_then(|k| config.deserialize(k).ok());
36            self.db_iter.next();
37            key.map(Ok)
38        } else {
39            match self.db_iter.status() {
40                Ok(_) => None,
41                Err(err) => Some(Err(TypedStoreError::RocksDB(format!("{err}")))),
42            }
43        }
44    }
45}
46
47impl<K: Serialize> Keys<'_, K> {
48    /// Skips all the elements that are smaller than the given key,
49    /// and either lands on the key or the first one greater than
50    /// the key.
51    pub fn skip_to(mut self, key: &K) -> Result<Self, TypedStoreError> {
52        self.db_iter.seek(be_fix_int_ser(key)?);
53        Ok(self)
54    }
55
56    /// Moves the iterator the element given or
57    /// the one prior to it if it does not exist. If there is
58    /// no element prior to it, it returns an empty iterator.
59    pub fn skip_prior_to(mut self, key: &K) -> Result<Self, TypedStoreError> {
60        self.db_iter.seek_for_prev(be_fix_int_ser(key)?);
61        Ok(self)
62    }
63
64    /// Seeks to the last key in the database (at this column family).
65    pub fn skip_to_last(mut self) -> Self {
66        self.db_iter.seek_to_last();
67        self
68    }
69}