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 read MOSAIC files
//!
//! Given a MOSAIC local file, use [MmapMosaicReader::new].

use crate::backends::mmap::MmapMosaicBackend;
use crate::backends::MosaicBackend;
use crate::creator::max_object_size;
use crate::ebml::{MosaicTag, DOCTYPE_READ_VERSION};
use crate::{CompressionMethod, IdxDescription, Position, Size};
use anyhow::{ensure, Result};
use ph::fmph::GOFunction;
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use thiserror::Error;
use zstd::bulk::decompress;

/// Typed errors for MosaicReader
#[derive(Error, Debug)]
pub enum MosaicReaderError {
    #[error("Bad checksum for tag {tag}")]
    BadChecksum { tag: MosaicTag },

    #[error("Bad key size, expected {expected} bytes, got {actual} bytes")]
    BadKeySize { expected: u64, actual: u64 },

    #[error("{filename}: Error, {broken_tiles} corrupted tiles and {broken_indexes} corrupted indexes found")]
    BrokenMosaic {
        filename: String,
        broken_tiles: i32,
        broken_indexes: i32,
    },

    #[error("Element {tag} not found")]
    ElementNotFound { tag: MosaicTag },

    #[error("Index {idx_description} not found")]
    IndexNotFound { idx_description: IdxDescription },

    #[error("Missing checksum for element {tag}")]
    MissingChecksum { tag: MosaicTag },

    #[error("Unknown index description {description}")]
    UnknownIdxDescription { description: String },
}

/// Wrapper around [`MosaicBackend`] that remembers the current position so users
/// don't have to manage offsets explicitly.
pub struct Cursor<B: MosaicBackend> {
    pub backend: Arc<B>,
    pub offset: Position,
    mark: Option<Position>,
}

impl<B: MosaicBackend> Cursor<B> {
    pub fn new(backend: Arc<B>) -> Self {
        Cursor {
            backend,
            offset: Position(0),
            mark: None,
        }
    }

    pub fn new_at(backend: Arc<B>, offset: Position) -> Self {
        Cursor {
            backend,
            offset,
            mark: None,
        }
    }

    /// Reads an EBML tag and its size at the current offset. Skips voids automatically.
    ///
    /// Returns the tag and element's size.
    pub fn next_tag_and_size(&mut self) -> Result<(MosaicTag, Size)> {
        let (tag, size, new_offset) = self.backend.find_tag_and_size(self.offset)?;
        self.offset = new_offset;
        Ok((tag, size))
    }

    /// Reads an EBML tag and its size at the current offset, may be a void element.
    ///
    /// Returns the tag and element's size.
    pub fn next_one_tag_and_size(&mut self) -> Result<(MosaicTag, Size)> {
        let (element, size, new_offset) = self.backend.read_one_tag_and_size(self.offset)?;
        self.offset = new_offset;
        Ok((element, size))
    }

    /// Reads binary data of given size at the current offset
    /// Returns the binary data.
    pub fn next_binary(&mut self, size: Size) -> Result<&[u8]> {
        let (binary, new_offset) = self.backend.read_binary(self.offset, size)?;
        self.offset = new_offset;
        Ok(binary)
    }

    /// Reads a single byte at the current offset
    /// Returns the byte
    pub fn next_u8(&mut self) -> Result<u8> {
        let (byte, new_offset) = self.backend.read_u8(self.offset)?;
        self.offset = new_offset;
        Ok(byte)
    }

    /// Reads an unsigned integer of given size from the current offset
    /// Returns the integer
    pub fn next_uint(&mut self, size: Size) -> Result<u64> {
        let (uint, new_offset) = self.backend.read_uint(self.offset, size)?;
        self.offset = new_offset;
        Ok(uint)
    }

    /// Reads a UTF-8 string of given size from the current offset
    /// Returns the string
    pub fn next_utf8(&mut self, size: Size) -> Result<String> {
        let (string, new_offset) = self.backend.read_utf8(self.offset, size)?;
        self.offset = new_offset;
        Ok(string)
    }

