liteboxfs 0.2.0

use std::{
    fmt::{self, Display},
    iter,
};

use rusqlite::{ToSql, types::FromSql};

use crate::{
    block::{BlockData, BlockLen, HoleLen, LogicalBlockData},
    errors::InternalError,
};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct BlockHash(blake3::Hash);

impl TryFrom<&[u8]> for BlockHash {
    type Error = crate::Error;

    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        blake3::Hash::from_slice(value)
            .map(BlockHash)
            .map_err(|_| InternalError::IncorrectBlockHashLen.into())
    }
}

impl AsRef<[u8]> for BlockHash {
    fn as_ref(&self) -> &[u8] {
        self.0.as_bytes()
    }
}

#[derive(Debug)]
pub struct BlockHasher;

impl BlockHasher {
    pub fn hash(input: &[u8]) -> BlockHash {
        BlockHash(blake3::hash(input))
    }
}

/// The block hash of a leaf node in the merkle tree.
///
/// There are two kinds of blocks: "physical" blocks (usually just called blocks) and "logical"
/// blocks.
///
/// Physical blocks represent data as stored in the database; they are variable-size and are
/// subject to content-defined chunking. Logical blocks represent data from the perspective of the
/// merkle tree; they are fixed-size.
///
/// Two ensure two files with the same contents always have the same merkle hash, logical blocks
/// must be the same size for all files. Because physical blocks are variable-size, a single
/// logical block may span multiple physical blocks or vice versa.
///
/// Ultimately, having a merkle hash for each file allows us to cheaply determine if two files have
/// identical contents.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct LogicalBlockHash(blake3::Hash);

/// The block hash of a non-leaf node in the merkle tree.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MerkleHash(blake3::Hash);

#[doc(hidden)]
impl ToSql for MerkleHash {
    fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
        Ok(rusqlite::types::ToSqlOutput::from(
            self.0.as_bytes().as_slice(),
        ))
    }
}

#[doc(hidden)]
impl FromSql for MerkleHash {
    fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
        Ok(Self(blake3::Hash::from_slice(value.as_blob()?).map_err(
            |_| rusqlite::types::FromSqlError::InvalidBlobSize {
                expected_size: blake3::OUT_LEN,
                blob_size: value.as_blob().map(|b| b.len()).unwrap_or(0),
            },
        )?))
    }
}

#[cfg_attr(coverage_nightly, coverage(off))]
impl Display for MerkleHash {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0.to_hex())
    }
}

impl From<LogicalBlockHash> for MerkleHash {
    fn from(value: LogicalBlockHash) -> Self {
        MerkleHash(value.0)
    }
}

#[derive(Debug)]
pub struct LogicalBlockHasher {
    logical_block_size: BlockLen,
    // Because logical blocks are fixed-size, logical blocks that are fully contained within a hole
    // must all have the same hash. To avoid needing to recompute it every time, we cache it here.
    hole_block_hash_cache: Option<LogicalBlockHash>,
}

impl LogicalBlockHasher {
    pub fn new(logical_block_size: BlockLen) -> Self {
        Self {
            logical_block_size,
            hole_block_hash_cache: None,
        }
    }

    pub fn hash(&mut self, data: &LogicalBlockData) -> LogicalBlockHash {
        match data.as_ref() {
            crate::block::BlockData::Data { bytes } => LogicalBlockHash(blake3::hash(bytes)),
            // For the purpose of hashing, a hole is equivalent to a block of null bytes.
            crate::block::BlockData::Hole { len } => {
                if *len == self.logical_block_size as HoleLen {
                    return self.hole_block_hash_cache.clone().unwrap_or_else(|| {
                        LogicalBlockHash(blake3::hash(&vec![0u8; self.logical_block_size]))
                    });
                }

                LogicalBlockHash(blake3::hash(&vec![0u8; *len as usize]))
            }
        }
    }
}

#[derive(Debug)]
pub struct MerkleNodeHasher {
    buf: Vec<u8>,
}

impl MerkleNodeHasher {
    pub fn new(branch_factor: u32) -> Self {
        Self {
            buf: Vec::with_capacity((branch_factor as usize) * blake3::OUT_LEN),
        }
    }

    pub fn hash<T, I>(&mut self, hashes: I) -> MerkleHash
    where
        T: Into<MerkleHash>,
        I: IntoIterator<Item = T>,
    {
        self.buf.clear();

        for hash in hashes {
            self.buf.extend_from_slice(hash.into().0.as_bytes());
        }

        MerkleHash(blake3::hash(&self.buf))
    }
}

/// A cache for merkle node hashes for subtrees where all leaf nodes are holes.
///
/// This allows computing the root hash for files with large hole regions without iterating through
/// each logical block.
#[derive(Debug)]
pub struct HoleMerkleCache {
    branch_factor: u32,
    logical_block_size: usize,
    /// Cached hashes for all-hole subtrees. Index `i` contains the hash for level `i`.
    level_hashes: Vec<MerkleHash>,
}

impl HoleMerkleCache {
    pub fn new(branch_factor: u32, logical_block_size: usize) -> Self {
        Self {
            branch_factor,
            logical_block_size,
            level_hashes: Vec::new(),
        }
    }

    pub fn get_or_compute(
        &mut self,
        level: usize,
        merkle_hasher: &mut MerkleNodeHasher,
        block_hasher: &mut LogicalBlockHasher,
    ) -> MerkleHash {
        // Ensure we have computed up to the requested level.
        while self.level_hashes.len() <= level {
            let current_level = self.level_hashes.len();

            if current_level == 0 {
                // Compute the hash of the leaf node.
                let block_data = BlockData::Hole {
                    len: self.logical_block_size as HoleLen,
                };

                self.level_hashes
                    .push(block_hasher.hash(&block_data.into()).into());
            } else {
                // Compute the hash of an interior node.
                let child_hash = self
                    .level_hashes
                    .get(current_level - 1)
                    .expect("We should have already computed the child level's merkle hash.");

                let parent_hash = merkle_hasher.hash(iter::repeat_n(
                    child_hash.clone(),
                    self.branch_factor as usize,
                ));

                self.level_hashes.push(parent_hash);
            }
        }

        self.level_hashes
            .get(level)
            .expect("We should have already computed this hash.")
            .clone()
    }
}