swh-mosaic 0.3.1

// Copyright (C) 2026  The Software Heritage developers
// See the AUTHORS file at the top-level directory of this distribution
// License: GNU General Public License version 3, or any later version
// See top-level LICENSE file for more information

//!
//! our EBML toolbox: read/write tags and VINTs.
//!

use crate::Size;
use anyhow::Result;
use crc32fast::hash;
use std::fmt;
use thiserror::Error;

/// how many bytes does a complete (tag + length + data) CRC32 element use ?
pub const CRC32_SIZE: Size = Size(6);

pub const DOCTYPE: &str = "mosaic";

pub const DOCTYPE_VERSION: u8 = 1;

pub const DOCTYPE_READ_VERSION: u8 = 1;

pub const EBML_MAX_ID_LENGTH: Size = Size(4);

pub const EBML_MAX_SIZE_LENGTH: Size = Size(8);

/// Enumeration of possible EBML tags for MOSAIC, including EBML header
#[derive(PartialEq, Debug, Clone, Copy)]
pub enum MosaicTag {
    // EBML Header as per RFC 8794
    Ebml,
    EbmlVersion,
    EbmlReadVersion,
    EbmlMaxIdLength,
    EbmlMaxSizeLength,
    DocType,
    DocTypeVersion,
    DocTypeReadVersion,
    DocTypeExtension,
    DocTypeExtensionName,
    DocTypeExtensionVersion,
    Crc32,
    Void,

    // MOSAIC elements
    Mosaic,
    ContainerMetaData,
    ObjectsCounter,
    ObjectsTotalSize,
    CompressionMethod,
    CompressionData,
    Comment,
    EndOfTilesOffset,
    Tile,
    Object,
    Index,
    IdxDescription,
    IdxUnrolled,
    Key,
    Offset,
    GoToConflicts,
    IdxUnrolledEntrySize,
    MapContainer,
    Map,
    Conflicts,
    Conflict,
    ConflictingKey,
    ConflictingOffset,
}

impl fmt::Display for MosaicTag {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl MosaicTag {
    /// Try to parse an EBML tag into a MOSAIC tag from the beginning of given slice.
    /// Returns the tag and its length (that also indicates the offset of the element's size).
    ///
    /// Note that en error is returned only on I/O error or if data is not a valid EBML
    /// tag; unknown tags will be parsed as EBML Void so the reader will skip them.
    pub fn parse(raw: &[u8]) -> Result<(MosaicTag, Size)> {
        let (&first, _) = raw
            .split_first()
            .ok_or_else(|| anyhow::anyhow!("empty input"))?;
        let width = match first {
            0x10..=0x1F => 4,
            0x40..=0x7F => 2,
            0x80..=0xFF => 1,
            _ => anyhow::bail!("invalid EBML element ID leading byte: 0x{:02x}", first),
        };
        anyhow::ensure!(raw.len() >= width, "not enough bytes for EBML element ID");
        let id = raw[..width]
            .iter()
            .fold(0u32, |acc, &b| (acc << 8) | u32::from(b));
        let tag = match id {
            // EBML Header as per RFC 8794
            0x1a45dfa3 => MosaicTag::Ebml,
            0x4286 => MosaicTag::EbmlVersion,
            0x42f7 => MosaicTag::EbmlReadVersion,
            0x42f2 => MosaicTag::EbmlMaxIdLength,
            0x42f3 => MosaicTag::EbmlMaxSizeLength,
            0x4282 => MosaicTag::DocType,
            0x4287 => MosaicTag::DocTypeVersion,
            0x4285 => MosaicTag::DocTypeReadVersion,
            0x4281 => MosaicTag::DocTypeExtension,
            0x4283 => MosaicTag::DocTypeExtensionName,
            0x4284 => MosaicTag::DocTypeExtensionVersion,
            0xBF => MosaicTag::Crc32,
            0xEC => MosaicTag::Void,
            // MOSAIC elements
            0x1C535748 => MosaicTag::Mosaic,
            0x1D535748 => MosaicTag::ContainerMetaData,
            0x5000 => MosaicTag::ObjectsCounter,
            0x5001 => MosaicTag::ObjectsTotalSize,
            0x5002 => MosaicTag::CompressionMethod,
            0x5003 => MosaicTag::CompressionData,
            0x5004 => MosaicTag::Comment,
            0x5099 => MosaicTag::EndOfTilesOffset,
            0x1E535748 => MosaicTag::Tile,
            0xF0 => MosaicTag::Object,
            0x1F535748 => MosaicTag::Index,
            0x6001 => MosaicTag::IdxDescription,
            0x6002 => MosaicTag::IdxUnrolled,
            0xA0 => MosaicTag::Key,
            0xA1 => MosaicTag::Offset,
            0xA2 => MosaicTag::GoToConflicts,
            0x6003 => MosaicTag::IdxUnrolledEntrySize,
            0x6004 => MosaicTag::MapContainer,
            0x6005 => MosaicTag::Map,
            0x6010 => MosaicTag::Conflicts,
            0x6011 => MosaicTag::Conflict,
            0x6012 => MosaicTag::ConflictingKey,
            0x6013 => MosaicTag::ConflictingOffset,
            _ => MosaicTag::Void,
        };
        Ok((tag, width.try_into()?))
    }