    /// Finds the next element after the current offset in the MOSAIC whose tag matches the `element_tag`.
    ///
    ///  Leaves the cursor at the beginning of the element's data and returns its size.
    ///
    /// It will raise `MosaicReaderError::ElementNotFound(...)` if no element matches `element_tag`
    /// after `offset`.
    pub fn seek_to_element(&mut self, element_tag: MosaicTag) -> Result<Size> {
        let (size, new_offset) = self.backend.find_element(self.offset, element_tag)?;
        self.offset = new_offset;
        Ok(size)
    }

    /// Finds the next element after the current offset in the MOSAIC whose tag matches the `element_tag`.
    ///
    ///  Leaves the cursor at the beginning of the element's data and returns its size
    ///  and the position of the beginning of the element's tag.
    ///
    /// It will raise `MosaicReaderError::ElementNotFound(...)` if no element matches `element_tag`
    /// after `offset`.
    pub fn seek_to_element_pre(&mut self, element_tag: MosaicTag) -> Result<(Position, Size)> {
        let (tag_offset, size, new_offset) =
            self.backend.find_element_pre(self.offset, element_tag)?;
        self.offset = new_offset;
        Ok((tag_offset, size))
    }

    /// Finds the next master element after the current offset in the MOSAIC whose tag matches the
    /// `master_tag`.
    ///
    /// It will also compute and validate the CRC32 of the found master element, and raise
    /// `MosaicReaderError::BadChecksum(...)` if the validation fails or
    /// `MosaicReaderError::MissingChecksum(...)` if no CRC32 element is found.
    ///
    /// Returns the size of the element's data and seeks to first offset of element's data, adjusted to
    /// skip the CRC32 checksum.
    ///
    /// It will raise `MosaicReaderError::ElementNotFound(...)` if no element matches `master_tag`
    /// after `offset`.
    pub fn seek_to_and_check(&mut self, master_tag: MosaicTag) -> Result<Size> {
        let (size, new_offset) = self.backend.find_and_check(self.offset, master_tag)?;
        self.offset = new_offset;
        Ok(size)
    }

    /// Seek forward by the given number of bytes
    pub fn skip(&mut self, size: Size) -> Result<()> {
        ensure!(self.offset + size <= self.backend.len());
        self.offset += size;
        Ok(())
    }

    /// After (successfully) seeking a master element, call this to remember where
    /// this element ends so you can call `skip_master` later.
    pub fn mark_master(&mut self, size: Size) -> Result<()> {
        self.mark = Some(self.offset + size);
        Ok(())
    }

    /// Seek at the end of a master element. This requires you called `mark_master`
    /// before. Marks are not stacked.
    pub fn skip_master(&mut self) -> Result<()> {
        match self.mark {
            Some(position) => {
                self.mark = None;
                self.offset = position;
                Ok(())
            }
            None => panic!("skip_master called without mark_master"),
        }
    }
}

/// Reader class for MOSAIC files
/// -----------------------------
///
/// This reader is optimized for randomized lookups of contents per key. A MOSAIC
/// file may be indexed over multiple keys, but `MosaicReader` will use at most one
/// index at a time. This index can be chosen by [MosaicReader::load_index], or at
/// construction time as in [MmapMosaicReader::new].
///
/// The `object_size_bound` is theoric and should only be used as a buffer boundary: it
/// is the maximum object size supported by the MOSAIC file (considering its settings),
/// it is _not_ the actual biggest object's size.
pub struct MosaicReader<B: MosaicBackend> {
    pub path: PathBuf,
    pub backend: Arc<B>, // wrapped in an Arc so MosaicReader can be cloned
    pub doctype: String,
    pub doctype_version: u8,
    pub doctype_read_version: u8,
    pub objects_counter: u64,
    pub objects_total_size: Size,
    pub end_of_tiles_offset: Position,
    pub comments: Vec<String>,
    pub idx: Option<IdxDescription>,
    mph: Option<Arc<GOFunction>>, // wrapped in an Arc so MosaicReader can be cloned cheaply
    idx_unrolled_entry_size: Size,
    idx_unrolled_start: Position,
    pub compression: CompressionMethod,
    pub object_size_bound: usize,
}

