sqlite-collections 0.1.0

use crate::serializer::Serializer;
use crate::Savepointable;
use crate::{db, identifier::Identifier};
use rusqlite::{params, Connection, OptionalExtension, Savepoint};

use std::ops::Range;
use std::{borrow::Borrow, marker::PhantomData};

mod error;
// mod iter;
// pub use iter::Iter;

pub use error::{Error, OpenError};

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

impl Default for Config<'static, 'static> {
    fn default() -> Self {
        Config {
            database: "main".try_into().unwrap(),
            table_base: "btree::set".try_into().unwrap(),
            max_elements_per_node: 64,
        }
    }
}

#[derive(Debug, Clone, Copy)]
enum ElementPosition {
    Exists { node: i64, index: u16 },
    Prospective { node: i64, index: u16 },
}

/// A b-tree set.  The database does not ensure uniqueness of the elements.
/// Uniqueness is checked at the application level via the Ord trait. This
/// is useful for non-deterministic serialization, but is slower and more
/// complicated than the DSSet.
#[derive(Debug)]
pub struct BTreeSet<'db, S, C>
where
    S: Serializer,
    C: Savepointable,
{
    connection: C,
    database: Identifier<'db>,
    nodes_table: Identifier<'static>,
    elements_table: Identifier<'static>,
    max_elements_per_node: u16,
    serializer: PhantomData<S>,
}

