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

//! Memory-mapped back-end for local MOSAIC file reading

use crate::ebml::{read_vint, MosaicTag};
use crate::{Position, Size};
use anyhow::{ensure, Context, Result};
use mmap_rs::{Mmap, MmapFlags, MmapOptions};
use std::fs::File;
use std::path::Path;

use super::MosaicBackend;

/// Memory-mapped backend implementation for MOSAIC file reading
pub struct MmapMosaicBackend {
    mmap: Mmap,
}

impl MosaicBackend for MmapMosaicBackend {
    fn read_one_tag_and_size(&self, mut offset: Position) -> Result<(MosaicTag, Size, Position)> {
        let (tag, tag_size) =
            MosaicTag::parse(self.mmap.get(offset.0..).context("EOF while reading tag")?)?;
        offset += tag_size;
        let (size, size_of_size) = read_vint(
            self.mmap
                .get(offset.0..)
                .context("EOF while reading vint")?,
        )?;
        offset += size_of_size;

        Ok((tag, size.into(), offset))
    }

    fn read_binary(&self, offset: Position, size: Size) -> Result<(&[u8], Position)> {
        let next_offset = offset + size;
        let element = &self
            .mmap
            .get(offset.0..next_offset.0)
            .context("EOF while reading binary")?;

        Ok((element, next_offset))
    }

    fn read_u8(&self, offset: Position) -> Result<(u8, Position)> {
        let element = self.mmap.get(offset.0).context("EOF while reading u8")?;
        Ok((*element, offset + 1.into()))
    }

    fn read_uint(&self, offset: Position, size: Size) -> Result<(u64, Position)> {
        if size == 0 {
            return Ok((0, offset));
        }

        ensure!(size <= 8, "Unsupported uint size: {size}");

        let mut bytes = [0u8; 8];
        bytes[8 - usize::try_from(size.0)?..].copy_from_slice(
            self.mmap
                .get(offset.0..(offset + size).0)
                .context("EOF while reading uint")?,
        );

        Ok((u64::from_be_bytes(bytes), offset + size))
    }

    fn read_utf8(&self, mut offset: Position, size: Size) -> Result<(String, Position)> {
        let element = self
            .mmap
            .get(offset.0..(offset + size).0)
            .context("EOF while reading utf8 string")?
            .to_vec();
        offset += size;

        Ok((String::from_utf8(element)?, offset))
    }

    fn len(&self) -> usize {
        self.mmap.len()
    }

    fn is_empty(&self) -> bool {
        self.mmap.is_empty()
    }
}

impl MmapMosaicBackend {
    /// Create a new memory-mapped backend from a file path
    pub fn new(path: &Path) -> Result<Self> {
        let file =
            File::open(path).with_context(|| format!("Could not open {}", path.display()))?;
        let file_size = file
            .metadata()
            .with_context(|| format!("Could not stat {}", path.display()))?
            .len() as usize;

        let mmap = unsafe {
            MmapOptions::new(file_size)
                .context("Could not initialize mmap")?
                .with_file(&file, 0)
        }
        .with_flags(MmapFlags::RANDOM_ACCESS)
        .map()
        .with_context(|| format!("Could not mmap {}", path.display()))?;

        Ok(Self { mmap })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn test_read_tag_and_size() {
        // Create a temporary file with EBML-like structure
        let mut temp_file = NamedTempFile::new().unwrap();
        // Simple EBML tag structure: 0x1A45DFA3 (EBML tag) + size
        let test_data = b"\x1A\x45\xDF\xA3\x81\x00"; // EBML tag (4 bytes) + size (1 byte)
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.find_tag_and_size(0.into());
        assert!(result.is_ok());
        let (tag, size, offset) = result.unwrap();
        assert_eq!(tag, MosaicTag::Ebml);
        assert_eq!(size, 1);
        assert_eq!(offset, 5);
    }

    #[test]
    fn test_read_binary() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"test binary data";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_binary(0.into(), test_data.len().try_into().unwrap());
        assert!(result.is_ok());
        let (data, offset) = result.unwrap();
        assert_eq!(data, test_data);
        assert_eq!(offset, test_data.len());
    }

    #[test]
    fn test_read_u8() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"A";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_u8(0.into());
        assert!(result.is_ok());
        let (byte, offset) = result.unwrap();
        assert_eq!(byte, b'A');
        assert_eq!(offset, 1);
    }

    #[test]
    fn test_read_uint_0_bytes() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\xAA\xBB\xCC\xDD\xEE\xFF\x11\x22\x01";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 0.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 0);
        assert_eq!(offset, 0);
    }

    #[test]
    fn test_read_uint_1_byte() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\x2A";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 1.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 42);
        assert_eq!(offset, 1);
    }

    #[test]
    fn test_read_uint_2_bytes() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\x03\xE8";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 2.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 1000);
        assert_eq!(offset, 2);
    }

    #[test]
    fn test_read_uint_3_bytes() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\xAA\xBB\xCC";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 3.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 11189196);
        assert_eq!(offset, 3);
    }

    #[test]
    fn test_read_uint_4_bytes() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\xAA\xBB\xCC\xDD";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 4.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 2864434397);
        assert_eq!(offset, 4);
    }

    #[test]
    fn test_read_uint_5_bytes() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\xAA\xBB\xCC\xDD\xEE";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 5.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 733295205870);
        assert_eq!(offset, 5);
    }

    #[test]
    fn test_read_uint_6_bytes() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\xAA\xBB\xCC\xDD\xEE\xFF";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 6.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 187723572702975);
        assert_eq!(offset, 6);
    }

    #[test]
    fn test_read_uint_7_bytes() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\xAA\xBB\xCC\xDD\xEE\xFF\x11";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 7.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 48057234611961617);
        assert_eq!(offset, 7);
    }

    #[test]
    fn test_read_uint_8_bytes() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\xAA\xBB\xCC\xDD\xEE\xFF\x11\x22";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 8.into());
        assert!(result.is_ok());
        let (value, offset) = result.unwrap();
        assert_eq!(value, 12302652060662173986);
        assert_eq!(offset, 8);
    }

    #[test]
    fn test_read_uint_unsupported_size() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"\xAA\xBB\xCC\xDD\xEE\xFF\x11\x22\x01";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_uint(0.into(), 9.into());
        assert!(result.is_err());
    }

    #[test]
    fn test_read_utf8() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"Hello, World!";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        let result = backend.read_utf8(0.into(), test_data.len().try_into().unwrap());
        assert!(result.is_ok());
        let (string, offset) = result.unwrap();
        assert_eq!(string, "Hello, World!");
        assert_eq!(offset, test_data.len());
    }

    #[test]
    fn test_len() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"test data for length";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        assert_eq!(backend.len(), test_data.len());
    }

    #[test]
    fn test_new() {
        let mut temp_file = NamedTempFile::new().unwrap();
        let test_data = b"test data";
        temp_file.write_all(test_data).unwrap();
        let path = temp_file.path();

        let backend = MmapMosaicBackend::new(path).unwrap();
        assert_eq!(backend.len(), test_data.len());
    }
}