liteboxfs 0.2.0

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

use crate::{
    block::{
        Block, BlockEncoding, BlockId, BlockIndex, BlockList, BlockSignature, BlockStore,
        BlockStoreInput, FileId as StoreFileId, FileOffset, HoleLen, ReadBlock,
    },
    errors::InternalError,
};

use super::{exclusive::ExclusiveFileId, store::SqlStore};

impl<'conn> SqlStore<'conn> {
    pub fn write_file_block(
        &mut self,
        exclusive: ExclusiveFileId,
        offset: FileOffset,
        block_list: BlockList,
        input: BlockStoreInput,
    ) -> crate::Result<BlockList> {
        self.data_store()
            .write_block(exclusive.file_id(), offset, block_list, input)
    }

    pub fn truncate_file(
        &mut self,
        exclusive: ExclusiveFileId,
        block_list: BlockList,
        len: FileOffset,
    ) -> crate::Result<BlockList> {
        self.data_store()
            .truncate(exclusive.file_id(), block_list, len)
    }
}

impl<'conn> SqlStore<'conn> {
    pub fn commit(self) -> crate::Result<()> {
        Ok(self.db.commit()?)
    }

    fn insert_block(&mut self, input: BlockStoreInput) -> crate::Result<BlockId> {
        let mut stmt = self.db.prepare_cached(
            r#"
            INSERT OR IGNORE INTO
                liteboxfs_blocks (length, encoding, hash, data)
            VALUES
                (?, ?, ?, ?);
            "#,
        )?;

        let encoding = input.encoding();
        let logical_len = input.len();

        let params = match &input {
            BlockStoreInput::Data { bytes, hash, .. } => {
                self.block_buf.clear();
                self.block_buf.extend_from_slice(bytes);

                rusqlite::params![
                    logical_len,
                    encoding,
                    hash.as_ref().to_owned(),
                    self.block_buf.as_slice()
                ]
            }
            BlockStoreInput::Hole { len } => {
                rusqlite::params![
                    *len as usize,
                    encoding,
                    rusqlite::types::Null,
                    rusqlite::types::Null
                ]
            }
        };

        let rows_inserted = stmt.execute(params)?;

        if rows_inserted > 0 {
            return Ok(BlockId::from(self.db.last_insert_rowid()));
        }

        // Block already exists.
        match input {
            BlockStoreInput::Data { hash, .. } => {
                let mut stmt = self.db.prepare_cached(
                    r#"
                    SELECT
                        id
                    FROM
                        liteboxfs_blocks
                    WHERE
                        hash = ?;
                    "#,
                )?;

                Ok(
                    stmt.query_one(rusqlite::params![hash.as_ref().to_owned()], |row| {
                        row.get(0)
                    })?,
                )
            }
            BlockStoreInput::Hole { len } => {
                let mut stmt = self.db.prepare_cached(
                    r#"
                    SELECT
                        id
                    FROM
                        liteboxfs_blocks
                    WHERE
                        hash IS NULL
                        AND length = ?;
                    "#,
                )?;

                Ok(stmt.query_one(rusqlite::params![len as usize], |row| row.get(0))?)
            }
        }
    }
}

impl<'conn> BlockStore for SqlStore<'conn> {
    fn list_blocks(&self, file: StoreFileId) -> crate::Result<BlockList> {
        let mut stmt = self.db.prepare_cached(
            r#"
            SELECT
                liteboxfs_blocks.id,
                liteboxfs_blocks.length,
                liteboxfs_blocks.hash
            FROM
                liteboxfs_file_blocks
            JOIN
                liteboxfs_blocks ON liteboxfs_file_blocks.block = liteboxfs_blocks.id
            WHERE
                liteboxfs_file_blocks.file = ?
            ORDER BY
                liteboxfs_file_blocks."index" ASC;
            "#,
        )?;

        let blocks = stmt.query_map(rusqlite::params![file], |row| {
            let block_id: i64 = row.get(0)?;
            let block_len: usize = row.get(1)?;
            let block_hash: Option<Vec<u8>> = row.get(2)?;

            Ok(match block_hash {
                Some(hash) => Ok(Block {
                    id: block_id.into(),
                    signature: BlockSignature::Data {
                        hash: match hash.as_slice().try_into() {
                            Ok(hash) => hash,
                            Err(err) => return Ok(Err(err)),
                        },
                        len: block_len,
                    },
                }),
                None => Ok(Block {
                    id: block_id.into(),
                    signature: BlockSignature::Hole {
                        len: block_len as u64,
                    },
                }),
            })
        })?;

        let blocks: Vec<_> = blocks.collect::<Result<Result<Vec<_>, _>, _>>()??;
        Ok(blocks.into())
    }

