Skip to main content

cherenkov_model_data/
source.rs

1use crate::{DataSpan, ObjectId};
2use anyhow::{Context, Result, ensure};
3use memmap2::Mmap;
4use std::{fs::File, io::Write, path::Path, sync::Arc};
5
6/// An object's identity and total size within a byte source.
7#[derive(Debug, Clone)]
8pub struct ObjectInfo {
9    /// Identity assigned by the source.
10    pub id: ObjectId,
11    /// Object length in bytes.
12    pub bytes: u64,
13}
14
15/// Object IDs belong to one source; they are neither paths nor content hashes.
16/// A copy/import operation must rebind them to its destination's identities.
17pub trait ByteSource {
18    /// List the objects addressable through this source.
19    fn objects(&self) -> Vec<ObjectInfo>;
20    /// Write the requested byte range to `output`, or return an error.
21    fn read(&self, span: &DataSpan, output: &mut dyn Write) -> Result<()>;
22
23    /// Return a retained view, or `None` if this source cannot map bytes.
24    /// Invalid ranges on a mappable source return an error.
25    fn map(&self, _span: &DataSpan) -> Result<Option<MappedBytes>> {
26        Ok(None)
27    }
28}
29
30/// Read-only file mappings, numbered in the order passed to [`Self::open`].
31#[derive(Clone, Default)]
32pub struct MappedObjects {
33    mappings: Vec<Arc<Mmap>>,
34}
35
36/// A byte range that keeps its backing mapping alive after the source is dropped.
37/// The backing file must remain unchanged for the lifetime of every clone.
38#[derive(Clone)]
39pub struct MappedBytes {
40    backing: Arc<Mmap>,
41    range: std::ops::Range<usize>,
42}
43
44impl AsRef<[u8]> for MappedBytes {
45    fn as_ref(&self) -> &[u8] {
46        &self.backing[self.range.clone()]
47    }
48}
49
50impl MappedObjects {
51    /// Map files read-only, assigning object IDs from zero in iteration order.
52    /// Files must remain unchanged while any returned mapping is alive.
53    pub fn open<'a>(paths: impl IntoIterator<Item = &'a Path>) -> Result<Self> {
54        let mappings = paths
55            .into_iter()
56            .map(|path| {
57                let file =
58                    File::open(path).with_context(|| format!("opening {}", path.display()))?;
59                // The caller owns the immutable-file contract, including views
60                // retained after this source is dropped.
61                let map = unsafe { Mmap::map(&file) }
62                    .with_context(|| format!("mapping {}", path.display()))?;
63
64                Ok(Arc::new(map))
65            })
66            .collect::<Result<_>>()?;
67
68        Ok(Self { mappings })
69    }
70
71    /// Borrow an entire object mapping; fail if the object index is out of range.
72    /// The caller must use IDs from this source.
73    pub fn mapping(&self, id: ObjectId) -> Result<&Mmap> {
74        self.mappings
75            .get(id.0)
76            .map(AsRef::as_ref)
77            .context("unknown checkpoint object")
78    }
79
80    /// Borrow a byte range after checking its object ID and bounds.
81    pub fn bytes(&self, span: &DataSpan) -> Result<&[u8]> {
82        let map = self.mapping(span.object)?;
83        let range = checked_range(span, map.len())?;
84
85        Ok(&map[range])
86    }
87
88    /// Retain a byte range without copying its contents, checking ID and bounds.
89    pub fn view(&self, span: &DataSpan) -> Result<MappedBytes> {
90        let backing = self
91            .mappings
92            .get(span.object.0)
93            .context("unknown checkpoint object")?
94            .clone();
95        let range = checked_range(span, backing.len())?;
96
97        Ok(MappedBytes { backing, range })
98    }
99}
100
101impl ByteSource for MappedObjects {
102    fn objects(&self) -> Vec<ObjectInfo> {
103        self.mappings
104            .iter()
105            .enumerate()
106            .map(|(id, map)| ObjectInfo {
107                id: ObjectId(id),
108                bytes: map.len() as u64,
109            })
110            .collect()
111    }
112
113    fn read(&self, span: &DataSpan, output: &mut dyn Write) -> Result<()> {
114        output.write_all(self.bytes(span)?)?;
115
116        Ok(())
117    }
118
119    fn map(&self, span: &DataSpan) -> Result<Option<MappedBytes>> {
120        Ok(Some(self.view(span)?))
121    }
122}
123
124fn checked_range(span: &DataSpan, size: usize) -> Result<std::ops::Range<usize>> {
125    let end = span
126        .offset
127        .checked_add(span.length)
128        .context("data span overflow")?;
129
130    ensure!(end <= size as u64, "data span exceeds checkpoint object");
131
132    Ok(span.offset as usize..end as usize)
133}