    pub fn to_be_bytes(&self) -> &'static [u8] {
        match self {
            // EBML Header as per RFC 8794
            MosaicTag::Ebml => &[0x1a, 0x45, 0xdf, 0xa3],
            MosaicTag::EbmlVersion => &[0x42, 0x86],
            MosaicTag::EbmlReadVersion => &[0x42, 0xf7],
            MosaicTag::EbmlMaxIdLength => &[0x42, 0xf2],
            MosaicTag::EbmlMaxSizeLength => &[0x42, 0xf3],
            MosaicTag::DocType => &[0x42, 0x82],
            MosaicTag::DocTypeVersion => &[0x42, 0x87],
            MosaicTag::DocTypeReadVersion => &[0x42, 0x85],
            MosaicTag::DocTypeExtension => &[0x42, 0x81],
            MosaicTag::DocTypeExtensionName => &[0x42, 0x83],
            MosaicTag::DocTypeExtensionVersion => &[0x42, 0x84],
            MosaicTag::Crc32 => &[0xBF],
            MosaicTag::Void => &[0xEC],
            // MOSAIC elements
            MosaicTag::Mosaic => &[0x1C, 0x53, 0x57, 0x48],
            MosaicTag::ContainerMetaData => &[0x1D, 0x53, 0x57, 0x48],
            MosaicTag::ObjectsCounter => &[0x50, 0x00],
            MosaicTag::ObjectsTotalSize => &[0x50, 0x01],
            MosaicTag::CompressionMethod => &[0x50, 0x02],
            MosaicTag::CompressionData => &[0x50, 0x03],
            MosaicTag::Comment => &[0x50, 0x04],
            MosaicTag::EndOfTilesOffset => &[0x50, 0x99],
            MosaicTag::Tile => &[0x1E, 0x53, 0x57, 0x48],
            MosaicTag::Object => &[0xF0],
            MosaicTag::Index => &[0x1F, 0x53, 0x57, 0x48],
            MosaicTag::IdxDescription => &[0x60, 0x01],
            MosaicTag::IdxUnrolled => &[0x60, 0x02],
            MosaicTag::Key => &[0xA0],
            MosaicTag::Offset => &[0xA1],
            MosaicTag::GoToConflicts => &[0xA2],
            MosaicTag::IdxUnrolledEntrySize => &[0x60, 0x03],
            MosaicTag::MapContainer => &[0x60, 0x04],
            MosaicTag::Map => &[0x60, 0x05],
            MosaicTag::Conflicts => &[0x60, 0x10],
            MosaicTag::Conflict => &[0x60, 0x11],
            MosaicTag::ConflictingKey => &[0x60, 0x12],
            MosaicTag::ConflictingOffset => &[0x60, 0x13],
        }
    }

