liteboxfs 0.1.0

use std::ops::{Bound, RangeBounds};

use crate::{
    block::{DirtyFilter, FileId, MerkleNode, MerkleNodePosition, MerkleNodeStore},
    errors::InternalError,
    hash::MerkleHash,
};

use super::SqlStore;

impl MerkleNodeStore for SqlStore<'_> {
    fn get_root(&mut self, file: FileId) -> crate::Result<Option<MerkleNode>> {
        let mut stmt = self.db.prepare_cached(
            r#"
            SELECT
                file,
                level,
                "index",
                hash,
                dirty
            FROM
                liteboxfs_merkle_nodes
            WHERE
                file = ?1
                AND level = (
                    SELECT
                        MAX(level)
                    FROM
                        liteboxfs_merkle_nodes
                    WHERE
                        file = ?1
                );
            "#,
        )?;

        let params = rusqlite::params![file];

        let result = stmt.query_one(params, |row| {
            Ok(MerkleNode {
                pos: MerkleNodePosition {
                    file: row.get(0)?,
                    level: row.get(1)?,
                    index: row.get(2)?,
                },
                hash: row.get(3)?,
                dirty: row.get(4)?,
            })
        });

        match result {
            Ok(node) => Ok(Some(node)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(rusqlite::Error::QueryReturnedMoreThanOneRow) => {
                Err(InternalError::MoreThanOneRootMerkleNode.into())
            }
            Err(err) => Err(err.into()),
        }
    }

    fn node_indices(
        &mut self,
        file: FileId,
        level: u32,
        filter: DirtyFilter,
    ) -> crate::Result<Vec<usize>> {
        let mut stmt = self.db.prepare_cached(
            r#"
            SELECT
                "index"
            FROM
                liteboxfs_merkle_nodes
            WHERE
                file = ?1
                AND level = ?2
                AND CASE ?3
                    WHEN 0 THEN dirty = 0
                    WHEN 1 THEN dirty = 1
                    ELSE 1
                END;
            "#,
        )?;

        let params = rusqlite::params![
            file,
            level,
            match filter {
                DirtyFilter::Dirty => Some(1),
                DirtyFilter::NotDirty => Some(0),
                DirtyFilter::All => None,
            },
        ];

        Ok(stmt
            .query_map(params, |row| row.get(0))?
            .collect::<Result<Vec<usize>, _>>()?)
    }

    fn list_nodes<R>(
        &mut self,
        file: FileId,
        level: u32,
        indices: R,
    ) -> crate::Result<Vec<MerkleNode>>
    where
        R: RangeBounds<usize>,
    {
        let mut stmt = self.db.prepare_cached(
            r#"
            SELECT
                file,
                level,
                "index",
                hash,
                dirty
            FROM
                liteboxfs_merkle_nodes
            WHERE
                file = ?1
                AND level = ?2
                AND "index" >= ?3
                AND "index" < COALESCE(
                    ?4,
                    (
                        SELECT
                            MAX("index") + 1
                        FROM
                            liteboxfs_merkle_nodes
                        WHERE
                            file = ?1
                            AND level = ?2
                    )
                );
            "#,
        )?;

        let start_index = match indices.start_bound() {
            Bound::Included(&idx) => idx,
            Bound::Excluded(&idx) => idx + 1,
            Bound::Unbounded => 0,
        };

        // SQL uses `< end_index`, so we need to convert to an exclusive upper bound
        let end_index = match indices.end_bound() {
            Bound::Included(&idx) => Some(idx + 1),
            Bound::Excluded(&idx) => Some(idx),
            Bound::Unbounded => None,
        };

        let params = rusqlite::params![file, level, start_index, end_index];

        Ok(stmt
            .query_map(params, |row| {
                Ok(MerkleNode {
                    pos: MerkleNodePosition {
                        file: row.get(0)?,
                        level: row.get(1)?,
                        index: row.get(2)?,
                    },
                    hash: row.get(3)?,
                    dirty: row.get(4)?,
                })
            })?
            .collect::<Result<Vec<_>, _>>()?)
    }

    fn put_node(&mut self, pos: MerkleNodePosition, hash: MerkleHash) -> crate::Result<()> {
        let mut stmt = self.db.prepare_cached(
            r#"
            INSERT OR REPLACE INTO
                liteboxfs_merkle_nodes (file, level, "index", hash, dirty)
            VALUES
                (?, ?, ?, ?, 0);
            "#,
        )?;

        stmt.execute(rusqlite::params![pos.file, pos.level, pos.index, hash])?;

        let mut stmt = self.db.prepare_cached(
            r#"
            WITH RECURSIVE ancestors AS (
                SELECT
                    file,
                    level,
                    "index"
                FROM
                    liteboxfs_merkle_nodes
                WHERE
                    file = ?1
                    AND level = ?2 + 1
                    AND "index" = ?3 / ?4

                UNION ALL

                SELECT
                    liteboxfs_merkle_nodes.file,
                    liteboxfs_merkle_nodes.level,
                    liteboxfs_merkle_nodes."index"
                FROM
                    ancestors
                JOIN
                    liteboxfs_merkle_nodes
                    ON liteboxfs_merkle_nodes.file = ancestors.file
                    AND liteboxfs_merkle_nodes.level = ancestors.level + 1
                    AND liteboxfs_merkle_nodes."index" = ancestors."index" / ?4
            )
            UPDATE
                liteboxfs_merkle_nodes
            SET
                dirty = 1
            WHERE
                (file, level, "index") IN (
                    SELECT file, level, "index" FROM ancestors
                );
            "#,
        )?;

        stmt.execute(rusqlite::params![
            pos.file,
            pos.level,
            pos.index,
            self.settings.merkle_branch_factor
        ])?;

        Ok(())
    }

    fn delete_node(&mut self, file: FileId, index: usize) -> crate::Result<()> {
        let mut stmt = self.db.prepare_cached(
            r#"
            WITH RECURSIVE deleted (level, "index") AS (
                SELECT
                    0 AS level,
                    ?2 AS "index"

                UNION ALL

                SELECT
                    deleted.level + 1 AS level,
                    deleted."index" / ?3 AS "index"
                FROM
                    deleted
                WHERE
                    NOT EXISTS (
                        SELECT
                            1
                        FROM
                            liteboxfs_merkle_nodes
                        WHERE
                            liteboxfs_merkle_nodes.file = ?1
                            AND liteboxfs_merkle_nodes.level = deleted.level
                            AND liteboxfs_merkle_nodes."index" / ?3 = deleted."index" / ?3
                            AND liteboxfs_merkle_nodes."index" != deleted."index"
                    )
            )
            DELETE FROM
                liteboxfs_merkle_nodes
            WHERE
                file = ?1
                AND (level, "index") IN (
                    SELECT
                        level,
                        "index"
                    FROM
                        deleted
                );
            "#,
        )?;

        stmt.execute(rusqlite::params![
            file,
            index,
            self.settings.merkle_branch_factor
        ])?;

        Ok(())
    }

    fn mark_dirty<R>(&mut self, file: FileId, indices: R) -> crate::Result<()>
    where
        R: RangeBounds<usize>,
    {
        let start_index = match indices.start_bound() {
            Bound::Included(&idx) => idx,
            Bound::Excluded(&idx) => idx + 1,
            Bound::Unbounded => 0,
        };

        // SQL uses `< end_index`, so we need to convert to an exclusive upper bound
        let end_index = match indices.end_bound() {
            Bound::Included(&idx) => Some(idx + 1),
            Bound::Excluded(&idx) => Some(idx),
            Bound::Unbounded => None,
        };

        // Only mark existing nodes and their ancestors as dirty. For new liteboxfs_files with no merkle
        // nodes, this is a no-op.
        let mut stmt = self.db.prepare_cached(
            r#"
            WITH RECURSIVE ancestors AS (
                SELECT
                    file,
                    level,
                    "index"
                FROM
                    liteboxfs_merkle_nodes
                WHERE
                    file = ?1
                    AND level = ?2
                    AND "index" >= ?3
                    AND "index" < COALESCE(
                        ?4,
                        (
                            SELECT
                                MAX("index") + 1
                            FROM
                                liteboxfs_merkle_nodes
                            WHERE
                                file = ?1
                                AND level = ?2
                        )
                    )

                UNION ALL

                SELECT
                    liteboxfs_merkle_nodes.file,
                    liteboxfs_merkle_nodes.level,
                    liteboxfs_merkle_nodes."index"
                FROM
                    ancestors
                JOIN
                    liteboxfs_merkle_nodes
                    ON liteboxfs_merkle_nodes.file = ancestors.file
                    AND liteboxfs_merkle_nodes.level = ancestors.level + 1
                    AND liteboxfs_merkle_nodes."index" = ancestors."index" / ?5
            )
            UPDATE
                liteboxfs_merkle_nodes
            SET
                dirty = 1
            WHERE
                (file, level, "index") IN (
                    SELECT file, level, "index" FROM ancestors
                );
            "#,
        )?;

        stmt.execute(rusqlite::params![
            file,
            0,
            start_index,
            end_index,
            self.settings.merkle_branch_factor,
        ])?;

        Ok(())
    }
}

impl MerkleNodeStore for &mut SqlStore<'_> {
    fn get_root(&mut self, file: FileId) -> crate::Result<Option<MerkleNode>> {
        (**self).get_root(file)
    }

    fn node_indices(
        &mut self,
        file: FileId,
        level: u32,
        filter: DirtyFilter,
    ) -> crate::Result<Vec<usize>> {
        (**self).node_indices(file, level, filter)
    }

    fn list_nodes<R>(
        &mut self,
        file: FileId,
        level: u32,
        indices: R,
    ) -> crate::Result<Vec<MerkleNode>>
    where
        R: RangeBounds<usize>,
    {
        (**self).list_nodes(file, level, indices)
    }

    fn put_node(&mut self, pos: MerkleNodePosition, hash: MerkleHash) -> crate::Result<()> {
        (**self).put_node(pos, hash)
    }

    fn delete_node(&mut self, file: FileId, index: usize) -> crate::Result<()> {
        (**self).delete_node(file, index)
    }

    fn mark_dirty<R>(&mut self, file: FileId, indices: R) -> crate::Result<()>
    where
        R: RangeBounds<usize>,
    {
        (**self).mark_dirty(file, indices)
    }
}