impl<B: MosaicBackend> Clone for MosaicReader<B> {
    fn clone(&self) -> Self {
        let MosaicReader {
            path,
            backend,
            doctype,
            doctype_version,
            doctype_read_version,
            objects_counter,
            objects_total_size,
            end_of_tiles_offset,
            comments,
            idx,
            mph,
            idx_unrolled_entry_size,
            idx_unrolled_start,
            compression,
            object_size_bound,
        } = self;

        MosaicReader {
            path: path.clone(),
            backend: backend.clone(),
            doctype: doctype.clone(),
            doctype_version: *doctype_version,
            doctype_read_version: *doctype_read_version,
            objects_counter: *objects_counter,
            objects_total_size: *objects_total_size,
            end_of_tiles_offset: *end_of_tiles_offset,
            comments: comments.clone(),
            idx: *idx,
            mph: mph.clone(),
            idx_unrolled_entry_size: *idx_unrolled_entry_size,
            idx_unrolled_start: *idx_unrolled_start,
            compression: *compression,
            object_size_bound: *object_size_bound,
        }
    }
}

impl<B: MosaicBackend> fmt::Debug for MosaicReader<B> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        f.debug_struct("MosaicReader")
            .field("path", &self.path)
            .field("doctype_version", &self.doctype_version)
            .field("objects_counter", &self.objects_counter)
            .field("objects_total_size", &self.objects_total_size)
            .finish_non_exhaustive()
    }
}

impl<B: MosaicBackend> MosaicReader<B> {
    /// Creates a new MosaicReader with the given backend
    pub fn with_backend(path: &Path, backend: Arc<B>) -> Result<Self> {
        let mut cursor = Cursor::new(backend.clone());

        let size = cursor.seek_to_element(MosaicTag::Ebml)?;
        assert!(size >= 33);

        cursor.seek_to_element(MosaicTag::EbmlVersion)?;
        let ebml_version = cursor.next_u8()?;

        cursor.seek_to_element(MosaicTag::EbmlReadVersion)?;
        let ebml_read_version = cursor.next_u8()?;

        assert!(
            ebml_version >= ebml_read_version,
            "EBML read version ({}) is bigger than EBML version ({})",
            ebml_read_version,
            ebml_version
        );

        cursor.seek_to_element(MosaicTag::EbmlMaxIdLength)?;
        let _ebml_max_id_length = cursor.next_u8()?;

        cursor.seek_to_element(MosaicTag::EbmlMaxSizeLength)?;
        let _ebml_max_size_length = cursor.next_u8()?;

        let size = cursor.seek_to_element(MosaicTag::DocType)?;
        let doctype = cursor.next_utf8(size)?;

        cursor.seek_to_element(MosaicTag::DocTypeVersion)?;
        let doctype_version = cursor.next_u8()?;

        cursor.seek_to_element(MosaicTag::DocTypeReadVersion)?;
        let doctype_read_version = cursor.next_u8()?;

        assert!(
            DOCTYPE_READ_VERSION >= doctype_read_version,
            "This version of swh-mosaic requires DocTypeReadVersion <= {},
            this MOSAIC file is version {}. If possible, install a newer version of swh-mosaic",
            DOCTYPE_READ_VERSION,
            doctype_read_version
        );
        assert!(
            doctype_version >= doctype_read_version,
            "EBML read version ({}) is bigger than EBML version ({})",
            doctype_version,
            doctype_read_version
        );
        assert!(cursor.offset >= 38);

        cursor.seek_to_element(MosaicTag::Mosaic)?;

        cursor.seek_to_and_check(MosaicTag::ContainerMetaData)?;

        let size = cursor.seek_to_element(MosaicTag::ObjectsCounter)?;
        let objects_counter = cursor.next_uint(size)?;

        let size = cursor.seek_to_element(MosaicTag::ObjectsTotalSize)?;
        let objects_total_size = cursor.next_uint(size)?;

        let size = cursor.seek_to_element(MosaicTag::EndOfTilesOffset)?;
        let end_of_tiles_offset = cursor.next_uint(size)?;

        let size = cursor.seek_to_element(MosaicTag::CompressionMethod)?;
        let compression_name = cursor.next_utf8(size)?;
        let compression = CompressionMethod::try_from(compression_name)?;

        let mut comments = Vec::new();
        let mut tag_and_size_or_err = cursor.seek_to_element(MosaicTag::Comment);
        while tag_and_size_or_err.is_ok() {
            let size = tag_and_size_or_err?;
            let comment = cursor.next_utf8(size)?;
            comments.push(comment);
            tag_and_size_or_err = cursor.seek_to_element(MosaicTag::Comment)
        }

        // cf. max_object_size computation: the size of any Tile should allow us to
        // re-compute the max_object_size used when writing this MOSAIC
        let tile_size_or_err = cursor.seek_to_element(MosaicTag::Tile);
        let object_size_bound = if let Ok(size) = tile_size_or_err {
            max_object_size(size.0)?.try_into()?
        } else {
            // the Err case (Tile not found) should only happen to empty MOSAICs
            0
        };

        Ok(MosaicReader {
            path: path.to_path_buf(),
            backend,
            doctype,
            doctype_version,
            doctype_read_version,
            objects_counter,
            objects_total_size: objects_total_size.into(),
            end_of_tiles_offset: (end_of_tiles_offset as usize).into(),
            comments,
            idx: None,
            mph: None,
            idx_unrolled_entry_size: 0.into(),
            idx_unrolled_start: 0.into(),
            compression,
            object_size_bound,
        })
    }

