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

//!
//! Tools to write MOSAIC files
//!

use crate::ebml::{
    crc32, read_vint, MosaicTag, Vint, CRC32_SIZE, EBML_MAX_ID_LENGTH, EBML_MAX_SIZE_LENGTH,
};
use crate::{Position, Size};
use anyhow::{Error, Result};
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};

/// MOSAIC file writer
///
/// EBML-level file writing helper
pub struct MosaicWriter {
    pub file: File,
}

impl MosaicWriter {
    /// Creates a MosaicWriter that will write to `file` (and fail if it does not exist).
    pub fn new(file: File) -> Result<Self> {
        Ok(MosaicWriter { file })
    }

    /// Computes buffer's CRC32 and write the CRC32 element to target file, returning the sum of the
    /// sizes of the CRC32 tag, the CRC32 tag size and the CRC32 element.
    pub fn write_crc32(&mut self, buffer: &[u8]) -> Result<Size, Error> {
        let mut size = Size(0);
        let crc32 = crc32(buffer);
        self.write_all(&mut size, MosaicTag::Crc32.to_be_bytes())?;
        self.write_all(&mut size, &4u8.as_vint()?)?;
        self.write_all(&mut size, &crc32)?;
        Ok(size)
    }

    /// Reads a master element at `offset` in current file, then update its CRC32
    pub fn rewrite_crc32(&mut self, offset: Position) -> Result<(), Error> {
        self.file.seek(SeekFrom::Start(offset.0.try_into()?))?;
        let mut buffer: Vec<u8> = vec![0; (EBML_MAX_ID_LENGTH + EBML_MAX_SIZE_LENGTH).0 as usize];
        self.file.read_exact(&mut buffer)?;
        let (_tag, size_offset) = read_vint(&buffer)?;
        let (element_size, crc_offset) = read_vint(&buffer[size_offset.0 as usize..])?;
        let content_offset = offset + size_offset + crc_offset + CRC32_SIZE;
        let content_size = Size(element_size) - CRC32_SIZE;

        let mut buffer: Vec<u8> = vec![0; content_size.0 as usize];
        self.file
            .seek(SeekFrom::Start(content_offset.0.try_into()?))?;
        self.file.read_exact(&mut buffer)?;
        let checksum = crc32(&buffer);
        self.file
            .seek(SeekFrom::Start((content_offset - CRC32_SIZE).0.try_into()?))?;
        self.file.write_all(MosaicTag::Crc32.to_be_bytes())?;
        self.file.write_all(&4u8.as_vint()?)?;
        self.file.write_all(&checksum)?;
        Ok(())
    }

    /// Writes a master element with a given tag, returns the size of what was written to file.
    ///
    /// A CRC32 element can be automatically prepended, and the length of the tag's size VINT can be
    /// forced if `size_length` is provided (otherwise it writes the smallest possible VINT).
    pub fn write_master(
        &mut self,
        tag: MosaicTag,
        buffer: &[u8],
        crc32: bool,
        size_length: Option<Size>,
    ) -> Result<Size> {
        let mut size = Size(0);
        self.write_all(&mut size, tag.to_be_bytes())?;
        let buffer_size: Size = if crc32 {
            Size(buffer.len().try_into()?) + CRC32_SIZE
        } else {
            buffer.len().try_into()?
        };

        if let Some(len) = size_length {
            self.write_all(&mut size, &buffer_size.0.as_vint_sized(len)?)?;
        } else {
            self.write_all(&mut size, &buffer_size.0.as_vint()?)?;
        }
        if crc32 {
            size += self.write_crc32(buffer)?;
        }
        self.write_all(&mut size, buffer)?;
        Ok(size)
    }

    fn write_all(&mut self, size: &mut Size, bytes: &[u8]) -> Result<(), Error> {
        *size += bytes.len().try_into()?;
        self.file.write_all(bytes)?;
        Ok(())
    }
}