    fn read_block(&mut self, block: BlockId, buf: &mut Vec<u8>) -> crate::Result<ReadBlock> {
        let mut stmt = self.db.prepare_cached(
            r#"
            SELECT
                length,
                encoding,
                hash,
                data
            FROM
                liteboxfs_blocks
            WHERE
                id = ?;
            "#,
        )?;

        let mut rows = stmt.query(rusqlite::params![block])?;

        if let Some(row) = rows.next()? {
            let block_len: usize = row.get(0)?;
            let encoding: BlockEncoding = row.get(1)?;

            #[cfg(not(feature = "compression"))]
            if encoding != BlockEncoding::NONE {
                return Err(InternalError::CompressionDisabled.into());
            }

            let block_hash: Option<Vec<u8>> = row.get(2)?;
            let block_data: Option<Vec<u8>> = row.get(3)?;

            match (block_hash, block_data) {
                (Some(hash), Some(data)) => {
                    buf.extend_from_slice(&data);

                    Ok(ReadBlock {
                        signature: BlockSignature::Data {
                            hash: hash.as_slice().try_into()?,
                            len: block_len,
                        },
                        encoding,
                    })
                }
                (None, None) => Ok(ReadBlock {
                    signature: BlockSignature::Hole {
                        len: block_len as HoleLen,
                    },
                    encoding,
                }),
                _ => Err(InternalError::MismatchedBlockHashAndData.into()),
            }
        } else {
            Err(InternalError::BlockNotFound { id: block }.into())
        }
    }

    fn replace_block(
        &mut self,
        file: StoreFileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        let block_id = self.insert_block(input)?;

        {
            let mut stmt = self.db.prepare_cached(
                r#"
                INSERT OR REPLACE INTO
                    liteboxfs_file_blocks (file, "index", block)
                VALUES
                    (?, ?, ?);
                "#,
            )?;

            stmt.execute(rusqlite::params![file, index, block_id])?;
        }

        Ok(block_id)
    }

    fn insert_block(
        &mut self,
        file: StoreFileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        let block_id = self.insert_block(input)?;

        // We need to shift each block index forward by one. However, SQLite doesn't allow you to
        // control which order rows are updated in, which means the naive approach would result in
        // a UNIQUE constraint violation when one row index is incremented before the one that
        // comes after it.
        //
        // To work around this, we shift in two phases. The shift offset has to be chosen so that
        // *both* phases' source and target ranges are disjoint from each other and from any
        // unshifted rows. Concretely, we shift up by `max_index + 2`, then back down by
        // `max_index + 1` for a net shift of +1. The `+2` ensures the phase-2 source range
        // `[index + max_index + 2, 2 * max_index + 2]` lies strictly above its target range
        // `[index + 1, max_index + 1]` — important for the `index == 0` case, where a smaller
        // offset would leave the largest shifted row's target equal to the smallest unshifted
        // row's source (e.g. shift `{0..=5}` to `{6..=11}` and back to `{1..=6}` — row 11→6
        // collides with row 6 if 11 moves first).
        let max_index: usize = {
            let mut stmt = self.db.prepare_cached(
                r#"
                SELECT
                    COALESCE(MAX("index"), 0)
                FROM
                    liteboxfs_file_blocks
                WHERE
                    file = ?;
                "#,
            )?;

            stmt.query_one(rusqlite::params![file], |row| row.get(0))?
        };

        {
            let mut stmt = self.db.prepare_cached(
                r#"
                UPDATE
                    liteboxfs_file_blocks
                SET
                    "index" = "index" + ? + 2
                WHERE
                    file = ?
                AND
                    "index" >= ?;
                "#,
            )?;

            stmt.execute(rusqlite::params![max_index, file, index])?;
        }

        {
            let mut stmt = self.db.prepare_cached(
                r#"
                UPDATE
                    liteboxfs_file_blocks
                SET
                    "index" = "index" - ? - 1
                WHERE
                    file = ?
                AND
                    "index" >= ?;
                "#,
            )?;

            stmt.execute(rusqlite::params![max_index, file, index + max_index + 2])?;
        }

        {
            let mut stmt = self.db.prepare_cached(
                r#"
                INSERT OR ABORT INTO
                    liteboxfs_file_blocks (file, "index", block)
                VALUES
                    (?, ?, ?);
                "#,
            )?;

            stmt.execute(rusqlite::params![file, index, block_id])?;
        }

        Ok(block_id)
    }