    /// Loads the required index and its corresponding MPH in memory.
    ///
    /// It will raise `MosaicReaderError::MasterElementNotFound(MosaicTag::Index)` if the required
    /// index is not found in the MOSAIC.
    pub fn load_index(&mut self, idx_description: IdxDescription) -> Result<()> {
        let mut cursor = Cursor::new_at(self.backend.clone(), self.end_of_tiles_offset);

        loop {
            let idx_seek = cursor.seek_to_element(MosaicTag::Index);
            if let Ok(idx_size) = idx_seek {
                cursor.mark_master(idx_size)?;
            } else {
                return Err(MosaicReaderError::IndexNotFound { idx_description }.into());
            }

            let description_size = cursor.seek_to_element(MosaicTag::IdxDescription)?;
            let description = cursor.next_utf8(description_size)?;
            if description == idx_description.description() {
                break;
            } else {
                cursor.skip_master()?;
            }
        }

        self.idx = Some(idx_description);

        let idx_unrolled_size = cursor.seek_to_element(MosaicTag::IdxUnrolled)?;
        self.idx_unrolled_start = cursor.offset;
        cursor.skip(idx_unrolled_size)?;

        let idx_unrolled_entry_size_size =
            cursor.seek_to_element(MosaicTag::IdxUnrolledEntrySize)?;
        self.idx_unrolled_entry_size = cursor.next_uint(idx_unrolled_entry_size_size)?.into();

        cursor.seek_to_and_check(MosaicTag::MapContainer)?;
        let mph_size = cursor.seek_to_element(MosaicTag::Map)?;
        let mut mph_vec = cursor.next_binary(mph_size)?;
        self.mph = Some(Arc::new(GOFunction::read(&mut mph_vec)?));

        Ok(())
    }

    /// Lists all indexes contained in the MOSAIC file and return the corresponding description and
    /// position (before the tag).
    pub fn list_indexes(&self) -> Result<Vec<(IdxDescription, Position)>> {
        let mut indexes = Vec::new();
        let mut cursor = Cursor::new_at(self.backend.clone(), self.end_of_tiles_offset);

        while let Ok((index_offset, size)) = cursor.seek_to_element_pre(MosaicTag::Index) {
            cursor.mark_master(size)?;

            let description_size = cursor.seek_to_element(MosaicTag::IdxDescription)?;
            let description = cursor.next_utf8(description_size)?;
            let idx_description = IdxDescription::from_description(description.as_str());
            match idx_description {
                Ok(val) => indexes.push((val, index_offset)),
                Err(err) => match err.downcast::<MosaicReaderError>()? {
                    MosaicReaderError::UnknownIdxDescription { description } => {
                        println!("Skipping unknown index description: {description}")
                    }
                    _ => panic!("Unexpected failure while trying to list indexes"),
                },
            }

            cursor.skip_master()?;
        }

        Ok(indexes)
    }