    pub fn is_master(&self) -> bool {
        matches!(
            self,
            MosaicTag::Ebml
                | MosaicTag::Mosaic
                | MosaicTag::ContainerMetaData
                | MosaicTag::Tile
                | MosaicTag::Index
                | MosaicTag::IdxUnrolled
                | MosaicTag::MapContainer
                | MosaicTag::Conflicts
                | MosaicTag::Conflict
        )
    }
}

pub trait ShortestBeBytes: Into<u64> + Copy {
    /// Like to_be_bytes, but cutting the first zeros of the bytes array. Note that
    /// zero is represented as an empty array.
    fn shortest_be_bytes(self) -> Vec<u8> {
        let casted: u64 = self.into();
        let sliced = casted.to_be_bytes();
        let mut real_start = 0;
        while real_start < 8 && sliced[real_start] == 0 {
            real_start += 1;
        }

        sliced[real_start..].to_vec()
    }
}

impl ShortestBeBytes for u64 {}
impl ShortestBeBytes for u32 {}
impl ShortestBeBytes for u16 {}
impl ShortestBeBytes for u8 {}

/// Specialized error return types
#[derive(Error, Debug)]
pub enum VIntError {
    #[error("{0} (0x{0:X}) is too large to be represented as a VInt")]
    WriteOverflow(u64),

    #[error("{0} (0x{0:X}) is too large to be represented as a {1}-bytes VInt")]
    ConstrainedWriteOverflow(u64, u64),

    #[error("Required a {0}-bytes VInt but their maximal size is {1}")]
    VIntTooLarge(u64, u64),

    #[error("cannot read the VInt: got {0} bytes, expected {1}.")]
    ReadUnderflow(usize, usize),

    #[error("Invalid start byte: this is not a valid VInt")]
    InvalidVInt,
}

/// greatest value we can represent as a VInt, a long as we use EbmlMaxSizeLength=8
const MAX_VINT: u64 = 0x00FFFFFFFFFFFFFF;

/// VInts as per RFC8794
///
/// The implementation assumes that EBMLMaxSizeLength=8
pub trait Vint: Into<u64> + Copy {
    /// Returns a representation of the current value as a VInt array.
    ///
    /// # Errors
    ///
    /// This can return an error if the value is too large to be representable as a VInt.
    fn as_vint(self) -> Result<Vec<u8>, VIntError> {
        let casted: u64 = self.into();
        if casted >= MAX_VINT {
            return Err(VIntError::WriteOverflow(casted));
        }

        let mut sliced = casted.to_be_bytes();
        let mut real_start: usize = 0;
        while real_start < 7 && sliced[real_start] == 0 {
            real_start += 1;
        }

        let mut vint_marker = 1u8 << real_start;
        if sliced[real_start] >= vint_marker {
            real_start = real_start
                .checked_add_signed(-1)
                .expect("casted < MAX_VINT, so real_start should be > 0");
            vint_marker = 1u8 << real_start;
        }

        sliced[real_start] |= vint_marker;

        Ok(sliced[real_start..].to_vec())
    }

    /// Returns a representation of the current value as an sized-bytes VInt array.
    ///
    /// # Errors
    ///
    /// This can return an error if the value is too large to be representable as a VInt.
    fn as_vint_sized(self, size: Size) -> Result<Vec<u8>, VIntError> {
        if size > EBML_MAX_SIZE_LENGTH {
            return Err(VIntError::VIntTooLarge(size.0, EBML_MAX_SIZE_LENGTH.0));
        }
        let casted: u64 = self.into();
        let largest: u64 = 0xFFFFFFFFFFFFFFFF >> ((EBML_MAX_SIZE_LENGTH - size).0 * 8 + size.0);

        if casted >= largest {
            Err(VIntError::ConstrainedWriteOverflow(casted, size.0))
        } else {
            let vint_marker = 1u8 << (EBML_MAX_SIZE_LENGTH - size).0;
            let first_byte = 8 - size.0 as usize;
            let mut sliced = casted.to_be_bytes()[first_byte..].to_vec();
            sliced[0] |= vint_marker;
            Ok(sliced)
        }
    }
}

