menhirkv 0.4.5

// Copyright (C) 2024 Christian Mauduit <ufoot@ufoot.org>

use crate::error::{Error, Result};
use crate::serial::*;
use ofilter::SyncStream;
use rocksdb::DBIteratorWithThreadMode;
use rocksdb::DB;
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::fmt;
use std::marker::PhantomData;

/// Iterator over a KV store.
///
///
/// IMPORTANT: the iterator will only iterate on entries which would
/// be removed by the next compaction. This means you could have
/// an item which is retrievable by `get()` or `peek()`, yet will
/// never be yielded by the iterator.
///
/// The idea is that when you
/// iterate, you care about what is "recent" and "important".
/// Old items which are still in the database are still available
/// if you know their keys, but they are not discoverable through
/// iterators.
///
/// ```
/// use menhirkv::Store;
///
/// let store = Store::<usize, usize>::open_temporary(100).unwrap();
/// store.put(&1, &2).unwrap();
/// store.put(&3, &4).unwrap();
/// store.put(&5, &6).unwrap();
/// for item in store.iter().unwrap() {
///     let item = item.unwrap();
///     println!("key: {}, value: {}", item.0, item.1);
/// }
/// ```
// #[derive(Debug)]
pub struct Iter<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub(crate) inner: DBIteratorWithThreadMode<'a, DB>,
    pub(crate) filter: SyncStream<Vec<u8>>,
    pub(crate) done: usize,
    pub(crate) phantom_data: PhantomData<(K, V)>,
}

impl<K, V> fmt::Display for Iter<'_, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}/?", self.done)
    }
}

impl<'a, K, V> Iterator for Iter<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    type Item = Result<(K, V)>;

    fn next(&mut self) -> Option<Result<(K, V)>> {
        loop {
            if let Some(res) = self.inner.next() {
                match res {
                    Ok(item) => {
                        let (kbuf, vbuf) = item;
                        let kvec: &Vec<u8> = &kbuf.to_vec();
                        if !self.filter.check(kvec) {
                            continue;
                        }
                        return match deserialize_key(kbuf) {
                            Ok(k) => match deserialize_value(vbuf) {
                                Ok(v) => Some(Ok((k, v))),
                                Err(e) => Some(Err(e)),
                            },
                            Err(e) => Some(Err(e)),
                        };
                    }
                    Err(e) => return Some(Err(Error::from(e))),
                }
            } else {
                break;
            }
        }
        None
    }
}

/// Iterator over the keys of a KV store.
///
/// Only yield keys that would still be here after a compaction.
///
/// ```
/// use menhirkv::Store;
///
/// let store = Store::<usize, usize>::open_temporary(100).unwrap();
/// store.put(&1, &2).unwrap();
/// store.put(&3, &4).unwrap();
/// store.put(&5, &6).unwrap();
/// for key in store.keys().unwrap() {
///     let key = key.unwrap();
///     println!("key: {}", key);
/// }
/// ```
// #[derive(Debug)]
pub struct Keys<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub(crate) inner: DBIteratorWithThreadMode<'a, DB>,
    pub(crate) filter: SyncStream<Vec<u8>>,
    pub(crate) done: usize,
    pub(crate) phantom_data: PhantomData<(K, V)>,
}

impl<K, V> fmt::Display for Keys<'_, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}/?", self.done)
    }
}

impl<'a, K, V> Iterator for Keys<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    type Item = Result<K>;

    fn next(&mut self) -> Option<Result<K>> {
        loop {
            if let Some(res) = self.inner.next() {
                match res {
                    Ok(item) => {
                        let kvec: &Vec<u8> = &item.0.to_vec();
                        if !self.filter.check(kvec) {
                            continue;
                        }
                        return Some(deserialize_key(item.0));
                    }
                    Err(e) => return Some(Err(Error::from(e))),
                }
            } else {
                break;
            }
        }
        None
    }
}

/// Iterator over the values of a KV store.
///
/// Only yield values that would still be here after a compaction.
///
/// ```
/// use menhirkv::Store;
///
/// let store = Store::<usize, usize>::open_temporary(100).unwrap();
/// store.put(&1, &2).unwrap();
/// store.put(&3, &4).unwrap();
/// store.put(&5, &6).unwrap();
/// for value in store.values().unwrap() {
///     let value = value.unwrap();
///     println!("value: {}", value);
/// }
/// ```
// #[derive(Debug)]
pub struct Values<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub(crate) inner: DBIteratorWithThreadMode<'a, DB>,
    pub(crate) filter: SyncStream<Vec<u8>>,
    pub(crate) done: usize,
    pub(crate) phantom_data: PhantomData<(K, V)>,
}