    /// Finds a key in the loaded MPH and returns the offset of the key as well as the offset of the
    /// key tag corresponding to it.
    ///
    /// Returns `Ok(None)` if the key is not found in the MPH, if something is found but is not a
    /// `MosaicTag::Key` or if a key is found but does not match the queried key.
    ///
    /// ⚠️ `MosaicReader::load_index()` must be called before calling this function.
    pub fn find_key(&self, key: &[u8]) -> Result<Option<(Position, Position)>> {
        let mph_key = self
            .mph
            .as_ref()
            .unwrap_or_else(|| panic!("No MPH loaded"))
            .get(key);
        let Some(mph_key) = mph_key else {
            return Ok(None);
        };
        let offset = self.idx_unrolled_start + Size(mph_key * self.idx_unrolled_entry_size.0);

        let (key_tag, key_size, key_offset) = self.backend.find_tag_and_size(offset)?;
        if key_tag != MosaicTag::Key {
            return Ok(None);
        }
        let (key_read, _) = self.backend.read_binary(key_offset, key_size)?;
        if key_read != key {
            return Ok(None);
        }

        Ok(Some((key_offset, offset)))
    }

    /// Finds the object corresponding to the key of size `key_size` located at `key_offset` in the
    /// MOSAIC file and returns its size, its offset in the file and the offset of the
    /// `MosaicTag::Object` tag corresponding to it.
    ///
    /// It will return `Ok(None)` if the offset corresponding to the given key is not a
    /// `MosaicTag::Offset` or if the item corresponding to the offset is not a `MosaicTag::Object`.
    pub fn find_object_offset(
        &self,
        key_size: Size,
        key_offset: Position,
    ) -> Result<Option<(Size, Position, Position)>> {
        let (_, offset) = self.backend.read_binary(key_offset, key_size)?;
        let (offset_tag, offset_size, offset) = self.backend.find_tag_and_size(offset)?;
        if offset_tag != MosaicTag::Offset {
            return Ok(None);
        }
        let offset = (self.backend.read_uint(offset, offset_size)?.0 as usize).into();

        let (object_tag, object_size, object_offset) = self.backend.find_tag_and_size(offset)?;
        if object_tag != MosaicTag::Object {
            return Ok(None);
        }
        Ok(Some((object_size, object_offset, offset)))
    }

    /// Returns the value corresponding to the given key, if contained in that file.
    ///
    /// ⚠️ `MosaicReader::load_index()` must be called before calling this function.
    pub fn lookup(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let expected_key_len = match self.idx {
            Some(val) => val.key_len(),
            None => panic!("Trying to look up a key with no index loaded. Make sure to call MosaicReader::load_index() first."),
        };
        if u64::try_from(key.len())? != expected_key_len.0 {
            return Err(MosaicReaderError::BadKeySize {
                expected: expected_key_len.0,
                actual: key.len().try_into()?,
            }
            .into());
        }

        let Some((key_offset, _)) = self.find_key(key)? else {
            return Ok(None);
        };
        let key_size = Size(key.len().try_into()?);

        let Some((object_size, offset, _)) = self.find_object_offset(key_size, key_offset)? else {
            return Ok(None);
        };

        let object = self.read_object(offset, object_size)?;

        Ok(Some(object))
    }

    /// multi-threaded version of `MosaicReader::lookup()`. It also requires a
    /// preliminary call to `load_index()`.
    pub fn get_batch(&self, keys: Vec<Vec<u8>>) -> Result<Vec<Option<Vec<u8>>>> {
        keys.par_iter().map(|k| self.lookup(k)).collect()
    }