impl Vint for u64 {}
impl Vint for u32 {}
impl Vint for u16 {}
impl Vint for u8 {}

/// Reads a VInt from the beginning of the input array slice.
///
/// The returned tuple contains the value of the VInt (`u64`) and its length (`usize`).
///
/// # Errors
///
/// This method can return a `VIntError` if the input array cannot be read as a VInt or
/// if the array is too short.
pub fn read_vint(buffer: &[u8]) -> Result<(u64, Size), VIntError> {
    if buffer.is_empty() {
        return Err(VIntError::ReadUnderflow(0, 1));
    }

    if buffer[0] == 0 {
        return Err(VIntError::InvalidVInt);
    }

    // this the fastest method we found (cf. benchmark_read_vint branch)
    //  - .ilog2() is slightly faster than .leading_zeros()
    //  - the "value loop" is faster than buffer.copy_from_slice(); .from_be_bytes()

    let length: Size = (8 - buffer[0].ilog2() as u64).into();

    if length.0 as usize > buffer.len() {
        // Not enough data in the buffer to read out the vint value
        return Err(VIntError::ReadUnderflow(buffer.len(), length.0 as usize));
    }

    let mut value: u64 = buffer[0].into();
    value -= 1 << (8 - length.0);

    for item in buffer.iter().take(length.0 as usize).skip(1) {
        value <<= 8;
        value += u64::from(*item);
    }

    Ok((value, length))
}