impl<K, V> fmt::Display for Values<'_, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}/?", self.done)
    }
}

impl<'a, K, V> Iterator for Values<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    type Item = Result<V>;

    fn next(&mut self) -> Option<Result<V>> {
        loop {
            if let Some(res) = self.inner.next() {
                match res {
                    Ok(item) => {
                        let (kbuf, vbuf) = item;
                        let kvec: &Vec<u8> = &kbuf.to_vec();
                        if !self.filter.check(kvec) {
                            continue;
                        }
                        return Some(deserialize_value(vbuf));
                    }
                    Err(e) => return Some(Err(Error::from(e))),
                }
            } else {
                break;
            }
        }
        None
    }
}

/// Export a KV store.
///
/// The difference between this and a standard iterator is that
/// this one returns `(K, V)` instead of `Result<(K, V)>`.
/// A consequence of this is that it may `panic()` at any iteration,
/// as there can always be an error, and it will not catch it.
///
/// If you want something failsafe, use the standard iterator.
/// This one is easier to use, but can crash your program.
///
/// ```
/// use menhirkv::Store;
///
/// let store = Store::<usize, usize>::open_temporary(100).unwrap();
/// store.put(&1, &2).unwrap();
/// store.put(&3, &4).unwrap();
/// store.put(&5, &6).unwrap();
/// let export = store.export().unwrap().collect::<Vec<(usize, usize)>>(); // may panic()
/// assert_eq!(3, export.len());
/// ```
// #[derive(Debug)]
pub struct Export<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub(crate) inner: DBIteratorWithThreadMode<'a, DB>,
    pub(crate) filter: SyncStream<Vec<u8>>,
    pub(crate) done: usize,
    pub(crate) phantom_data: PhantomData<(K, V)>,
}

impl<K, V> fmt::Display for Export<'_, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}/?", self.done)
    }
}

impl<'a, K, V> Export<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub fn try_next(&mut self) -> Result<Option<(K, V)>> {
        loop {
            if let Some(res) = self.inner.next() {
                match res {
                    Ok(item) => {
                        let (kbuf, vbuf) = item;
                        let kvec: &Vec<u8> = &kbuf.to_vec();
                        if !self.filter.check(kvec) {
                            continue;
                        }
                        return match deserialize_key(kbuf) {
                            Ok(k) => match deserialize_value(vbuf) {
                                Ok(v) => Ok(Some((k, v))),
                                Err(e) => Err(e),
                            },
                            Err(e) => Err(e),
                        };
                    }
                    Err(e) => return Err(Error::from(e)),
                }
            } else {
                break;
            }
        }
        Ok(None)
    }
}

impl<'a, K, V> Iterator for Export<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    type Item = (K, V);

    fn next(&mut self) -> Option<(K, V)> {
        match self.try_next() {
            Ok(item) => item,
            Err(e) => panic!("export error: {}", e),
        }
    }
}

/// Export the keys of a KV store.
///
/// The difference between this and a standard iterator is that
/// this one returns `K` instead of `Result<K>`.
/// A consequence of this is that it may `panic()` at any iteration,
/// as there can always be an error, and it will not catch it.
///
/// If you want something failsafe, use the standard iterator.
/// This one is easier to use, but can crash your program.
///
/// ```
/// use menhirkv::Store;
///
/// let store = Store::<usize, usize>::open_temporary(100).unwrap();
/// store.put(&1, &2).unwrap();
/// store.put(&3, &4).unwrap();
/// store.put(&5, &6).unwrap();
/// let mut keys = store.export_keys().unwrap().collect::<Vec<usize>>(); // may panic()
/// keys.sort();
/// assert_eq!(vec![1, 3, 5], keys);
/// ```
// #[derive(Debug)]
pub struct ExportKeys<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub(crate) inner: DBIteratorWithThreadMode<'a, DB>,
    pub(crate) filter: SyncStream<Vec<u8>>,
    pub(crate) done: usize,
    pub(crate) phantom_data: PhantomData<(K, V)>,
}

impl<K, V> fmt::Display for ExportKeys<'_, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}/?", self.done)
    }
}

