trine-kv 0.5.13

use crate::{
    codec::{self, CodecId},
    error::{Error, Result},
    storage::StorageReadBuffer,
};
use bytes::Bytes;
use std::ops::Range;

const BLOCK_HEADER_LEN: usize = 13;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct BlockHandle {
    pub(crate) offset: u64,
    pub(crate) len: u64,
}

#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct BlockManager;

#[derive(Debug, Clone)]
pub(crate) struct DecodedBlock {
    codec: CodecId,
    bytes: Bytes,
    payload_range: Range<usize>,
}

impl DecodedBlock {
    pub(crate) const fn codec(&self) -> CodecId {
        self.codec
    }

    pub(crate) fn payload(&self) -> &[u8] {
        &self.bytes[self.payload_range.clone()]
    }

    pub(crate) fn into_shared_payload(self) -> (Bytes, Range<usize>) {
        (self.bytes, self.payload_range)
    }
}

pub(crate) trait BlockReadSource {
    fn read_exact_at(&self, offset: usize, bytes: &mut [u8]) -> Result<()>;

    /// Reads `len` bytes at `offset` into an owned, `Arc`-backed completion that
    /// is decoupled from the borrowed-buffer path. Block decode reads through
    /// this seam so the read completion no longer borrows the decode call frame,
    /// which is the precondition for driving decode through the runtime's
    /// owned-read boundary in a later phase. The default fills a heap buffer via
    /// the borrowed path; native-file sources override it to use the storage
    /// object's owned blocking read.
    fn read_exact_at_owned(&self, offset: usize, len: usize) -> Result<StorageReadBuffer> {
        let mut bytes = vec![0_u8; len];
        self.read_exact_at(offset, &mut bytes)?;
        Ok(StorageReadBuffer::from_vec(offset, bytes))
    }
}

impl BlockManager {
    pub(crate) fn append_checked(
        bytes: &mut Vec<u8>,
        codec: CodecId,
        block_payload: &[u8],
    ) -> Result<BlockHandle> {
        let section_start = bytes.len();
        let block = Self::encode_checked(codec, block_payload)?;
        bytes.extend_from_slice(&block);

        Ok(BlockHandle {
            offset: usize_to_u64(section_start, "block offset")?,
            len: usize_to_u64(bytes.len() - section_start, "block length")?,
        })
    }

    pub(crate) fn encode_checked(codec: CodecId, block_payload: &[u8]) -> Result<Vec<u8>> {
        ensure_writable_block_len(block_payload.len())?;
        let encoded = codec::encode_block(codec, block_payload)?;
        let mut bytes = Vec::with_capacity(BLOCK_HEADER_LEN + encoded.len());
        bytes.push(codec.tag());
        put_u32(
            &mut bytes,
            usize_to_u32(block_payload.len(), "block payload length")?,
        );
        put_u32(
            &mut bytes,
            usize_to_u32(encoded.len(), "encoded block length")?,
        );
        put_u32(&mut bytes, checksum(&encoded));
        bytes.extend_from_slice(&encoded);
        Ok(bytes)
    }

    #[cfg(test)]
    pub(crate) fn read_checked(payload: &[u8], block: BlockHandle) -> Result<(CodecId, Vec<u8>)> {
        let (start, end) = block_bounds(block)?;
        let block_bytes = payload
            .get(start..end)
            .ok_or_else(|| invalid_table("block outside table payload"))?;
        Self::decode_checked(block_bytes)
    }

    #[cfg(test)]
    pub(crate) fn read_checked_at_payload_offset(
        payload: &[u8],
        offset: usize,
    ) -> Result<(BlockHandle, CodecId, Vec<u8>)> {
        let header_end = offset
            .checked_add(BLOCK_HEADER_LEN)
            .ok_or_else(|| invalid_table("block offset overflow"))?;
        let header = payload
            .get(offset..header_end)
            .ok_or_else(|| invalid_table("short block header"))?;
        let len = checked_block_len(header)?;
        let end = offset
            .checked_add(len)
            .ok_or_else(|| invalid_table("block offset overflow"))?;
        let block = BlockHandle {
            offset: usize_to_u64(offset, "block offset")?,
            len: usize_to_u64(len, "block length")?,
        };
        let block_bytes = payload
            .get(offset..end)
            .ok_or_else(|| invalid_table("block outside table payload"))?;
        let (codec, decoded) = Self::decode_checked(block_bytes)?;
        Ok((block, codec, decoded))
    }