    /// Finds the tile the object at position `offset` is in and returns the offset of the
    /// `MosaicTag::Tile` tag.
    pub fn find_object_tile(&self, offset: Position) -> Result<Position> {
        let (_, mosaic_offset) = self.backend.find_element(0.into(), MosaicTag::Mosaic)?;
        let (md_size, md_offset) = self
            .backend
            .find_element(mosaic_offset, MosaicTag::ContainerMetaData)?;
        let mut tile_offset = md_offset + md_size;
        let (mut next_tile_offset, _, _) = self
            .backend
            .find_element_pre(tile_offset, MosaicTag::Tile)?;
        while tile_offset < offset && tile_offset != next_tile_offset {
            tile_offset = next_tile_offset;
            (next_tile_offset, _, _) = self
                .backend
                .find_element_pre(tile_offset, MosaicTag::Tile)?;
        }
        Ok(tile_offset)
    }

    pub fn read_object(&self, offset: Position, size: Size) -> Result<Vec<u8>> {
        let raw = self.backend.read_binary(offset, size)?.0;
        match self.compression {
            // thanks to zstd's experimental feature the capacity is actually estimated
            // from the data and bounded by object_size_bound
            CompressionMethod::Zstd => {
                let decompressed = decompress(raw, self.object_size_bound)?;
                Ok(decompressed)
            }
            CompressionMethod::None => Ok(raw.to_vec()),
            _ => anyhow::bail!("Not implemented yet: decompression of {}", self.compression),
        }
    }

    pub fn close(&self) -> Result<()> {
        Ok(())
    }

    /// Returns an iterator on `(key, content)` pairs, where `key` comes from the
    /// currently selected index.
    ///
    /// Pairs are returned in the same order as the [MosaicReader::keys] iterators, but
    /// not as [MosaicReader::values].
    pub fn iter(&self) -> MosaicIterator<B> {
        MosaicIterator::new(self.clone())
    }

    /// Returns an iterator on keys from the currently selected index, according to
    /// the index's order (which can be considered random).
    pub fn keys(&self) -> MosaicKeysIterator<B> {
        MosaicKeysIterator::new(self.clone())
    }

    /// This iterates contents by reading the file linearly. This is faster, but will
    /// not provide contents in the same order as [MosaicReader::iter] or
    /// [MosaicReader::keys].
    pub fn values(&self) -> MosaicValuesIterator<B> {
        MosaicValuesIterator::new(self.clone())
    }
}

/// See [MosaicReader::iter]
pub struct MosaicIterator<B: MosaicBackend> {
    reader: MosaicReader<B>,
    offset: Position,
}

impl<B: MosaicBackend> MosaicIterator<B> {
    pub fn new(reader: MosaicReader<B>) -> Self {
        let offset = reader.idx_unrolled_start;
        Self { reader, offset }
    }
}

impl<B: MosaicBackend> Iterator for MosaicIterator<B> {
    type Item = Result<(Vec<u8>, Vec<u8>)>;

    /// Iterate (key, object) pairs over the currently loaded index.
    fn next(&mut self) -> Option<Self::Item> {
        while let Ok((tag, size, offset)) = self.reader.backend.find_tag_and_size(self.offset) {
            if tag == MosaicTag::Key {
                if let Ok((key, offset)) = self.reader.backend.read_binary(offset, size) {
                    if let Ok((tag, size, offset)) = self.reader.backend.find_tag_and_size(offset) {
                        if tag == MosaicTag::Offset {
                            if let Ok((obj_offset, offset)) =
                                self.reader.backend.read_uint(offset, size)
                            {
                                self.offset = offset;
                                if let Ok((object_tag, object_size, obj_offset)) = self
                                    .reader
                                    .backend
                                    .find_tag_and_size((obj_offset as usize).into())
                                {
                                    if object_tag == MosaicTag::Object {
                                        let key = key.to_vec();
                                        if let Ok(object) =
                                            self.reader.read_object(obj_offset, object_size)
                                        {
                                            return Some(Ok((key, object)));
                                        }
                                    }
                                }
                                continue;
                            }
                        }
                    }
                }
            }
            return None;
        }
        None
    }
}

/// See [MosaicReader::keys]
pub struct MosaicKeysIterator<B: MosaicBackend> {
    cursor: Cursor<B>,
}