impl<'a, K, V> ExportKeys<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub fn try_next(&mut self) -> Result<Option<K>> {
        loop {
            if let Some(res) = self.inner.next() {
                match res {
                    Ok(item) => {
                        let kvec: &Vec<u8> = &item.0.to_vec();
                        if !self.filter.check(kvec) {
                            continue;
                        }
                        return match deserialize_key(item.0) {
                            Ok(k) => Ok(Some(k)),
                            Err(e) => Err(e),
                        };
                    }
                    Err(e) => return Err(Error::from(e)),
                }
            } else {
                break;
            }
        }
        Ok(None)
    }
}

impl<'a, K, V> Iterator for ExportKeys<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    type Item = K;

    fn next(&mut self) -> Option<K> {
        match self.try_next() {
            Ok(item) => item,
            Err(e) => panic!("export error: {}", e),
        }
    }
}

/// Export the values of a KV store.
///
/// The difference between this and a standard iterator is that
/// this one returns `V` instead of `Result<V>`.
/// A consequence of this is that it may `panic()` at any iteration,
/// as there can always be an error, and it will not catch it.
///
/// If you want something failsafe, use the standard iterator.
/// This one is easier to use, but can crash your program.
///
/// ```
/// use menhirkv::Store;
///
/// let store = Store::<usize, usize>::open_temporary(100).unwrap();
/// store.put(&1, &2).unwrap();
/// store.put(&3, &4).unwrap();
/// store.put(&5, &6).unwrap();
/// let mut values = store.export_values().unwrap().collect::<Vec<usize>>(); // may panic()
/// values.sort();
/// assert_eq!(vec![2, 4, 6], values);
/// ```
// #[derive(Debug)]
pub struct ExportValues<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub(crate) inner: DBIteratorWithThreadMode<'a, DB>,
    pub(crate) filter: SyncStream<Vec<u8>>,
    pub(crate) done: usize,
    pub(crate) phantom_data: PhantomData<(K, V)>,
}

impl<K, V> fmt::Display for ExportValues<'_, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}/?", self.done)
    }
}

impl<'a, K, V> ExportValues<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    pub fn try_next(&mut self) -> Result<Option<V>> {
        loop {
            if let Some(res) = self.inner.next() {
                match res {
                    Ok(item) => {
                        let (kbuf, vbuf) = item;
                        let kvec: &Vec<u8> = &kbuf.to_vec();
                        if !self.filter.check(kvec) {
                            continue;
                        };
                        return match deserialize_value(vbuf) {
                            Ok(v) => Ok(Some(v)),
                            Err(e) => Err(e),
                        };
                    }
                    Err(e) => return Err(Error::from(e)),
                }
            } else {
                break;
            }
        }
        Ok(None)
    }
}

impl<'a, K, V> Iterator for ExportValues<'a, K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    type Item = V;

    fn next(&mut self) -> Option<V> {
        match self.try_next() {
            Ok(item) => item,
            Err(e) => panic!("export error: {}", e),
        }
    }
}

/// Iterator over the column family names of a KV store.
///
/// Use this to list the name of all column family names.
/// It will also yield the empty string, corresponding to
/// the main/root column family you open by default.
///
/// These are not necessarily exactly the names used by RocksDB
/// internally, there can be some marshalling. Typically,
/// a prefix is added. Unless you inspect the content of the
/// inner database yourself, do not worry about this.
///
/// ```
/// use menhirkv::Store;
///
/// let store = Store::<usize, usize>::open_cf_temporary(&["a", "b"], 100).unwrap();
/// let cf_names = store.iter_cf_names().collect::<Vec<&str>>();
/// assert_eq!(vec!["", "a", "b"], cf_names);
/// ```
pub struct IterCfNames<'a> {
    pub(crate) idx: isize,
    pub(crate) other_cfs: &'a Vec<String>,
}

impl fmt::Display for IterCfNames<'_> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}/{}", self.idx + 1, self.other_cfs.len() + 1)
    }
}

impl<'a> Iterator for IterCfNames<'a> {
    type Item = &'a str;

    fn next(&mut self) -> Option<&'a str> {
        if self.idx < 0 {
            self.idx = 0;
            return Some("");
        }
        if self.idx as usize >= self.other_cfs.len() {
            return None;
        }
        let cf_name = self.other_cfs[self.idx as usize].as_str();
        self.idx += 1;
        Some(cf_name)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.other_cfs.len() + 1 - ((self.idx + 1) as usize);
        (remaining, Some(remaining))
    }
}

impl<'a> ExactSizeIterator for IterCfNames<'a> {}