    #[cfg(test)]
    pub(crate) fn read_checked_from_source(
        payload_len: usize,
        payload_base_offset: usize,
        block: BlockHandle,
        source: &impl BlockReadSource,
    ) -> Result<(CodecId, Vec<u8>)> {
        let (start, end) = block_bounds(block)?;
        if end > payload_len {
            return Err(invalid_table("block outside table payload"));
        }

        let source_offset = payload_base_offset
            .checked_add(start)
            .ok_or_else(|| invalid_table("block file offset overflow"))?;
        let block_bytes = source.read_exact_at_owned(source_offset, end - start)?;
        Self::decode_checked_owned(block_bytes).map(|block| {
            let codec = block.codec();
            let payload = block.payload().to_vec();
            (codec, payload)
        })
    }

    pub(crate) fn read_checked_from_source_shared(
        payload_len: usize,
        payload_base_offset: usize,
        block: BlockHandle,
        source: &impl BlockReadSource,
    ) -> Result<DecodedBlock> {
        let (start, end) = block_bounds(block)?;
        if end > payload_len {
            return Err(invalid_table("block outside table payload"));
        }

        let source_offset = payload_base_offset
            .checked_add(start)
            .ok_or_else(|| invalid_table("block file offset overflow"))?;
        let block_bytes = source.read_exact_at_owned(source_offset, end - start)?;
        Self::decode_checked_owned(block_bytes)
    }

    #[cfg(test)]
    pub(crate) fn read_checked_at_source_offset(
        payload_len: usize,
        payload_base_offset: usize,
        offset: usize,
        source: &impl BlockReadSource,
    ) -> Result<(BlockHandle, CodecId, Vec<u8>)> {
        if offset >= payload_len {
            return Err(invalid_table("block outside table payload"));
        }
        let source_offset = payload_base_offset
            .checked_add(offset)
            .ok_or_else(|| invalid_table("block file offset overflow"))?;
        let header = source.read_exact_at_owned(source_offset, BLOCK_HEADER_LEN)?;
        let len = checked_block_len(header.as_slice())?;
        let end = offset
            .checked_add(len)
            .ok_or_else(|| invalid_table("block offset overflow"))?;
        if end > payload_len {
            return Err(invalid_table("block outside table payload"));
        }
        let block_bytes = source.read_exact_at_owned(source_offset, len)?;
        let decoded_block = Self::decode_checked_owned(block_bytes)?;
        let codec = decoded_block.codec();
        let decoded = decoded_block.payload().to_vec();
        Ok((
            BlockHandle {
                offset: usize_to_u64(offset, "block offset")?,
                len: usize_to_u64(len, "block length")?,
            },
            codec,
            decoded,
        ))
    }

    pub(crate) fn read_checked_at_source_offset_shared(
        payload_len: usize,
        payload_base_offset: usize,
        offset: usize,
        source: &impl BlockReadSource,
    ) -> Result<(BlockHandle, DecodedBlock)> {
        if offset > payload_len {
            return Err(invalid_table("block offset outside table payload"));
        }
        let source_offset = payload_base_offset
            .checked_add(offset)
            .ok_or_else(|| invalid_table("block file offset overflow"))?;
        let header = source.read_exact_at_owned(source_offset, BLOCK_HEADER_LEN)?;
        let len = checked_block_len(header.as_slice())?;
        let end = offset
            .checked_add(len)
            .ok_or_else(|| invalid_table("block length overflow"))?;
        if end > payload_len {
            return Err(invalid_table("block outside table payload"));
        }
        let block_bytes = source.read_exact_at_owned(source_offset, len)?;
        let decoded_block = Self::decode_checked_owned(block_bytes)?;
        Ok((
            BlockHandle {
                offset: usize_to_u64(offset, "block offset")?,
                len: usize_to_u64(len, "block length")?,
            },
            decoded_block,
        ))
    }