    // This optimizes for the append case by saving us from needing to shift any indices.
    fn append_block(
        &mut self,
        file: StoreFileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        let block_id = self.insert_block(input)?;

        {
            let mut stmt = self.db.prepare_cached(
                r#"
                INSERT INTO
                    liteboxfs_file_blocks (file, "index", block)
                VALUES
                    (?, ?, ?);
                "#,
            )?;

            stmt.execute(rusqlite::params![file, index, block_id])?;
        }

        Ok(block_id)
    }

    fn remove_blocks<R: RangeBounds<usize>>(
        &mut self,
        file: StoreFileId,
        blocks: R,
    ) -> crate::Result<Vec<BlockId>> {
        let start = match blocks.start_bound() {
            Bound::Included(index) => *index,
            Bound::Excluded(index) => *index + 1,
            Bound::Unbounded => 0,
        };

        let end = match blocks.end_bound() {
            Bound::Included(index) => *index + 1,
            Bound::Excluded(index) => *index,
            Bound::Unbounded => usize::MAX,
        };

        // The two-phase shift below relies on `max_index` being larger than every index in the
        // shift's WHERE clause, so it has to be sampled *before* the DELETE.
        let max_index: usize = {
            let mut stmt = self.db.prepare_cached(
                r#"
                SELECT
                    COALESCE(MAX("index"), 0)
                FROM
                    liteboxfs_file_blocks
                WHERE
                    file = ?;
                "#,
            )?;

            stmt.query_one(rusqlite::params![file], |row| row.get(0))?
        };

        let removed = {
            let mut stmt = self.db.prepare_cached(
                r#"
                DELETE FROM
                    liteboxfs_file_blocks
                WHERE
                    file = ?
                    AND "index" >= ?
                    AND "index" < ?
                RETURNING
                    block;
                "#,
            )?;

            let blocks = stmt.query_map(rusqlite::params![file, start, end], |row| {
                let block_id: i64 = row.get(0)?;
                Ok(BlockId::from(block_id))
            })?;

            blocks.collect::<Result<Vec<_>, _>>()?
        };

        // Shift subsequent blocks down to keep the per-file index sequence contiguous, matching
        // the in-memory block list's view of the file. The shift is split into two phases for
        // the same reason as `insert_block`: a single naive `UPDATE` could trigger a UNIQUE
        // constraint violation when a shifted value collides with another unshifted row in the
        // same statement.
        if !removed.is_empty() && end < usize::MAX {
            let shift = end - start;

            {
                let mut stmt = self.db.prepare_cached(
                    r#"
                    UPDATE
                        liteboxfs_file_blocks
                    SET
                        "index" = "index" + ?
                    WHERE
                        file = ?
                    AND
                        "index" >= ?;
                    "#,
                )?;

                stmt.execute(rusqlite::params![max_index, file, end])?;
            }

            {
                let mut stmt = self.db.prepare_cached(
                    r#"
                    UPDATE
                        liteboxfs_file_blocks
                    SET
                        "index" = "index" - ? - ?
                    WHERE
                        file = ?
                    AND
                        "index" >= ?;
                    "#,
                )?;

                stmt.execute(rusqlite::params![max_index, shift, file, max_index + end])?;
            }
        }

        Ok(removed)
    }
}

impl<'conn> BlockStore for &mut SqlStore<'conn> {
    fn list_blocks(&self, file: StoreFileId) -> crate::Result<BlockList> {
        <SqlStore<'conn> as BlockStore>::list_blocks(*self, file)
    }

    fn read_block(&mut self, block: BlockId, buf: &mut Vec<u8>) -> crate::Result<ReadBlock> {
        <SqlStore<'conn> as BlockStore>::read_block(*self, block, buf)
    }

    fn replace_block(
        &mut self,
        file: StoreFileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        <SqlStore<'conn> as BlockStore>::replace_block(*self, file, index, input)
    }

    fn insert_block(
        &mut self,
        file: StoreFileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        <SqlStore<'conn> as BlockStore>::insert_block(*self, file, index, input)
    }

    fn append_block(
        &mut self,
        file: StoreFileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        <SqlStore<'conn> as BlockStore>::append_block(*self, file, index, input)
    }

    fn remove_blocks<R: RangeBounds<usize>>(
        &mut self,
        file: StoreFileId,
        blocks: R,
    ) -> crate::Result<Vec<BlockId>> {
        <SqlStore<'conn> as BlockStore>::remove_blocks(*self, file, blocks)
    }
}