/// Computes an EBML-compliant checksum from given slice
pub fn crc32(buf: &[u8]) -> [u8; 4] {
    hash(buf).to_le_bytes()
}

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

    #[test]
    fn test_parse_roundtrip() {
        let variants = [
            // EBML Header as per RFC 8794
            MosaicTag::Ebml,
            MosaicTag::EbmlVersion,
            MosaicTag::EbmlReadVersion,
            MosaicTag::EbmlMaxIdLength,
            MosaicTag::EbmlMaxSizeLength,
            MosaicTag::DocType,
            MosaicTag::DocTypeVersion,
            MosaicTag::DocTypeReadVersion,
            MosaicTag::DocTypeExtension,
            MosaicTag::DocTypeExtensionName,
            MosaicTag::DocTypeExtensionVersion,
            MosaicTag::Crc32,
            MosaicTag::Void,
            // MOSAIC elements
            MosaicTag::Mosaic,
            MosaicTag::ContainerMetaData,
            MosaicTag::ObjectsCounter,
            MosaicTag::ObjectsTotalSize,
            MosaicTag::CompressionMethod,
            MosaicTag::CompressionData,
            MosaicTag::Comment,
            MosaicTag::EndOfTilesOffset,
            MosaicTag::Tile,
            MosaicTag::Object,
            MosaicTag::Index,
            MosaicTag::IdxDescription,
            MosaicTag::IdxUnrolled,
            MosaicTag::Key,
            MosaicTag::Offset,
            MosaicTag::GoToConflicts,
            MosaicTag::IdxUnrolledEntrySize,
            MosaicTag::MapContainer,
            MosaicTag::Map,
            MosaicTag::Conflicts,
            MosaicTag::Conflict,
            MosaicTag::ConflictingKey,
            MosaicTag::ConflictingOffset,
        ];
        for tag in &variants {
            let bytes = tag.to_be_bytes();
            let n: Size = bytes
                .len()
                .try_into()
                .unwrap_or_else(|_| panic!("Failed to convert {} to Size", bytes.len()));
            assert_eq!(
                MosaicTag::parse(bytes).unwrap(),
                (*tag, n),
                "roundtrip failed for {:?}",
                tag
            );
        }
    }

    #[test]
    fn test_shortest_be_bytes() -> Result<()> {
        let bytes = 0x1234567890123456u64.shortest_be_bytes();
        assert_eq!(bytes, vec![0x12, 0x34, 0x56, 0x78, 0x90, 0x12, 0x34, 0x56]);

        let bytes = 16u64.shortest_be_bytes();
        assert_eq!(bytes, vec![0x10]);

        let bytes = 1u8.shortest_be_bytes();
        assert_eq!(bytes, vec![1]);

        // uInt of size zero is zero
        // www.rfc-editor.org/rfc/rfc8794.html#name-unsigned-integer-element
        let bytes = 0u8.shortest_be_bytes();
        assert_eq!(bytes, vec![]);

        Ok(())
    }

    #[test]
    fn read_vint_sixteen() {
        let buffer = [144];
        let result = read_vint(&buffer).unwrap();

        assert_eq!(16, result.0);
        assert_eq!(1, result.1);
    }

    #[test]
    fn write_vint_sixteen() {
        let result = 16u64.as_vint().expect("Writing vint failed");
        assert_eq!(vec![144u8], result);
    }

    #[test]
    fn read_vint_one_twenty_seven() {
        let buffer = [255u8];
        let result = read_vint(&buffer).unwrap();

        assert_eq!(127, result.0);
        assert_eq!(1, result.1);
    }

    #[test]
    fn write_vint_one_twenty_seven() {
        let result = 127u64.as_vint().expect("Writing vint failed");
        assert_eq!(vec![255u8], result);
    }

    #[test]
    fn read_vint_two_hundred() {
        // 200 is 11001000, can't fit in a 1-byte VInt
        //  64 is 01000000, the 2-bytes marker.
        let buffer = [64, 200];
        let result = read_vint(&buffer).unwrap();

        assert_eq!(200, result.0);
        assert_eq!(2, result.1);
    }

    #[test]
    fn write_vint_two_hundred() {
        let result = 200u64.as_vint().expect("Writing vint failed");
        assert_eq!(vec![64u8, 200u8], result);
    }

    #[test]
    fn read_vint_for_ebml_tag() {
        let buffer = [0x1a, 0x45, 0xdf, 0xa3];
        let result = read_vint(&buffer).unwrap();

        assert_eq!(0x0a45dfa3, result.0);
        assert_eq!(4, result.1);
    }

    #[test]
    fn read_vint_very_long() {
        let buffer = [1, 0, 0, 0, 0, 0, 0, 1];
        let result = read_vint(&buffer).unwrap();

        assert_eq!(1, result.0);
        assert_eq!(8, result.1);
    }

    #[test]
    fn write_vint_sized() {
        let result = 1u64.as_vint_sized(3.into()).expect("Writing vint failed");
        assert_eq!(vec![0x20, 0, 1], result);

        let result = 0x1FFFFEu64
            .as_vint_sized(3.into())
            .expect("Writing vint failed");
        assert_eq!(vec![0x3F, 0xFF, 0xFE], result);

        let result = 1u64
            .as_vint_sized(EBML_MAX_SIZE_LENGTH)
            .expect("Writing vint failed");
        assert_eq!(vec![1, 0, 0, 0, 0, 0, 0, 1], result);
    }

    #[test]
    fn write_vint_sized_errors() {
        let result = 1u64.as_vint_sized(EBML_MAX_SIZE_LENGTH + 1.into());
        assert_matches!(result.err().unwrap(), VIntError::VIntTooLarge(_, _));

        let result = 0x8FFFu64.as_vint_sized(2.into());
        assert_matches!(
            result.err().unwrap(),
            VIntError::ConstrainedWriteOverflow(_, _)
        );

        let result = 0xEFFFFFFFFFu64.as_vint_sized(5.into());
        assert_matches!(
            result.err().unwrap(),
            VIntError::ConstrainedWriteOverflow(_, _)
        );
    }

    #[test]
    fn read_vint_overflow() {
        let buffer = [1, 0, 0, 0];
        let result = read_vint(&buffer);

        assert!(result.is_err());
    }

    #[test]
    #[should_panic]
    fn too_big_for_vint() {
        (1u64 << 56).as_vint().expect("Writing vint failed");
    }

    #[test]
    fn vint_encode_decode_range() {
        for val in 0..500_000 {
            let bytes = val.as_vint().unwrap();
            let result = read_vint(bytes.as_slice()).unwrap().0;
            assert_eq!(val, result);
        }
    }
}