    #[cfg(test)]
    pub(crate) fn decode_checked(block_bytes: &[u8]) -> Result<(CodecId, Vec<u8>)> {
        if block_bytes.len() < BLOCK_HEADER_LEN {
            return Err(invalid_table("short block header"));
        }

        let codec = CodecId::from_tag(block_bytes[0])?;
        let uncompressed_len = read_u32_at(block_bytes, 1)? as usize;
        codec::ensure_decoded_block_len(uncompressed_len)?;
        let encoded_len = read_u32_at(block_bytes, 5)? as usize;
        let expected_checksum = read_u32_at(block_bytes, 9)?;
        if block_bytes.len() != BLOCK_HEADER_LEN + encoded_len {
            return Err(Error::Corruption {
                message: "block length mismatch".to_owned(),
            });
        }

        let encoded = &block_bytes[BLOCK_HEADER_LEN..];
        if checksum(encoded) != expected_checksum {
            return Err(Error::Corruption {
                message: "block checksum mismatch".to_owned(),
            });
        }

        Ok((
            codec,
            codec::decode_block(codec, encoded, uncompressed_len)?,
        ))
    }

    pub(crate) fn decode_checked_owned(block_bytes: StorageReadBuffer) -> Result<DecodedBlock> {
        let bytes = block_bytes.as_slice();
        let block = decoded_block_header(bytes)?;
        let encoded = &bytes[BLOCK_HEADER_LEN..];
        if checksum(encoded) != block.expected_checksum {
            return Err(Error::Corruption {
                message: "block checksum mismatch".to_owned(),
            });
        }

        if block.codec == CodecId::None {
            if block.encoded_len != block.uncompressed_len {
                return Err(Error::InvalidFormat {
                    message: "uncompressed block length mismatch".to_owned(),
                });
            }
            return Ok(DecodedBlock {
                codec: block.codec,
                bytes: block_bytes.into_bytes(),
                payload_range: BLOCK_HEADER_LEN..BLOCK_HEADER_LEN + block.encoded_len,
            });
        }

        let decoded = codec::decode_block(block.codec, encoded, block.uncompressed_len)?;
        let payload_len = decoded.len();
        Ok(DecodedBlock {
            codec: block.codec,
            bytes: Bytes::from(decoded),
            payload_range: 0..payload_len,
        })
    }
}

#[derive(Debug, Clone, Copy)]
struct DecodedBlockHeader {
    codec: CodecId,
    uncompressed_len: usize,
    encoded_len: usize,
    expected_checksum: u32,
}

pub(crate) fn block_bounds(handle: BlockHandle) -> Result<(usize, usize)> {
    bounds(handle.offset, handle.len)
}

pub(crate) fn bounds(offset: u64, len: u64) -> Result<(usize, usize)> {
    let start = usize::try_from(offset).map_err(|_| invalid_table("offset exceeds usize"))?;
    let len = usize::try_from(len).map_err(|_| invalid_table("length exceeds usize"))?;
    let end = start
        .checked_add(len)
        .ok_or_else(|| invalid_table("offset plus length overflows usize"))?;
    Ok((start, end))
}

pub(crate) fn checksum(bytes: &[u8]) -> u32 {
    crate::checksum::crc32c(bytes)
}

fn checked_block_len(header: &[u8]) -> Result<usize> {
    let uncompressed_len = read_u32_at(header, 1)? as usize;
    codec::ensure_decoded_block_len(uncompressed_len)?;
    let encoded_len = read_u32_at(header, 5)? as usize;
    BLOCK_HEADER_LEN
        .checked_add(encoded_len)
        .ok_or_else(|| invalid_table("block length overflow"))
}

fn decoded_block_header(block_bytes: &[u8]) -> Result<DecodedBlockHeader> {
    if block_bytes.len() < BLOCK_HEADER_LEN {
        return Err(invalid_table("short block header"));
    }

    let codec = CodecId::from_tag(block_bytes[0])?;
    let uncompressed_len = read_u32_at(block_bytes, 1)? as usize;
    codec::ensure_decoded_block_len(uncompressed_len)?;
    let encoded_len = read_u32_at(block_bytes, 5)? as usize;
    let expected_checksum = read_u32_at(block_bytes, 9)?;
    if block_bytes.len() != BLOCK_HEADER_LEN + encoded_len {
        return Err(Error::Corruption {
            message: "block length mismatch".to_owned(),
        });
    }

    Ok(DecodedBlockHeader {
        codec,
        uncompressed_len,
        encoded_len,
        expected_checksum,
    })
}

