sqlite-collections 0.1.0

use crate::format::Format;
use crate::OpenError;
use crate::Savepointable;
use crate::{db, identifier::Identifier};
use rusqlite::Savepoint;
use rusqlite::{params, Connection, OptionalExtension};

use std::marker::PhantomData;

mod error;
pub mod iter;

use iter::KeyIter;
use iter::KeyValueIter;
use iter::ValueIter;

pub use error::Error;

#[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Eq, Hash)]
pub struct Config<'db, 'tbl> {
    pub database: Identifier<'db>,
    pub table: Identifier<'tbl>,
}

impl Default for Config<'static, 'static> {
    fn default() -> Self {
        Config {
            database: "main".try_into().unwrap(),
            table: "ds::map".try_into().unwrap(),
        }
    }
}

/// Deterministic store set.
pub struct Map<'db, 'tbl, K, V, C>
where
    K: Format,
    V: Format,
    C: Savepointable,
{
    connection: C,
    database: Identifier<'db>,
    table: Identifier<'tbl>,
    key_serializer: PhantomData<K>,
    value_serializer: PhantomData<V>,
}

impl<K, V, C> Map<'static, 'static, K, V, C>
where
    K: Format,
    V: Format,
    C: Savepointable,
{
    pub fn open(connection: C) -> Result<Self, OpenError> {
        Map::open_with_config(connection, Config::default())
    }

    /// Open a set without creating it or checking if it exists.  This is safe
    /// if you call a safe open in (or under) the same transaction or savepoint
    /// beforehand.
    pub fn unchecked_open(connection: C) -> Self {
        Map::unchecked_open_with_config(connection, Config::default())
    }
}