impl<B: MosaicBackend> MosaicKeysIterator<B> {
    pub fn new(reader: MosaicReader<B>) -> Self {
        let cursor = Cursor::new_at(reader.backend.clone(), reader.idx_unrolled_start);
        Self { cursor }
    }
}

impl<B: MosaicBackend> Iterator for MosaicKeysIterator<B> {
    type Item = Result<Vec<u8>>;

    /// Iterate keys over the currently loaded index.
    fn next(&mut self) -> Option<Self::Item> {
        let key_size = match self.cursor.next_tag_and_size() {
            Ok((tag, size)) => {
                if tag != MosaicTag::Key {
                    return None;
                }
                size
            }
            Err(_) => return None,
        };
        let key = match self.cursor.next_binary(key_size) {
            Ok(key) => key.to_vec(),
            Err(_) => return None,
        };
        let offset_size = match self.cursor.next_tag_and_size() {
            Ok((tag, size)) => {
                if tag != MosaicTag::Offset {
                    return None;
                }
                size
            }
            Err(_) => return None,
        };
        if let Err(e) = self.cursor.skip(offset_size) {
            return Some(Err(e));
        }
        Some(Ok(key))
    }
}

/// See [MosaicReader::values]
pub struct MosaicValuesIterator<B: MosaicBackend> {
    reader: MosaicReader<B>,
    cursor: Cursor<B>,
}

impl<B: MosaicBackend> MosaicValuesIterator<B> {
    pub fn new(reader: MosaicReader<B>) -> Self {
        let mut cursor = Cursor::new(reader.backend.clone());

        if MosaicValuesIterator::prepare_cursor(&mut cursor).is_err() {
            cursor.offset = reader.backend.len().into();
        }

        Self { reader, cursor }
    }

    /// From the beginning of the file, put cursor at the beginning of the first tile
    fn prepare_cursor(cursor: &mut Cursor<B>) -> Result<()> {
        if cursor.seek_to_element(MosaicTag::Mosaic).is_ok()
            && cursor.seek_to_element(MosaicTag::Tile).is_ok()
        {
            if let Ok((tag, size)) = cursor.next_tag_and_size() {
                if tag == MosaicTag::Crc32 {
                    cursor.skip(size)?;
                    return Ok(());
                }
            }
        }
        anyhow::bail!("Failed to find a Tile");
    }
}

impl<B: MosaicBackend> Iterator for MosaicValuesIterator<B> {
    type Item = Result<Vec<u8>>;

    /// Iterate objects over the currently loaded index.
    ///
    /// Objects will be iterated according to their order in the file
    fn next(&mut self) -> Option<Self::Item> {
        let object_size = match self.cursor.next_tag_and_size() {
            Ok((tag, size)) => match tag {
                MosaicTag::Object => size,
                MosaicTag::Tile => {
                    return self.next();
                }
                MosaicTag::Crc32 => {
                    if let Err(e) = self.cursor.skip(size) {
                        return Some(Err(e));
                    }
                    return self.next();
                }
                _ => return None,
            },
            Err(_) => return None,
        };
        let obj = self.reader.read_object(self.cursor.offset, object_size);
        if let Ok(object) = obj {
            if self.cursor.skip(object_size).is_ok() {
                Some(Ok(object))
            } else {
                None
            }
        } else {
            None
        }
    }
}

/// MosaicReader of a local, memory-mapped file
pub type MmapMosaicReader = MosaicReader<MmapMosaicBackend>;

impl MmapMosaicReader {
    /// Create a new MosaicReader from a local, memory-mapped file, and prepare it to use
    /// the given index.
    ///
    /// This returns `MosaicReaderError::IndexNotFound` if the
    /// required index is not found in that file.
    pub fn new(path: &Path, indexed_by: IdxDescription) -> Result<Self> {
        let backend = MmapMosaicBackend::new(path)?;
        let mut reader = MosaicReader::with_backend(path, Arc::new(backend))?;
        reader.load_index(indexed_by)?;
        Ok(reader)
    }
}