Skip to main content

lora_snapshot/
view.rs

1use crate::body::BodyReader;
2use crate::codec::SnapshotInfo;
3use crate::errors::{Result, SnapshotCodecError};
4use crate::format::BODY_FORMAT_VERSION;
5
6#[derive(Debug, Clone)]
7pub struct SnapshotView<'a> {
8    info: SnapshotInfo,
9    body: BodyView<'a>,
10}
11
12impl<'a> SnapshotView<'a> {
13    pub(crate) fn parse(info: SnapshotInfo, bytes: &'a [u8]) -> Result<Self> {
14        Ok(Self {
15            info,
16            body: BodyView::parse(bytes)?,
17        })
18    }
19
20    #[must_use]
21    pub fn info(&self) -> &SnapshotInfo {
22        &self.info
23    }
24
25    #[must_use]
26    pub fn next_node_id(&self) -> u64 {
27        self.body.next_node_id
28    }
29
30    #[must_use]
31    pub fn next_rel_id(&self) -> u64 {
32        self.body.next_rel_id
33    }
34
35    #[must_use]
36    pub fn node_ids(&self) -> U64ColumnView<'a> {
37        self.body.node_ids
38    }
39
40    #[must_use]
41    pub fn relationship_ids(&self) -> U64ColumnView<'a> {
42        self.body.rel_ids
43    }
44
45    #[must_use]
46    pub fn relationship_sources(&self) -> U64ColumnView<'a> {
47        self.body.rel_src
48    }
49
50    #[must_use]
51    pub fn relationship_targets(&self) -> U64ColumnView<'a> {
52        self.body.rel_dst
53    }
54
55    #[must_use]
56    pub fn relationship_type_ids(&self) -> U32ColumnView<'a> {
57        self.body.rel_type_ids
58    }
59
60    pub fn labels_for_node_index(
61        &self,
62        index: usize,
63    ) -> Result<impl Iterator<Item = &'a str> + '_> {
64        let start = u32_to_usize(
65            self.body.node_label_offsets.get(index).ok_or_else(|| {
66                SnapshotCodecError::Decode("node label offset out of bounds".into())
67            })?,
68            "node label offset",
69        )?;
70        let end = u32_to_usize(
71            self.body.node_label_offsets.get(index + 1).ok_or_else(|| {
72                SnapshotCodecError::Decode("node label offset out of bounds".into())
73            })?,
74            "node label offset",
75        )?;
76        if start > end || end > self.body.node_label_ids.len() {
77            return Err(SnapshotCodecError::Decode(
78                "invalid node label offset".into(),
79            ));
80        }
81        Ok(self.body.node_label_ids.slice(start, end).map(move |id| {
82            u32_to_usize(id, "label id")
83                .ok()
84                .and_then(|index| self.body.label_dictionary.get(index))
85                .unwrap_or("<invalid-label>")
86        }))
87    }
88
89    #[must_use]
90    pub fn relationship_type(&self, type_id: u32) -> Option<&'a str> {
91        self.body
92            .rel_type_dictionary
93            .get(usize::try_from(type_id).ok()?)
94    }
95}
96
97#[derive(Debug, Clone)]
98struct BodyView<'a> {
99    next_node_id: u64,
100    next_rel_id: u64,
101    node_ids: U64ColumnView<'a>,
102    label_dictionary: StringTableView<'a>,
103    node_label_offsets: U32ColumnView<'a>,
104    node_label_ids: U32ColumnView<'a>,
105    rel_ids: U64ColumnView<'a>,
106    rel_src: U64ColumnView<'a>,
107    rel_dst: U64ColumnView<'a>,
108    rel_type_dictionary: StringTableView<'a>,
109    rel_type_ids: U32ColumnView<'a>,
110}
111
112impl<'a> BodyView<'a> {
113    fn parse(bytes: &'a [u8]) -> Result<Self> {
114        let mut reader = BodyReader::new(bytes);
115        let version = reader.read_u32()?;
116        if version != BODY_FORMAT_VERSION {
117            return Err(SnapshotCodecError::Decode(format!(
118                "unsupported snapshot body format version {version}"
119            )));
120        }
121        let next_node_id = reader.read_u64()?;
122        let next_rel_id = reader.read_u64()?;
123        let node_ids = reader.read_u64_column_view()?;
124        let label_dictionary = reader.read_string_table_view()?;
125        let node_label_offsets = reader.read_u32_column_view()?;
126        let node_label_ids = reader.read_u32_column_view()?;
127        let rel_ids = reader.read_u64_column_view()?;
128        let rel_src = reader.read_u64_column_view()?;
129        let rel_dst = reader.read_u64_column_view()?;
130        let rel_type_dictionary = reader.read_string_table_view()?;
131        let rel_type_ids = reader.read_u32_column_view()?;
132        let view = Self {
133            next_node_id,
134            next_rel_id,
135            node_ids,
136            label_dictionary,
137            node_label_offsets,
138            node_label_ids,
139            rel_ids,
140            rel_src,
141            rel_dst,
142            rel_type_dictionary,
143            rel_type_ids,
144        };
145        view.validate()?;
146        Ok(view)
147    }
148
149    fn validate(&self) -> Result<()> {
150        let expected_offsets = self
151            .node_ids
152            .len()
153            .checked_add(1)
154            .ok_or_else(|| SnapshotCodecError::Decode("node column length overflow".into()))?;
155        if self.node_label_offsets.len() != expected_offsets {
156            return Err(SnapshotCodecError::Decode(
157                "node label offset length mismatch".into(),
158            ));
159        }
160        if self.rel_ids.len() != self.rel_src.len()
161            || self.rel_ids.len() != self.rel_dst.len()
162            || self.rel_ids.len() != self.rel_type_ids.len()
163        {
164            return Err(SnapshotCodecError::Decode(
165                "relationship column length mismatch".into(),
166            ));
167        }
168        let mut previous = 0usize;
169        for index in 0..self.node_label_offsets.len() {
170            let offset = u32_to_usize(
171                self.node_label_offsets.get(index).ok_or_else(|| {
172                    SnapshotCodecError::Decode("node label offset out of bounds".into())
173                })?,
174                "node label offset",
175            )?;
176            if offset < previous || offset > self.node_label_ids.len() {
177                return Err(SnapshotCodecError::Decode(
178                    "invalid node label offset".into(),
179                ));
180            }
181            previous = offset;
182        }
183        if previous != self.node_label_ids.len() {
184            return Err(SnapshotCodecError::Decode(
185                "node label offsets do not cover all label ids".into(),
186            ));
187        }
188        for type_id in self.rel_type_ids.iter() {
189            let type_id = u32_to_usize(type_id, "relationship type id")?;
190            if self.rel_type_dictionary.get(type_id).is_none() {
191                return Err(SnapshotCodecError::Decode(
192                    "invalid relationship type id".into(),
193                ));
194            }
195        }
196        for label_id in self.node_label_ids.iter() {
197            let label_id = u32_to_usize(label_id, "label id")?;
198            if self.label_dictionary.get(label_id).is_none() {
199                return Err(SnapshotCodecError::Decode("invalid label id".into()));
200            }
201        }
202        Ok(())
203    }
204}
205
206fn u32_to_usize(value: u32, label: &str) -> Result<usize> {
207    usize::try_from(value)
208        .map_err(|_| SnapshotCodecError::Decode(format!("{label} does not fit in usize")))
209}
210
211#[derive(Debug, Clone, Copy)]
212pub struct U64ColumnView<'a> {
213    bytes: &'a [u8],
214    len: usize,
215}
216
217impl<'a> U64ColumnView<'a> {
218    pub(crate) fn new(bytes: &'a [u8], len: usize) -> Self {
219        Self { bytes, len }
220    }
221
222    #[must_use]
223    pub fn len(&self) -> usize {
224        self.len
225    }
226
227    #[must_use]
228    pub fn is_empty(&self) -> bool {
229        self.len == 0
230    }
231
232    #[must_use]
233    pub fn get(&self, index: usize) -> Option<u64> {
234        if index >= self.len {
235            return None;
236        }
237        let start = index.checked_mul(8)?;
238        let end = start.checked_add(8)?;
239        Some(u64::from_le_bytes(
240            self.bytes.get(start..end)?.try_into().ok()?,
241        ))
242    }
243
244    pub fn iter(&self) -> impl Iterator<Item = u64> + '_ {
245        (0..self.len).filter_map(|index| self.get(index))
246    }
247}
248
249#[derive(Debug, Clone, Copy)]
250pub struct U32ColumnView<'a> {
251    bytes: &'a [u8],
252    len: usize,
253}
254
255impl<'a> U32ColumnView<'a> {
256    pub(crate) fn new(bytes: &'a [u8], len: usize) -> Self {
257        Self { bytes, len }
258    }
259
260    #[must_use]
261    pub fn len(&self) -> usize {
262        self.len
263    }
264
265    #[must_use]
266    pub fn is_empty(&self) -> bool {
267        self.len == 0
268    }
269
270    #[must_use]
271    pub fn get(&self, index: usize) -> Option<u32> {
272        if index >= self.len {
273            return None;
274        }
275        let start = index.checked_mul(4)?;
276        let end = start.checked_add(4)?;
277        Some(u32::from_le_bytes(
278            self.bytes.get(start..end)?.try_into().ok()?,
279        ))
280    }
281
282    pub fn iter(&self) -> impl Iterator<Item = u32> + '_ {
283        (0..self.len).filter_map(|index| self.get(index))
284    }
285
286    pub(crate) fn slice(&self, start: usize, end: usize) -> impl Iterator<Item = u32> + '_ {
287        (start..end).filter_map(|index| self.get(index))
288    }
289}
290
291#[derive(Debug, Clone)]
292pub struct StringTableView<'a> {
293    entries: Vec<&'a str>,
294}
295
296impl<'a> StringTableView<'a> {
297    pub(crate) fn new(entries: Vec<&'a str>) -> Self {
298        Self { entries }
299    }
300
301    #[must_use]
302    pub fn len(&self) -> usize {
303        self.entries.len()
304    }
305
306    #[must_use]
307    pub fn is_empty(&self) -> bool {
308        self.len() == 0
309    }
310
311    #[must_use]
312    pub fn get(&self, index: usize) -> Option<&'a str> {
313        self.entries.get(index).copied()
314    }
315
316    pub fn iter(&self) -> impl Iterator<Item = &'a str> + '_ {
317        self.entries.iter().copied()
318    }
319}