sqlite-collections 0.1.0

use crate::ds::map::Error;
use crate::format::Format;
use crate::{identifier::Identifier, AsConnection};
use rusqlite::{params, OptionalExtension};

use std::marker::PhantomData;

#[derive(Debug)]
pub struct Iter<'db, 'tbl, K, V, C>
where
    K: Format,
    V: Format,
    C: AsConnection,
{
    connection: C,
    database: Identifier<'db>,
    table: Identifier<'tbl>,
    front: Option<K::Buffer>,
    back: Option<K::Buffer>,
    size: usize,
    key_serializer: PhantomData<K>,
    value_serializer: PhantomData<V>,
}

impl<'db, 'tbl, K, V, C> Iter<'db, 'tbl, K, V, C>
where
    K: Format,
    V: Format,
    C: AsConnection,
{
    pub fn new(
        connection: C,
        database: Identifier<'db>,
        table: Identifier<'tbl>,
    ) -> rusqlite::Result<Self> {
        let mut new = Self {
            connection,
            database,
            table,
            size: 0,
            front: None,
            back: None,
            key_serializer: PhantomData,
            value_serializer: PhantomData,
        };
        new.update_size()?;
        Ok(new)
    }

    fn update_size(&mut self) -> rusqlite::Result<usize> {
        self.size = self.select("COUNT(*)", "", |row| row.get(0))?;
        Ok(self.size)
    }

    fn select<U, F>(&self, selection: &str, order: &str, f: F) -> rusqlite::Result<U>
    where
        F: FnOnce(&rusqlite::Row<'_>) -> rusqlite::Result<U>,
    {
        let database = &self.database;
        let table = &self.table;

        match (self.front.as_ref(), self.back.as_ref()) {
            (None, None) => self
                .connection
                .as_connection()
                .prepare_cached(&format!(
                    "SELECT {selection} FROM {database}.{table} {order} LIMIT 1"
                ))?
                .query_row([], f),
            (Some(front), None) => self
                .connection
                .as_connection()
                .prepare_cached(&format!(
                    "SELECT {selection} FROM {database}.{table} WHERE key > ? {order} LIMIT 1"
                ))?
                .query_row(params![front], f),
            (None, Some(back)) => self
                .connection
                .as_connection()
                .prepare_cached(&format!(
                    "SELECT {selection} FROM {database}.{table} WHERE key < ? {order} LIMIT 1"
                ))?
                .query_row(params![back], f),
            (Some(front), Some(back)) => self
                .connection
                .as_connection()
                .prepare_cached(&format!(
                    "SELECT {selection} FROM {database}.{table} WHERE key > ? AND key < ? {order} LIMIT 1"
                ))?
                .query_row(params![front, back], f),
        }
    }

    fn inner_next(&mut self, key: &K::Buffer, value: &V::Buffer) -> <Self as Iterator>::Item {
        let key = K::deserialize(key).map_err(Error::KeyDeserialize)?;
        let value = V::deserialize(value).map_err(Error::ValueDeserialize)?;
        Ok((key, value))
    }
}

impl<'db, 'tbl, K, V, C> Iterator for Iter<'db, 'tbl, K, V, C>
where
    K: Format,
    V: Format,
    C: AsConnection,
{
    type Item = Result<(K::Out, V::Out), Error<K, V>>;

    fn next(&mut self) -> Option<Self::Item> {
        let next: rusqlite::Result<Option<(K::Buffer, V::Buffer)>> = self
            .select("key, value", "ORDER BY key ASC", |row| {
                Ok((row.get(0)?, row.get(1)?))
            })
            .optional();
        match next {
            Ok(Some((key, value))) => {
                let res = self.inner_next(&key, &value);
                self.front = Some(key);
                Some(self.update_size().map_err(Into::into).and(res))
            }
            Ok(None) => {
                self.size = 0;
                None
            }
            Err(e) => Some(Err(e.into())),
        }
    }

    /// return how many elements are in left in the iterator as of this call.
    /// This can change during iteration if you are iterating in autocommit
    /// mode, or hold the iterator across transactions. This is why this
    /// iterator doesn't implement ExactSizeIterator
    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.size, Some(self.size))
    }
}

impl<'db, 'tbl, K, V, C> DoubleEndedIterator for Iter<'db, 'tbl, K, V, C>
where
    K: Format,
    V: Format,
    C: AsConnection,
{
    fn next_back(&mut self) -> Option<Self::Item> {
        let next: rusqlite::Result<Option<(K::Buffer, V::Buffer)>> = self
            .select("key, value", "ORDER BY key DESC", |row| {
                Ok((row.get(0)?, row.get(1)?))
            })
            .optional();
        match next {
            Ok(Some((key, value))) => {
                let res = self.inner_next(&key, &value);
                self.back = Some(key);
                Some(self.update_size().map_err(Into::into).and(res))
            }
            Ok(None) => {
                self.size = 0;
                None
            }
            Err(e) => Some(Err(e.into())),
        }
    }
}