fn put_u32(bytes: &mut Vec<u8>, value: u32) {
    bytes.extend_from_slice(&value.to_le_bytes());
}

fn read_u32_at(bytes: &[u8], offset: usize) -> Result<u32> {
    let value = bytes
        .get(offset..offset + 4)
        .ok_or_else(|| invalid_table("short u32"))?;
    Ok(u32::from_le_bytes([value[0], value[1], value[2], value[3]]))
}

fn usize_to_u32(value: usize, field: &'static str) -> Result<u32> {
    u32::try_from(value).map_err(|_| Error::invalid_options(format!("{field} exceeds u32::MAX")))
}

fn usize_to_u64(value: usize, field: &'static str) -> Result<u64> {
    u64::try_from(value).map_err(|_| Error::invalid_options(format!("{field} exceeds u64::MAX")))
}

fn ensure_writable_block_len(len: usize) -> Result<()> {
    if len <= codec::MAX_DECODED_BLOCK_BYTES {
        return Ok(());
    }

    Err(Error::invalid_options(format!(
        "block payload length {len} exceeds maximum {}",
        codec::MAX_DECODED_BLOCK_BYTES
    )))
}

fn invalid_table(message: &'static str) -> Error {
    Error::InvalidFormat {
        message: format!("invalid table: {message}"),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::cell::Cell;

    /// Slice-backed source that exercises the default (borrowed-fallback) owned
    /// read seam: it only implements the borrowed read, so `read_exact_at_owned`
    /// runs the trait default that copies into an owned buffer.
    struct SliceSource<'a> {
        bytes: &'a [u8],
    }

    impl BlockReadSource for SliceSource<'_> {
        fn read_exact_at(&self, offset: usize, bytes: &mut [u8]) -> Result<()> {
            let end = offset
                .checked_add(bytes.len())
                .ok_or_else(|| invalid_table("read offset overflow"))?;
            let slice = self
                .bytes
                .get(offset..end)
                .ok_or_else(|| invalid_table("read past end"))?;
            bytes.copy_from_slice(slice);
            Ok(())
        }
    }

    struct RecordingOwnedSource<'a> {
        bytes: &'a [u8],
        full_payload_ptr: Cell<*const u8>,
    }

    impl BlockReadSource for RecordingOwnedSource<'_> {
        fn read_exact_at(&self, offset: usize, bytes: &mut [u8]) -> Result<()> {
            let end = offset
                .checked_add(bytes.len())
                .ok_or_else(|| invalid_table("read offset overflow"))?;
            let slice = self
                .bytes
                .get(offset..end)
                .ok_or_else(|| invalid_table("read past end"))?;
            bytes.copy_from_slice(slice);
            Ok(())
        }

        fn read_exact_at_owned(&self, offset: usize, len: usize) -> Result<StorageReadBuffer> {
            let end = offset
                .checked_add(len)
                .ok_or_else(|| invalid_table("read offset overflow"))?;
            let bytes = self
                .bytes
                .get(offset..end)
                .ok_or_else(|| invalid_table("read past end"))?
                .to_vec();
            if len > BLOCK_HEADER_LEN {
                self.full_payload_ptr
                    .set(bytes.as_ptr().wrapping_add(BLOCK_HEADER_LEN));
            }
            Ok(StorageReadBuffer::from_vec(offset, bytes))
        }
    }

    fn block_bytes_with_header(codec: CodecId, uncompressed_len: u32, encoded: &[u8]) -> Vec<u8> {
        let mut bytes = Vec::new();
        bytes.push(codec.tag());
        put_u32(&mut bytes, uncompressed_len);
        put_u32(
            &mut bytes,
            u32::try_from(encoded.len()).expect("test encoded length fits"),
        );
        put_u32(&mut bytes, checksum(encoded));
        bytes.extend_from_slice(encoded);
        bytes
    }

    fn assert_oversized_decoded_block(error: &Error) {
        assert!(
            matches!(error, Error::InvalidFormat { .. }),
            "unexpected error kind: {error}"
        );
        assert!(
            error.to_string().contains("decoded block length"),
            "unexpected error: {error}"
        );
    }

    #[test]
    fn owned_default_fallback_matches_borrowed_read() {
        let payload = b"trine block decode owned seam".to_vec();
        let source = SliceSource { bytes: &payload };

        let owned = source.read_exact_at_owned(4, 5).expect("owned read");
        assert_eq!(owned.offset(), 4);
        assert_eq!(owned.as_slice(), &payload[4..9]);
    }

    #[test]
    fn read_checked_from_source_decodes_through_owned_seam() {
        let mut payload = Vec::new();
        let block_payload = b"checked block body";
        let handle = BlockManager::append_checked(&mut payload, CodecId::None, block_payload)
            .expect("append block");
        let source = SliceSource { bytes: &payload };

        let (codec, decoded) =
            BlockManager::read_checked_from_source(payload.len(), 0, handle, &source)
                .expect("read checked block");
        assert_eq!(codec, CodecId::None);
        assert_eq!(decoded, block_payload);
    }

    #[test]
    fn none_codec_owned_decode_reuses_payload_bytes() {
        let block_payload = b"owned none payload";
        let block =
            BlockManager::encode_checked(CodecId::None, block_payload).expect("block encodes");
        let buffer = StorageReadBuffer::from_vec(0, block);
        let expected_payload_ptr = buffer.as_slice().as_ptr().wrapping_add(BLOCK_HEADER_LEN);

        let decoded = BlockManager::decode_checked_owned(buffer).expect("block decodes");

        assert_eq!(decoded.codec(), CodecId::None);
        assert_eq!(decoded.payload(), block_payload);
        assert_eq!(decoded.payload().as_ptr(), expected_payload_ptr);
    }

    #[test]
    fn borrowed_decode_rejects_oversized_uncompressed_len_before_lz4_allocation() {
        let block = block_bytes_with_header(CodecId::FastLz4Block, u32::MAX, &[0]);

        let error = BlockManager::decode_checked(&block)
            .expect_err("oversized decoded block should fail before lz4 decode");

        assert_oversized_decoded_block(&error);
    }

    #[test]
    fn owned_decode_rejects_oversized_uncompressed_len_before_lz4_allocation() {
        let block = block_bytes_with_header(CodecId::FastLz4Block, u32::MAX, &[0]);
        let buffer = StorageReadBuffer::from_vec(0, block);

        let error = BlockManager::decode_checked_owned(buffer)
            .expect_err("oversized decoded block should fail before lz4 decode");

        assert_oversized_decoded_block(&error);
    }

    #[test]
    fn offset_decode_rejects_oversized_uncompressed_len_before_payload_read() {
        let block = block_bytes_with_header(CodecId::FastLz4Block, u32::MAX, &[0]);
        let source = SliceSource { bytes: &block };

        let error = BlockManager::read_checked_at_source_offset(block.len(), 0, 0, &source)
            .expect_err("oversized decoded block should fail from header");

        assert_oversized_decoded_block(&error);
    }

    #[test]
    fn read_checked_at_source_offset_decodes_through_owned_seam() {
        let mut payload = Vec::new();
        let block_payload = b"offset addressed block";
        let handle = BlockManager::append_checked(&mut payload, CodecId::None, block_payload)
            .expect("append block");
        let source = SliceSource { bytes: &payload };

        let offset = usize::try_from(handle.offset).expect("offset fits usize");
        let (read_handle, codec, decoded) =
            BlockManager::read_checked_at_source_offset(payload.len(), 0, offset, &source)
                .expect("read checked block at offset");
        assert_eq!(read_handle, handle);
        assert_eq!(codec, CodecId::None);
        assert_eq!(decoded, block_payload);
    }

    #[test]
    fn shared_source_offset_read_reuses_owned_payload_bytes() {
        let mut payload = Vec::new();
        let block_payload = b"shared offset addressed block";
        let handle = BlockManager::append_checked(&mut payload, CodecId::None, block_payload)
            .expect("append block");
        let source = RecordingOwnedSource {
            bytes: &payload,
            full_payload_ptr: Cell::new(std::ptr::null()),
        };

        let offset = usize::try_from(handle.offset).expect("offset fits usize");
        let (read_handle, decoded) =
            BlockManager::read_checked_at_source_offset_shared(payload.len(), 0, offset, &source)
                .expect("read checked block at offset");

        assert_eq!(read_handle, handle);
        assert_eq!(decoded.codec(), CodecId::None);
        assert_eq!(decoded.payload(), block_payload);
        assert_eq!(decoded.payload().as_ptr(), source.full_payload_ptr.get());
    }
}