impl<'db, 'tbl, K, V, C> Map<'db, 'tbl, K, V, C>
where
    K: Format,
    V: Format,
    C: Savepointable,
{
    pub fn open_with_config(
        mut connection: C,
        config: Config<'db, 'tbl>,
    ) -> Result<Self, OpenError> {
        let database = config.database;
        let table = config.table;

        {
            let sp = connection.savepoint()?;

            let mut version = db::setup(&sp, &database, &table, "ds::map")?;
            if version < 0 {
                return Err(OpenError::TableVersion(version));
            }
            let prev_version = version;
            if version < 1 {
                let trailer = db::strict_without_rowid();
                let sql_type = K::sql_type();

                sp.execute(
                    &format!(
                        "CREATE TABLE {database}.{table} (
                            key {sql_type} UNIQUE PRIMARY KEY NOT NULL,
                            value {sql_type} NOT NULL
                        ){trailer}"
                    ),
                    [],
                )?;
                version = 1;
            }
            if version > 1 {
                return Err(OpenError::TableVersion(version));
            }
            if prev_version != version {
                db::set_version(&sp, &database, &table, version)?;
            }

            sp.commit()?;
        }
        Ok(Self {
            connection,
            database,
            table,
            key_serializer: PhantomData,
            value_serializer: PhantomData,
        })
    }

    /// Open a set without creating it or checking if it exists.  This is safe
    /// if you call a safe open in (or under) the same transaction or savepoint
    /// beforehand.
    pub fn unchecked_open_with_config(connection: C, config: Config<'db, 'tbl>) -> Self {
        let database = config.database;
        let table = config.table;

        Self {
            connection,
            database,
            table,
            key_serializer: PhantomData,
            value_serializer: PhantomData,
        }
    }

    pub fn insert(&mut self, key: &K::In, value: &V::In) -> Result<Option<V::Out>, Error<K, V>> {
        let database = &self.database;
        let table = &self.table;
        let key = K::serialize(key).map_err(|e| Error::KeySerialize(e))?;
        let value = V::serialize(value).map_err(|e| Error::ValueSerialize(e))?;

        let sp = self.connection.savepoint()?;
        let prev_value = Self::get_from_serialized(database, table, &sp, &key)?;
        if db::has_upsert() {
            sp.prepare_cached(&format!(
                "INSERT INTO {database}.{table} (key, value) VALUES (?, ?) ON CONFLICT DO UPDATE SET value=excluded.value"
            ))?
            .execute(params![key, value])?;
        } else if prev_value.is_some() {
            sp.prepare_cached(&format!(
                "UPDATE {database}.{table} SET value=? WHERE key=?"
            ))?
            .execute(params![value, key])?;
        } else {
            sp.prepare_cached(&format!(
                "INSERT INTO {database}.{table} (key, value) VALUES (?, ?)"
            ))?
            .execute(params![key, value])?;
        };
        sp.commit()?;
        Ok(prev_value)
    }

    pub fn len(&mut self) -> Result<u64, Error<K, V>> {
        let database = &self.database;
        let table = &self.table;
        Ok(self
            .connection
            .savepoint()?
            .prepare_cached(&format!("SELECT COUNT(*) FROM {database}.{table}"))?
            .query_row([], |row| row.get(0))?)
    }

    pub fn is_empty(&mut self) -> Result<bool, Error<K, V>> {
        let database = &self.database;
        let table = &self.table;
        Ok(self
            .connection
            .savepoint()?
            .prepare_cached(&format!("SELECT 1 FROM {database}.{table} LIMIT 1"))?
            .query_row([], |_| Ok(()))
            .optional()?
            .is_none())
    }

    fn contains_from_serialized(
        database: &Identifier,
        table: &Identifier,
        connection: &Connection,
        key: &K::Buffer,
    ) -> Result<bool, Error<K, V>> {
        Ok(connection
            .prepare_cached(&format!("SELECT 1 FROM {database}.{table} WHERE key = ?"))?
            .query_row(params![key], |_| Ok(()))
            .optional()?
            .is_some())
    }

    fn get_from_serialized(
        database: &Identifier,
        table: &Identifier,
        connection: &Connection,
        key: &K::Buffer,
    ) -> Result<Option<V::Out>, Error<K, V>> {
        Ok(
            Self::get_serialized_from_serialized(database, table, connection, key)?
                .map(|b| V::deserialize(&b).map_err(|e| Error::ValueDeserialize(e)))
                .transpose()?,
        )
    }

    fn get_serialized_from_serialized(
        database: &Identifier,
        table: &Identifier,
        connection: &Connection,
        key: &K::Buffer,
    ) -> Result<Option<V::Buffer>, Error<K, V>> {
        Ok(connection
            .prepare_cached(&format!(
                "SELECT value FROM {database}.{table} WHERE key = ?"
            ))?
            .query_row(params![key], |row| row.get(0))
            .optional()?)
    }

    pub fn get(&mut self, key: &K::In) -> Result<Option<V::Out>, Error<K, V>> {
        let database = &self.database;
        let table = &self.table;
        let key = K::serialize(key).map_err(|e| Error::KeySerialize(e))?;

        let sp = self.connection.savepoint()?;
        let result = Self::get_from_serialized(database, table, &sp, &key)?;
        sp.commit()?;
        Ok(result)
    }

    pub fn remove(&mut self, key: &K::In) -> Result<Option<V::Out>, Error<K, V>> {
        let database = &self.database;
        let table = &self.table;
        let key = K::serialize(key).map_err(|e| Error::KeySerialize(e))?;

        let sp = self.connection.savepoint()?;
        let result = Self::get_from_serialized(database, table, &sp, &key)?;
        sp.prepare_cached(&format!("DELETE FROM {database}.{table} WHERE key = ?"))?
            .execute(params![key])?;
        sp.commit()?;
        Ok(result)
    }

    pub fn contains_key(&mut self, key: &K::In) -> Result<bool, Error<K, V>> {
        let database = &self.database;
        let table = &self.table;
        let key = K::serialize(key).map_err(|e| Error::KeySerialize(e))?;

        let sp = self.connection.savepoint()?;
        let result = Self::contains_from_serialized(database, table, &sp, &key)?;
        sp.commit()?;
        Ok(result)
    }

    pub fn clear(&mut self) -> Result<(), Error<K, V>> {
        let database = &self.database;
        let table = &self.table;
        let sp = self.connection.savepoint()?;
        sp.prepare_cached(&format!("DELETE FROM {database}.{table}"))?
            .execute([])?;
        sp.commit()?;
        Ok(())
    }

    pub fn iter(&mut self) -> Result<KeyValueIter<'db, 'tbl, K, V, Savepoint<'_>>, Error<K, V>> {
        Ok(KeyValueIter::new(
            self.connection.savepoint()?,
            self.database.clone(),
            self.table.clone(),
        )?)
    }

    pub fn keys(&mut self) -> Result<KeyIter<'db, 'tbl, K, V, Savepoint<'_>>, Error<K, V>> {
        Ok(KeyIter::new(
            self.connection.savepoint()?,
            self.database.clone(),
            self.table.clone(),
        )?)
    }
    pub fn values(&mut self) -> Result<ValueIter<'db, 'tbl, K, V, Savepoint<'_>>, Error<K, V>> {
        Ok(ValueIter::new(
            self.connection.savepoint()?,
            self.database.clone(),
            self.table.clone(),
        )?)
    }

    /// Retains only the elements specified by the predicate.
    ///
    /// In other words, remove all pairs (k, v) for which f(k, v) returns false.
    /// The elements are visited in unsorted (and unspecified) order.
    ///
    /// This is all done in a single transaction.
    pub fn retain<F>(&mut self, mut f: F) -> Result<(), Error<K, V>>
    where
        F: FnMut(K::Out, V::Out) -> bool,
    {
        let database = &self.database;
        let table = &self.table;

        let sp = self.connection.savepoint()?;
        {
            let mut maybe_serialized = sp
                .prepare_cached(&format!(
                    "SELECT key, value FROM {database}.{table} ORDER BY key ASC LIMIT 1"
                ))?
                .query_row([], |row| Ok((row.get(0)?, row.get(1)?)))
                .optional()?;
            let mut deleter =
                sp.prepare_cached(&format!("DELETE FROM {database}.{table} WHERE key = ?"))?;
            let mut select_next = sp.prepare_cached(&format!(
                "SELECT key, value FROM {database}.{table} WHERE key > ? ORDER BY key ASC LIMIT 1"
            ))?;
            while let Some((serialized_key, value)) = maybe_serialized {
                let key = K::deserialize(&serialized_key).map_err(Error::KeyDeserialize)?;
                let value = V::deserialize(&value).map_err(Error::ValueDeserialize)?;
                if !f(key, value) {
                    deleter.execute(params![serialized_key])?;
                }
                maybe_serialized = select_next
                    .query_row(params![serialized_key], |row| {
                        Ok((row.get(0)?, row.get(1)?))
                    })
                    .optional()?;
            }
        }
        sp.commit()?;
        Ok(())
    }
}

#[cfg(test)]
mod test;