impl<S, C> BTreeSet<'static, S, C>
where
    S: Serializer,
    C: Savepointable,
{
    pub fn open(connection: C) -> Result<Self, OpenError> {
        BTreeSet::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 {
        BTreeSet::unchecked_open_with_config(connection, Config::default())
    }
}

impl<'db, S, C> BTreeSet<'db, S, C>
where
    S: Serializer,
    C: Savepointable,
{
    pub fn open_with_config(mut connection: C, config: Config<'db, '_>) -> Result<Self, OpenError> {
        let database = config.database;
        let nodes_table = config.table_base.clone() + &Identifier::try_from("::nodes").unwrap();
        let elements_table = config.table_base + &Identifier::try_from("::elements").unwrap();
        let nodes_index = nodes_table.clone() + &Identifier::try_from("::index").unwrap();
        let max_elements_per_node = config.max_elements_per_node;

        {
            let sp = connection.savepoint()?;

            let mut nodes_version = db::setup(
                &sp,
                &database,
                &nodes_table,
                &format!("btree::set::nodes<{max_elements_per_node}>"),
            )?;
            let mut elements_version = db::setup(
                &sp,
                &database,
                &nodes_table,
                &format!("btree::set::elements<{max_elements_per_node}>"),
            )?;
            if nodes_version < 0 {
                return Err(OpenError::TableVersion(nodes_version));
            }
            if elements_version < 0 {
                return Err(OpenError::TableVersion(elements_version));
            }
            let prev_nodes_version = nodes_version;
            if nodes_version < 1 {
                let trailer = db::strict();
                let sql_type = S::sql_type();

                sp.execute(
                    &format!(
                        "CREATE TABLE {database}.{nodes_table} (
                            id INTEGER PRIMARY KEY NOT NULL,
                            parent INTEGER NULL DEFAULT NULL REFERENCES {database}.{nodes_table} (id) ON DELETE CASCADE,
                            index INTEGER NULL DEFAULT NULL,
                            CHECK ((parent IS NULL AND index IS NULL) OR (parent IS NOT NULL AND index IS NOT NULL))
                        ){trailer}"
                    ),
                    [],
                )?;

                sp.execute(
                    &format!(
                        "CREATE UNIQUE INDEX {database}.{nodes_index} ON (parent ASC, index ASC)"
                    ),
                    [],
                )?;

                // Create root node.
                sp.execute(
                    &format!("INSERT INTO {database}.{nodes_table} DEFAULT VALUES"),
                    [],
                )?;
                nodes_version = 1;
            }
            if nodes_version > 1 {
                return Err(OpenError::TableVersion(nodes_version));
            }
            if prev_nodes_version != nodes_version {
                db::set_version(&sp, &database, &nodes_table, nodes_version)?;
            }
            let prev_elements_version = elements_version;
            if elements_version < 1 {
                let trailer = db::strict_without_rowid();
                let sql_type = S::sql_type();

                sp.execute(
                    &format!(
                        "CREATE TABLE {database}.{elements_table} (
                            node INTEGER NOT NULL REFERENCES {database}.{nodes_table} (id) ON DELETE RESTRICT,
                            index INTEGER NOT NULL,
                            key {sql_type} NOT NULL,
                            PRIMARY KEY (node, index)
                        ){trailer}"
                    ),
                    [],
                )?;
                elements_version = 1;
            }
            if elements_version > 1 {
                return Err(OpenError::TableVersion(elements_version));
            }
            if prev_elements_version != elements_version {
                db::set_version(&sp, &database, &elements_table, elements_version)?;
            }

            sp.commit()?;
        }
        Ok(Self {
            connection,
            database,
            elements_table,
            nodes_table,
            max_elements_per_node,
            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, '_>) -> Self {
        let database = config.database;
        let nodes_table = config.table_base;
        let nodes_table = config.table_base.clone() + &Identifier::try_from("::nodes").unwrap();
        let elements_table = config.table_base + &Identifier::try_from("::elements").unwrap();

        Self {
            connection,
            database,
            elements_table,
            nodes_table,
            max_elements_per_node: config.max_elements_per_node,
            serializer: PhantomData,
        }
    }

    /// Binary search for element insertion point in node
    fn find_element_in_node(
        database: &Identifier,
        nodes_table: &Identifier,
        elements_table: &Identifier,
        connection: &Connection,
        value: &S::TargetBorrowed,
        node: i64,
        search: Option<Range<u16>>,
    ) -> Result<ElementPosition, Error<S>>
    where
        S::TargetBorrowed: Ord,
    {
        match search {
            Some(Range { start, end }) => {
                if start <= end {
                    Ok(ElementPosition::Prospective { node, index: start })
                } else {
                    let mid = start / 2 + end / 2;
                    let buffer: S::Buffer = connection
                        .prepare_cached(&format!(
                            "SELECT value FROM {database}.{elements_table} WHERE node = ? AND index = ?"
                        ))?
                        .query_row(params![node, mid], |row| row.get(0))?;
                    let checker = S::deserialize(buffer).map_err(|e| Error::Deserialize(e))?;

                    match checker.cmp(value) {
                        std::cmp::Ordering::Less => Self::find_element_in_node(
                            database,
                            nodes_table,
                            elements_table,
                            connection,
                            value,
                            node,
                            Some(start..mid),
                        ),
                        std::cmp::Ordering::Equal => {
                            Ok(ElementPosition::Exists { node, index: mid })
                        }
                        std::cmp::Ordering::Greater => Self::find_element_in_node(
                            database,
                            nodes_table,
                            elements_table,
                            connection,
                            value,
                            node,
                            Some((mid + 1)..end),
                        ),
                    }
                }
            }
            None => {
                let node_length = connection
                    .prepare_cached(&format!(
                        "SELECT COUNT(*) FROM {database}.{elements_table} WHERE node = ?"
                    ))?
                    .query_row(params![node], |row| row.get(0))?;

                Self::find_element_in_node(
                    database,
                    nodes_table,
                    elements_table,
                    connection,
                    value,
                    node,
                    Some(0..node_length),
                )
            }
        }
    }

    /// Find the position where the element already exists, or where it would
    /// be inserted.
    /// Exists may be in any node, Prospective will always be a leaf node.
    fn find_element_position(
        database: &Identifier,
        nodes_table: &Identifier,
        elements_table: &Identifier,
        connection: &Connection,
        value: &S::TargetBorrowed,
        node: Option<i64>,
    ) -> Result<ElementPosition, Error<S>>
    where
        S::TargetBorrowed: Ord,
    {
        match node {
            Some(node) => {
                match Self::find_element_in_node(
                    database,
                    nodes_table,
                    elements_table,
                    connection,
                    value,
                    node,
                    None,
                )? {
                    ElementPosition::Exists { node, index } => {
                        Ok(ElementPosition::Exists { node, index })
                    }
                    ElementPosition::Prospective { node, index } => {
                        let node_id = connection
                            .prepare_cached(&format!("SELECT id FROM {database}.{nodes_table} WHERE parent = ? AND index = ?"))?
                            .query_row(params![node, index], |row| row.get(0))
                            .optional()?;
                        match node_id {
                            Some(node) => Self::find_element_position(
                                database,
                                nodes_table,
                                elements_table,
                                connection,
                                value,
                                node,
                            ),
                            None => Ok(ElementPosition::Prospective { node, index }),
                        }
                    }
                }
            }
            None => {
                let id = connection
                    .prepare_cached(&format!("SELECT id FROM {database}.{nodes_table} WHERE parent IS NULL AND index IS NULL"))?
                    .query_row([], |row| row.get(0))?;

                Self::find_element_position(
                    database,
                    nodes_table,
                    elements_table,
                    connection,
                    value,
                    Some(id),
                )
            }
        }
    }

    pub fn insert(&mut self, value: &S::TargetBorrowed) -> Result<bool, Error<S>> {
        // TODO
        let database = &self.database;
        let table = &self.table;
        let serialized = match S::serialize(value) {
            Ok(s) => s,
            Err(e) => return Err(Error::Serialize(e)),
        };

        let sp = self.connection.savepoint()?;
        let ret = if db::has_upsert() {
            sp.prepare_cached(&format!(
                "INSERT INTO {database}.{table} (key) VALUES (?) ON CONFLICT DO NOTHING"
            ))?
            .execute(params![serialized])?;
            sp.changes() > 0
        } else if Self::contains_serialized(database, table, &sp, serialized)? {
            false
        } else {
            sp.prepare_cached(&format!("INSERT INTO {database}.{table} (key) VALUES (?)"))?
                .execute(params![serialized])?;
            true
        };
        sp.commit()?;
        Ok(ret)
    }

    pub fn contains(&mut self, value: &S::TargetBorrowed) -> Result<bool, Error<S>> {
        // TODO
        let serialized = match S::serialize(value) {
            Ok(s) => s,
            Err(e) => return Err(Error::Serialize(e)),
        };
        Self::contains_serialized(
            &self.database,
            &self.table,
            &*self.connection.savepoint()?,
            serialized,
        )
    }

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

    pub fn remove<Q>(&mut self, value: &S::TargetBorrowed) -> Result<bool, Error<S>> {
        // TODO
        let database = &self.database;
        let table = &self.table;
        let serialized = match S::serialize(value) {
            Ok(s) => s,
            Err(e) => return Err(Error::Serialize(e)),
        };

        let sp = self.connection.savepoint()?;
        let changes = sp
            .prepare_cached(&format!("DELETE FROM {database}.{table} WHERE key = ?"))?
            .execute(params![serialized])?;

        sp.commit()?;

        Ok(changes > 0)
    }

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

    pub fn first(&mut self) -> Result<Option<S::Target>, Error<S>> {
        // TODO
        let database = &self.database;
        let table = &self.table;

        let serialized: Option<S::Buffer> = self
            .connection
            .savepoint()?
            .prepare_cached(&format!(
                "SELECT key FROM {database}.{table} ORDER BY key ASC"
            ))?
            .query_row([], |row| row.get(0))
            .optional()?;

        match serialized.map(|s| S::deserialize(s)).transpose() {
            Ok(s) => Ok(s),
            Err(e) => Err(Error::Deserialize(e)),
        }
    }

    pub fn last(&mut self) -> Result<Option<S::Target>, Error<S>> {
        // TODO
        let database = &self.database;
        let table = &self.table;
        let serialized: Option<S::Buffer> = self
            .connection
            .savepoint()?
            .prepare_cached(&format!(
                "SELECT key FROM {database}.{table} ORDER BY key DESC"
            ))?
            .query_row([], |row| row.get(0))
            .optional()?;
        match serialized.map(|s| S::deserialize(s)).transpose() {
            Ok(s) => Ok(s),
            Err(e) => Err(Error::Deserialize(e)),
        }
    }

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

    pub fn iter(&mut self) -> Result<Iter<'db, S, Savepoint<'_>>, Error<S>> {
        // TODO
        Ok(Iter::new(
            self.connection.savepoint()?,
            self.database.clone(),
            self.table.clone(),
        )?)
    }
}

#[cfg(test)]
mod test;