argus-chunk 0.1.0

//! Chunk payload codecs shared by Tetration and Tessera.
//!
//! Wire tags: **`0` = raw** (`stored_byte_len == raw_byte_len`), **`1` = zstd**
//! (frame that decompresses to `raw_byte_len` bytes).
//!
//! Decoding validates the logical length recorded by the caller. It does not
//! validate file offsets or stored spans; use [`crate::le`] before passing a
//! slice into [`decode`].

use std::borrow::Cow;

/// Payload storage codec represented as a `u32` on disk.
///
/// The numeric discriminants are part of the shared wire contract and must
/// not be renumbered.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u32)]
pub enum PayloadCodec {
    /// Stored bytes are the raw payload (`stored_byte_len == raw_byte_len`).
    Raw = 0,
    /// Stored bytes are one zstd frame decoding to `raw_byte_len` bytes.
    Zstd = 1,
}

impl PayloadCodec {
    /// Convert a wire discriminant into a supported codec.
    ///
    /// # Errors
    ///
    /// Returns [`CodecError::Unsupported`] for unknown tags.
    pub fn from_u32(value: u32) -> Result<Self, CodecError> {
        match value {
            0 => Ok(Self::Raw),
            1 => Ok(Self::Zstd),
            other => Err(CodecError::Unsupported { codec: other }),
        }
    }

    /// Return the stable `u32` discriminant stored on disk.
    #[must_use]
    pub const fn as_u32(self) -> u32 {
        self as u32
    }
}

/// Failures while selecting, encoding, or decoding a payload codec.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum CodecError {
    /// Unknown codec tag.
    #[error("unsupported payload codec {codec}")]
    Unsupported {
        /// Wire `u32` tag.
        codec: u32,
    },
    /// Raw codec requires equal raw and stored lengths.
    #[error("raw codec: raw_byte_len {raw} != stored_byte_len {stored}")]
    RawStoredMismatch {
        /// Declared uncompressed length.
        raw: u64,
        /// Declared on-disk length.
        stored: u64,
    },
    /// Decompressed length did not match the index.
    #[error("decoded {decoded} bytes, index says raw_byte_len {raw}")]
    DecodedLengthMismatch {
        /// Actual decompressed length.
        decoded: usize,
        /// Expected `raw_byte_len`.
        raw: u64,
    },
    /// zstd compress or decompress failed.
    #[error("zstd: {0}")]
    Zstd(String),
}

/// Encode uncompressed bytes using the selected wire codec.
///
/// Raw encoding copies `raw` unchanged. Zstd uses the library's default
/// compression level (`0`) so the format does not prescribe a tuning level.
///
/// # Errors
///
/// Returns [`CodecError::Zstd`] when compression fails.
pub fn encode(codec: PayloadCodec, raw: &[u8]) -> Result<Vec<u8>, CodecError> {
    match codec {
        PayloadCodec::Raw => Ok(raw.to_vec()),
        PayloadCodec::Zstd => zstd::encode_all(raw, 0).map_err(|e| CodecError::Zstd(e.to_string())),
    }
}

/// Decode stored chunk bytes to uncompressed payload (`raw_byte_len` bytes).
///
/// Raw → [`Cow::Borrowed`]; zstd → owned. For raw, `stored.len()` must equal
/// `raw_byte_len`.
///
/// The `Cow` return avoids allocating for memory-mapped raw chunks while
/// preserving one API for compressed and uncompressed payloads.
///
/// # Errors
///
/// Returns [`CodecError`] on length mismatch or zstd failure.
pub fn decode<'a>(
    codec: PayloadCodec,
    stored: &'a [u8],
    raw_byte_len: u64,
) -> Result<Cow<'a, [u8]>, CodecError> {
    match codec {
        PayloadCodec::Raw => {
            let stored_len = stored.len() as u64;
            if stored_len != raw_byte_len {
                return Err(CodecError::RawStoredMismatch {
                    raw: raw_byte_len,
                    stored: stored_len,
                });
            }
            // Keep raw mmap-backed payloads zero-copy.
            Ok(Cow::Borrowed(stored))
        }
        PayloadCodec::Zstd => {
            let dec = zstd::decode_all(stored).map_err(|e| CodecError::Zstd(e.to_string()))?;
            // A valid zstd frame is not enough: its logical size must agree
            // with the format-specific index row supplied by the caller.
            if dec.len() as u64 != raw_byte_len {
                return Err(CodecError::DecodedLengthMismatch {
                    decoded: dec.len(),
                    raw: raw_byte_len,
                });
            }
            Ok(Cow::Owned(dec))
        }
    }
}

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

    #[test]
    fn raw_round_trip_borrows() {
        let raw = b"hello chunk";
        let stored = encode(PayloadCodec::Raw, raw).unwrap();
        assert_eq!(stored.as_slice(), raw);
        let out = decode(PayloadCodec::Raw, &stored, raw.len() as u64).unwrap();
        assert!(matches!(out, Cow::Borrowed(_)));
        assert_eq!(&*out, raw);
    }

    #[test]
    fn zstd_round_trip() {
        let raw = b"compress me please........";
        let stored = encode(PayloadCodec::Zstd, raw).unwrap();
        assert_ne!(stored.as_slice(), raw);
        let out = decode(PayloadCodec::Zstd, &stored, raw.len() as u64).unwrap();
        assert_eq!(&*out, raw);
    }

    #[test]
    fn raw_rejects_length_mismatch() {
        let err = decode(PayloadCodec::Raw, b"ab", 3).unwrap_err();
        assert_eq!(err, CodecError::RawStoredMismatch { raw: 3, stored: 2 });
    }

    #[test]
    fn from_u32_rejects_unknown() {
        assert!(matches!(
            PayloadCodec::from_u32(99),
            Err(CodecError::Unsupported { codec: 99 })
        ));
    }
}