Skip to main content

zeppelin_embed/graph/
block.rs

1//! Fixed-stride, cache-line-aligned graph node-block persistence.
2
3use crate::graph::refine::RefinementPasses;
4use crate::kernels::Bit4Row;
5use crate::quant::Bit4Factors;
6use xxhash_rust::xxh3::xxh3_64;
7
8/// Cache-line width frozen by the graph node-block v1 format.
9pub const CACHE_LINE_BYTES: usize = 128;
10/// Fixed-width metadata trailer following all directly addressable blocks.
11pub const NODE_BLOCK_TRAILER_LEN: usize = CACHE_LINE_BYTES;
12/// Sentinel occupying every unused neighbour slot.
13pub const UNUSED_NEIGHBOR_ID: u32 = u32::MAX;
14
15const NODE_BLOCK_MAGIC: [u8; 8] = *b"ZEGRNB01";
16const NODE_BLOCK_VERSION: u16 = 1;
17const NODE_METADATA_BYTES: usize = 16;
18const ALLOWED_NODE_FLAGS: u8 = 0b0000_0011;
19
20/// Validated geometry shared by every node block in one region.
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
22pub struct GraphNodeLayout {
23    dims: u32,
24    padded_dims: u32,
25    max_degree: u8,
26    stride: u32,
27}
28
29impl GraphNodeLayout {
30    /// Derives the frozen stride from logical/padded dimensions and maximum degree.
31    pub fn new(dims: u32, padded_dims: u32, max_degree: u8) -> Result<Self, GraphNodeError> {
32        if dims == 0 {
33            return Err(GraphNodeError::Layout(
34                "logical dimensions must not be zero".to_owned(),
35            ));
36        }
37        if padded_dims < dims {
38            return Err(GraphNodeError::Layout(format!(
39                "padded dimensions {padded_dims} are below logical dimensions {dims}"
40            )));
41        }
42        if !padded_dims.is_multiple_of(CACHE_LINE_BYTES as u32) {
43            return Err(GraphNodeError::Layout(format!(
44                "padded dimensions {padded_dims} are not a multiple of {CACHE_LINE_BYTES}"
45            )));
46        }
47        let code_bytes = usize::try_from(padded_dims)
48            .map_err(|_| GraphNodeError::Layout("padded dimensions exceed usize".to_owned()))?
49            .div_ceil(2);
50        let neighbor_bytes = usize::from(max_degree)
51            .checked_mul(std::mem::size_of::<u32>())
52            .ok_or_else(|| GraphNodeError::Layout("neighbour bytes overflow".to_owned()))?;
53        let unaligned = code_bytes
54            .checked_add(NODE_METADATA_BYTES)
55            .and_then(|value| value.checked_add(neighbor_bytes))
56            .ok_or_else(|| GraphNodeError::Layout("node block size overflow".to_owned()))?;
57        let stride = round_up_cache_line(unaligned)?;
58        Ok(Self {
59            dims,
60            padded_dims,
61            max_degree,
62            stride: u32::try_from(stride)
63                .map_err(|_| GraphNodeError::Layout("node stride exceeds u32".to_owned()))?,
64        })
65    }
66
67    /// Returns the logical vector dimension.
68    #[must_use]
69    pub const fn dims(self) -> u32 {
70        self.dims
71    }
72
73    /// Returns the stored, zero-padded vector dimension.
74    #[must_use]
75    pub const fn padded_dims(self) -> u32 {
76        self.padded_dims
77    }
78
79    /// Returns the fixed neighbour-slot capacity.
80    #[must_use]
81    pub const fn max_degree(self) -> u8 {
82        self.max_degree
83    }
84
85    /// Returns packed Bit4 bytes stored at the start of every block.
86    #[must_use]
87    pub const fn code_bytes(self) -> usize {
88        (self.padded_dims as usize).div_ceil(2)
89    }
90
91    /// Returns the cache-line-rounded bytes between adjacent node ids.
92    #[must_use]
93    pub const fn stride(self) -> u32 {
94        self.stride
95    }
96
97    /// Returns the byte offset for a dense segment-local node/row id.
98    pub fn block_offset(self, node_id: u32) -> Result<u32, GraphNodeError> {
99        node_id
100            .checked_mul(self.stride)
101            .ok_or(GraphNodeError::ArithmeticOverflow)
102    }
103}
104
105/// Borrowed values written into one fixed-stride node block.
106#[derive(Clone, Copy, Debug)]
107pub struct GraphNodeBlockInput<'a> {
108    /// Existing frozen MSB-first packed Bit4 row, including zero padding.
109    pub codes: &'a [u8],
110    /// Persisted Bit4 factor fields in their frozen declaration order.
111    pub factors: Bit4Factors,
112    /// Bit 0 marks an entry seed; bit 1 marks a hub.
113    pub flags: u8,
114    /// Dense segment-local neighbour row ids.
115    pub neighbors: &'a [u32],
116}
117
118/// Complete borrowed node-block set supplied to the production writer.
119#[derive(Clone, Copy, Debug)]
120pub struct GraphNodeBlockBuild<'a> {
121    /// Shared fixed-stride geometry.
122    pub layout: GraphNodeLayout,
123    /// Nodes in dense row-id order; array position is the graph node id.
124    pub nodes: &'a [GraphNodeBlockInput<'a>],
125}
126
127/// Owned encoded graph region produced by the production writer.
128#[derive(Clone, Debug, Eq, PartialEq)]
129pub struct EncodedNodeBlocks {
130    bytes: Vec<u8>,
131    layout: GraphNodeLayout,
132    node_count: u32,
133}
134
135/// Validated borrowed view over one complete mmap-backed graph region.
136#[derive(Clone, Copy, Debug)]
137pub struct GraphNodeBlocks<'a> {
138    bytes: &'a [u8],
139    layout: GraphNodeLayout,
140    node_count: u32,
141    block_bytes: usize,
142    refinement_passes: RefinementPasses,
143}
144
145/// Owned geometry retained after one complete immutable graph validation.
146#[derive(Clone, Copy, Debug)]
147pub(crate) struct ValidatedGraphNodeBlocks {
148    layout: GraphNodeLayout,
149    node_count: u32,
150    block_bytes: usize,
151    region_bytes: usize,
152    refinement_passes: RefinementPasses,
153}
154
155/// Dense node id whose fixed-stride byte offset was checked once.
156#[derive(Clone, Copy, Debug, Eq, PartialEq)]
157pub(super) struct CheckedNodeId {
158    raw: u32,
159    offset: usize,
160}
161
162impl CheckedNodeId {
163    #[must_use]
164    pub(super) const fn raw(self) -> u32 {
165        self.raw
166    }
167}
168
169impl GraphNodeBlocks<'_> {
170    /// Returns the geometry validated from the fixed metadata trailer.
171    #[must_use]
172    pub const fn layout(&self) -> GraphNodeLayout {
173        self.layout
174    }
175
176    /// Returns the dense segment-local graph node / row count.
177    #[must_use]
178    pub const fn node_count(&self) -> u32 {
179        self.node_count
180    }
181
182    /// Returns the authenticated record of post-consolidation passes applied.
183    #[must_use]
184    pub const fn refinement_passes(&self) -> RefinementPasses {
185        self.refinement_passes
186    }
187
188    pub(crate) const fn validated_descriptor(&self) -> ValidatedGraphNodeBlocks {
189        ValidatedGraphNodeBlocks {
190            layout: self.layout,
191            node_count: self.node_count,
192            block_bytes: self.block_bytes,
193            region_bytes: self.bytes.len(),
194            refinement_passes: self.refinement_passes,
195        }
196    }
197
198    /// Returns one validated node block by its identical segment row id.
199    pub fn block(&self, node_id: u32) -> Result<GraphNodeBlock<'_>, GraphNodeError> {
200        self.block_checked(self.checked_node_id(node_id)?)
201    }
202
203    pub(super) fn checked_node_id(&self, node_id: u32) -> Result<CheckedNodeId, GraphNodeError> {
204        if node_id >= self.node_count {
205            return Err(GraphNodeError::NodeIdOutOfRange {
206                node_id,
207                node_count: self.node_count,
208            });
209        }
210        let offset = usize::try_from(self.layout.block_offset(node_id)?)
211            .map_err(|_| GraphNodeError::ArithmeticOverflow)?;
212        Ok(CheckedNodeId {
213            raw: node_id,
214            offset,
215        })
216    }
217
218    pub(super) fn block_checked(
219        &self,
220        node_id: CheckedNodeId,
221    ) -> Result<GraphNodeBlock<'_>, GraphNodeError> {
222        let end = node_id
223            .offset
224            .checked_add(self.layout.stride as usize)
225            .ok_or(GraphNodeError::ArithmeticOverflow)?;
226        let block = self
227            .bytes
228            .get(node_id.offset..end.min(self.block_bytes))
229            .ok_or_else(|| GraphNodeError::InvalidNode {
230                node_id: node_id.raw,
231                detail: "validated block range is unavailable".to_owned(),
232            })?;
233        decode_block_view(self.layout, node_id.raw, block)
234    }
235
236    pub(super) fn score_row_checked(
237        &self,
238        node_id: CheckedNodeId,
239    ) -> Result<(Bit4Row<'_>, Bit4Factors), GraphNodeError> {
240        let code_bytes = self.layout.code_bytes();
241        let row_bytes = code_bytes
242            .checked_add(12)
243            .ok_or(GraphNodeError::ArithmeticOverflow)?;
244        let end = node_id
245            .offset
246            .checked_add(row_bytes)
247            .ok_or(GraphNodeError::ArithmeticOverflow)?;
248        let score_row =
249            self.bytes
250                .get(node_id.offset..end)
251                .ok_or_else(|| GraphNodeError::InvalidNode {
252                    node_id: node_id.raw,
253                    detail: "validated score row is unavailable".to_owned(),
254                })?;
255        let Some((codes, factor_bytes)) = score_row.split_at_checked(code_bytes) else {
256            return Err(GraphNodeError::InvalidNode {
257                node_id: node_id.raw,
258                detail: "validated score row width is invalid".to_owned(),
259            });
260        };
261        let factors: [u8; 12] =
262            factor_bytes
263                .try_into()
264                .map_err(|_| GraphNodeError::InvalidNode {
265                    node_id: node_id.raw,
266                    detail: "validated factor width is invalid".to_owned(),
267                })?;
268        let [s0, s1, s2, s3, n0, n1, n2, n3, c0, c1, c2, c3] = factors;
269        Ok((
270            Bit4Row::from_validated_bytes(codes),
271            Bit4Factors::from_persisted(
272                f32::from_le_bytes([s0, s1, s2, s3]),
273                f32::from_le_bytes([n0, n1, n2, n3]),
274                f32::from_le_bytes([c0, c1, c2, c3]),
275            ),
276        ))
277    }
278
279    pub(super) fn adjacency_checked(
280        &self,
281        node_id: CheckedNodeId,
282    ) -> Result<(u8, NeighborIds<'_>), GraphNodeError> {
283        let degree_offset = self
284            .layout
285            .code_bytes()
286            .checked_add(12)
287            .ok_or(GraphNodeError::ArithmeticOverflow)?;
288        let adjacency_bytes = usize::from(self.layout.max_degree)
289            .checked_mul(std::mem::size_of::<u32>())
290            .and_then(|value| value.checked_add(4))
291            .ok_or(GraphNodeError::ArithmeticOverflow)?;
292        let start = node_id
293            .offset
294            .checked_add(degree_offset)
295            .ok_or(GraphNodeError::ArithmeticOverflow)?;
296        let end = start
297            .checked_add(adjacency_bytes)
298            .ok_or(GraphNodeError::ArithmeticOverflow)?;
299        let adjacency = self
300            .bytes
301            .get(start..end)
302            .ok_or_else(|| GraphNodeError::InvalidNode {
303                node_id: node_id.raw,
304                detail: "validated adjacency is unavailable".to_owned(),
305            })?;
306        let [degree, _flags, _reserved0, _reserved1, neighbors @ ..] = adjacency else {
307            return Err(GraphNodeError::InvalidNode {
308                node_id: node_id.raw,
309                detail: "validated adjacency width is invalid".to_owned(),
310            });
311        };
312        Ok((
313            *degree,
314            NeighborIds {
315                chunks: neighbors.chunks_exact(std::mem::size_of::<u32>()),
316            },
317        ))
318    }
319
320    pub(super) fn prefetch_line0(&self, node_id: CheckedNodeId) {
321        if let Some(bytes) = self.bytes.get(node_id.offset..) {
322            prefetch_address(bytes.as_ptr());
323        }
324    }
325
326    pub(super) fn prefetch_head_block(&self, node_id: CheckedNodeId) {
327        let Some(bytes) = self.bytes.get(node_id.offset..) else {
328            return;
329        };
330        prefetch_address(bytes.as_ptr());
331        if self.layout.stride as usize > CACHE_LINE_BYTES
332            && let Some(second_line) = bytes.get(CACHE_LINE_BYTES..)
333        {
334            prefetch_address(second_line.as_ptr());
335        }
336    }
337}
338
339impl ValidatedGraphNodeBlocks {
340    pub(crate) fn bind<'a>(self, bytes: &'a [u8]) -> Result<GraphNodeBlocks<'a>, GraphNodeError> {
341        if bytes.len() != self.region_bytes {
342            return Err(GraphNodeError::InvalidHeader(format!(
343                "cached graph region length {}, current length {}",
344                self.region_bytes,
345                bytes.len()
346            )));
347        }
348        Ok(GraphNodeBlocks {
349            bytes,
350            layout: self.layout,
351            node_count: self.node_count,
352            block_bytes: self.block_bytes,
353            refinement_passes: self.refinement_passes,
354        })
355    }
356}
357
358#[inline]
359fn prefetch_address(address: *const u8) {
360    #[cfg(target_arch = "aarch64")]
361    // SAFETY: callers derive the address from a live immutable graph-region slice.
362    unsafe {
363        std::arch::asm!(
364            "prfm pldl1keep, [{address}]",
365            address = in(reg) address,
366            options(readonly, nostack)
367        );
368    }
369    #[cfg(not(target_arch = "aarch64"))]
370    let _ = address;
371}
372
373/// Borrowed fields from one validated fixed-stride node block.
374#[derive(Clone, Copy, Debug)]
375pub struct GraphNodeBlock<'a> {
376    codes: &'a [u8],
377    factors: Bit4Factors,
378    degree: u8,
379    flags: u8,
380    reserved: [u8; 2],
381    neighbors: &'a [u8],
382    padding: &'a [u8],
383}
384
385impl GraphNodeBlock<'_> {
386    /// Returns the existing frozen MSB-first packed Bit4 row.
387    #[must_use]
388    pub const fn codes(&self) -> &[u8] {
389        self.codes
390    }
391
392    /// Returns the three persisted Bit4 factor fields.
393    #[must_use]
394    pub const fn factors(&self) -> Bit4Factors {
395        self.factors
396    }
397
398    /// Returns the active neighbour count.
399    #[must_use]
400    pub const fn degree(&self) -> u8 {
401        self.degree
402    }
403
404    /// Returns the entry-seed/hub flag byte.
405    #[must_use]
406    pub const fn flags(&self) -> u8 {
407        self.flags
408    }
409
410    /// Returns the two must-be-zero bytes after degree and flags.
411    #[must_use]
412    pub const fn reserved_bytes(&self) -> [u8; 2] {
413        self.reserved
414    }
415
416    /// Iterates all fixed slots, including `u32::MAX` unused sentinels.
417    #[must_use]
418    pub fn neighbors_padded(&self) -> NeighborIds<'_> {
419        NeighborIds {
420            chunks: self.neighbors.chunks_exact(std::mem::size_of::<u32>()),
421        }
422    }
423
424    /// Returns cache-line rounding bytes, all validated as zero.
425    #[must_use]
426    pub const fn padding_bytes(&self) -> &[u8] {
427        self.padding
428    }
429}
430
431/// Exact-size iterator over little-endian fixed neighbour slots.
432#[derive(Clone, Debug)]
433pub struct NeighborIds<'a> {
434    chunks: std::slice::ChunksExact<'a, u8>,
435}
436
437impl Iterator for NeighborIds<'_> {
438    type Item = u32;
439
440    fn next(&mut self) -> Option<Self::Item> {
441        let raw: [u8; 4] = self.chunks.next()?.try_into().ok()?;
442        Some(u32::from_le_bytes(raw))
443    }
444
445    fn size_hint(&self) -> (usize, Option<usize>) {
446        let length = self.chunks.len();
447        (length, Some(length))
448    }
449}
450
451impl ExactSizeIterator for NeighborIds<'_> {}
452
453impl EncodedNodeBlocks {
454    /// Returns the complete region bytes, including the fixed metadata trailer.
455    #[must_use]
456    pub fn as_bytes(&self) -> &[u8] {
457        &self.bytes
458    }
459
460    /// Consumes the wrapper and returns the complete persisted region bytes.
461    #[must_use]
462    pub fn into_bytes(self) -> Vec<u8> {
463        self.bytes
464    }
465
466    /// Returns the byte offset of a present dense node id.
467    pub fn block_offset(&self, node_id: u32) -> Result<u32, GraphNodeError> {
468        if node_id >= self.node_count {
469            return Err(GraphNodeError::NodeIdOutOfRange {
470                node_id,
471                node_count: self.node_count,
472            });
473        }
474        self.layout.block_offset(node_id)
475    }
476}
477
478/// Typed graph node-block encode/decode failure.
479#[derive(Clone, Debug, Eq, PartialEq)]
480pub enum GraphNodeError {
481    /// The region is shorter than its bounded fixed trailer or block geometry.
482    Truncated {
483        /// Minimum bytes required by the declared format.
484        minimum: usize,
485        /// Bytes supplied by the caller.
486        actual: usize,
487    },
488    /// Fixed magic, version, flags, lengths, or reserved trailer bytes are invalid.
489    InvalidHeader(String),
490    /// Layout geometry cannot describe a bounded fixed-stride region.
491    Layout(String),
492    /// A node id is outside the dense row-id range.
493    NodeIdOutOfRange {
494        /// Requested graph node / segment row id.
495        node_id: u32,
496        /// Number of nodes present in the region.
497        node_count: u32,
498    },
499    /// One node's production input violates the persisted contract.
500    InvalidNode {
501        /// Dense graph node / segment row id.
502        node_id: u32,
503        /// Contract violation.
504        detail: String,
505    },
506    /// Checked offset or length arithmetic overflowed.
507    ArithmeticOverflow,
508}
509
510impl std::fmt::Display for GraphNodeError {
511    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
512        match self {
513            Self::Truncated { minimum, actual } => write!(
514                formatter,
515                "graph node region is truncated: need {minimum} bytes, got {actual}"
516            ),
517            Self::InvalidHeader(detail) => {
518                write!(formatter, "graph node region header is invalid: {detail}")
519            }
520            Self::Layout(detail) => write!(formatter, "graph node layout is invalid: {detail}"),
521            Self::NodeIdOutOfRange {
522                node_id,
523                node_count,
524            } => write!(
525                formatter,
526                "graph node id {node_id} is outside dense row count {node_count}"
527            ),
528            Self::InvalidNode { node_id, detail } => {
529                write!(formatter, "graph node {node_id} is invalid: {detail}")
530            }
531            Self::ArithmeticOverflow => {
532                formatter.write_str("graph node block arithmetic overflowed")
533            }
534        }
535    }
536}
537
538impl std::error::Error for GraphNodeError {}
539
540/// Encodes node blocks in dense row-id order followed by a 128-byte trailer.
541pub fn encode_node_blocks(
542    build: GraphNodeBlockBuild<'_>,
543) -> Result<EncodedNodeBlocks, GraphNodeError> {
544    encode_node_blocks_with_refinement_passes(build, RefinementPasses::default())
545}
546
547/// Encodes node blocks with an authenticated additive refinement-pass record.
548pub fn encode_node_blocks_with_refinement_passes(
549    build: GraphNodeBlockBuild<'_>,
550    refinement_passes: RefinementPasses,
551) -> Result<EncodedNodeBlocks, GraphNodeError> {
552    let node_count = u32::try_from(build.nodes.len())
553        .map_err(|_| GraphNodeError::Layout("node count exceeds u32".to_owned()))?;
554    let block_bytes = usize::try_from(build.layout.stride)
555        .map_err(|_| GraphNodeError::ArithmeticOverflow)?
556        .checked_mul(build.nodes.len())
557        .ok_or(GraphNodeError::ArithmeticOverflow)?;
558    let total_bytes = block_bytes
559        .checked_add(NODE_BLOCK_TRAILER_LEN)
560        .ok_or(GraphNodeError::ArithmeticOverflow)?;
561    let mut output = Vec::with_capacity(total_bytes);
562    for (position, node) in build.nodes.iter().enumerate() {
563        let node_id = u32::try_from(position).map_err(|_| GraphNodeError::ArithmeticOverflow)?;
564        encode_node(build.layout, node_id, node, node_count, &mut output)?;
565    }
566    encode_trailer(build.layout, node_count, refinement_passes, &mut output);
567    Ok(EncodedNodeBlocks {
568        bytes: output,
569        layout: build.layout,
570        node_count,
571    })
572}
573
574/// Validates and borrows a complete fixed-stride graph node-block region.
575pub fn decode_node_blocks(bytes: &[u8]) -> Result<GraphNodeBlocks<'_>, GraphNodeError> {
576    if bytes.len() < NODE_BLOCK_TRAILER_LEN {
577        return Err(GraphNodeError::Truncated {
578            minimum: NODE_BLOCK_TRAILER_LEN,
579            actual: bytes.len(),
580        });
581    }
582    let trailer_start = bytes
583        .len()
584        .checked_sub(NODE_BLOCK_TRAILER_LEN)
585        .ok_or(GraphNodeError::ArithmeticOverflow)?;
586    let trailer = bytes
587        .get(trailer_start..)
588        .ok_or(GraphNodeError::Truncated {
589            minimum: NODE_BLOCK_TRAILER_LEN,
590            actual: bytes.len(),
591        })?;
592    if trailer.get(..8) != Some(NODE_BLOCK_MAGIC.as_slice()) {
593        return Err(GraphNodeError::InvalidHeader(
594            "magic does not match ZEGRNB01".to_owned(),
595        ));
596    }
597    let version = read_u16(trailer, 8, "version")?;
598    if version != NODE_BLOCK_VERSION {
599        return Err(GraphNodeError::InvalidHeader(format!(
600            "version {version}, expected {NODE_BLOCK_VERSION}"
601        )));
602    }
603    let refinement_bits = read_u16(trailer, 10, "refinement passes")?;
604    let refinement_passes = RefinementPasses::from_bits(refinement_bits).ok_or_else(|| {
605        GraphNodeError::InvalidHeader(format!(
606            "unknown refinement-pass bits {refinement_bits:#06x}"
607        ))
608    })?;
609    let dims = read_u32(trailer, 12, "logical dimensions")?;
610    let padded_dims = read_u32(trailer, 16, "padded dimensions")?;
611    let max_degree = *trailer.get(20).ok_or(GraphNodeError::Truncated {
612        minimum: 21,
613        actual: trailer.len(),
614    })?;
615    if trailer
616        .get(21..24)
617        .is_none_or(|reserved| reserved.iter().any(|byte| *byte != 0))
618    {
619        return Err(GraphNodeError::InvalidHeader(
620            "degree padding is non-zero".to_owned(),
621        ));
622    }
623    let stored_stride = read_u32(trailer, 24, "stride")?;
624    let node_count = read_u32(trailer, 28, "node count")?;
625    let stored_checksum = read_u64(trailer, 32, "region checksum")?;
626    let checksummed_end = trailer_start
627        .checked_add(32)
628        .ok_or(GraphNodeError::ArithmeticOverflow)?;
629    let checksummed = bytes
630        .get(..checksummed_end)
631        .ok_or(GraphNodeError::Truncated {
632            minimum: checksummed_end,
633            actual: bytes.len(),
634        })?;
635    let actual_checksum = xxh3_64(checksummed);
636    if actual_checksum != stored_checksum {
637        return Err(GraphNodeError::InvalidHeader(format!(
638            "xxh3-64 expected {stored_checksum:#018x}, computed {actual_checksum:#018x}"
639        )));
640    }
641    if trailer
642        .get(40..)
643        .is_none_or(|reserved| reserved.iter().any(|byte| *byte != 0))
644    {
645        return Err(GraphNodeError::InvalidHeader(
646            "reserved trailer bytes are non-zero".to_owned(),
647        ));
648    }
649    let layout = GraphNodeLayout::new(dims, padded_dims, max_degree)?;
650    if stored_stride != layout.stride {
651        return Err(GraphNodeError::InvalidHeader(format!(
652            "stored stride {stored_stride}, formula yields {}",
653            layout.stride
654        )));
655    }
656    let block_bytes = usize::try_from(layout.stride)
657        .map_err(|_| GraphNodeError::ArithmeticOverflow)?
658        .checked_mul(node_count as usize)
659        .ok_or(GraphNodeError::ArithmeticOverflow)?;
660    if trailer_start != block_bytes {
661        let minimum = block_bytes
662            .checked_add(NODE_BLOCK_TRAILER_LEN)
663            .ok_or(GraphNodeError::ArithmeticOverflow)?;
664        if bytes.len() < minimum {
665            return Err(GraphNodeError::Truncated {
666                minimum,
667                actual: bytes.len(),
668            });
669        }
670        return Err(GraphNodeError::InvalidHeader(format!(
671            "region length {}, expected {minimum}",
672            bytes.len()
673        )));
674    }
675    for node_id in 0..node_count {
676        let start = usize::try_from(layout.block_offset(node_id)?)
677            .map_err(|_| GraphNodeError::ArithmeticOverflow)?;
678        let end = start
679            .checked_add(layout.stride as usize)
680            .ok_or(GraphNodeError::ArithmeticOverflow)?;
681        let block = bytes.get(start..end).ok_or(GraphNodeError::Truncated {
682            minimum: end,
683            actual: bytes.len(),
684        })?;
685        validate_decoded_node(layout, node_id, node_count, block)?;
686    }
687    Ok(GraphNodeBlocks {
688        bytes,
689        layout,
690        node_count,
691        block_bytes,
692        refinement_passes,
693    })
694}
695
696fn validate_decoded_node(
697    layout: GraphNodeLayout,
698    node_id: u32,
699    node_count: u32,
700    block: &[u8],
701) -> Result<(), GraphNodeError> {
702    let view = decode_block_view(layout, node_id, block)?;
703    validate_code_padding(layout, node_id, view.codes)?;
704    if view.flags & !ALLOWED_NODE_FLAGS != 0 {
705        return Err(GraphNodeError::InvalidNode {
706            node_id,
707            detail: format!("reserved flag bits are non-zero: {:#04x}", view.flags),
708        });
709    }
710    if view.reserved != [0, 0] {
711        return Err(GraphNodeError::InvalidNode {
712            node_id,
713            detail: "reserved bytes are non-zero".to_owned(),
714        });
715    }
716    if view.padding.iter().any(|byte| *byte != 0) {
717        return Err(GraphNodeError::InvalidNode {
718            node_id,
719            detail: "cache-line padding is non-zero".to_owned(),
720        });
721    }
722    let degree = usize::from(view.degree);
723    if degree > usize::from(layout.max_degree) {
724        return Err(GraphNodeError::InvalidNode {
725            node_id,
726            detail: format!("degree {degree}, maximum {}", layout.max_degree),
727        });
728    }
729    for (slot, neighbor) in view.neighbors_padded().enumerate() {
730        if slot < degree {
731            if neighbor >= node_count {
732                return Err(GraphNodeError::InvalidNode {
733                    node_id,
734                    detail: format!(
735                        "active neighbour slot {slot} id {neighbor} is outside row count {node_count}"
736                    ),
737                });
738            }
739        } else if neighbor != UNUSED_NEIGHBOR_ID {
740            return Err(GraphNodeError::InvalidNode {
741                node_id,
742                detail: format!("unused neighbour slot {slot} holds {neighbor}, expected u32::MAX"),
743            });
744        }
745    }
746    Ok(())
747}
748
749fn decode_block_view<'a>(
750    layout: GraphNodeLayout,
751    node_id: u32,
752    block: &'a [u8],
753) -> Result<GraphNodeBlock<'a>, GraphNodeError> {
754    if block.len() != layout.stride as usize {
755        return Err(GraphNodeError::InvalidNode {
756            node_id,
757            detail: format!("block bytes {}, expected {}", block.len(), layout.stride),
758        });
759    }
760    let code_end = layout.code_bytes();
761    let codes = block.get(..code_end).ok_or(GraphNodeError::Truncated {
762        minimum: code_end,
763        actual: block.len(),
764    })?;
765    let scale = read_f32(block, code_end, "scale")?;
766    let norm = read_f32(block, code_end.saturating_add(4), "normalized norm")?;
767    let correction = read_f32(block, code_end.saturating_add(8), "normalized correction")?;
768    if [scale, norm, correction]
769        .iter()
770        .any(|value| !value.is_finite())
771    {
772        return Err(GraphNodeError::InvalidNode {
773            node_id,
774            detail: "factor record contains a non-finite value".to_owned(),
775        });
776    }
777    let degree_offset = code_end.saturating_add(12);
778    let degree = *block.get(degree_offset).ok_or(GraphNodeError::Truncated {
779        minimum: degree_offset.saturating_add(1),
780        actual: block.len(),
781    })?;
782    let flags = *block
783        .get(degree_offset.saturating_add(1))
784        .ok_or(GraphNodeError::Truncated {
785            minimum: degree_offset.saturating_add(2),
786            actual: block.len(),
787        })?;
788    let reserved: [u8; 2] = block
789        .get(degree_offset.saturating_add(2)..degree_offset.saturating_add(4))
790        .ok_or(GraphNodeError::Truncated {
791            minimum: degree_offset.saturating_add(4),
792            actual: block.len(),
793        })?
794        .try_into()
795        .map_err(|_| GraphNodeError::InvalidNode {
796            node_id,
797            detail: "reserved field width is invalid".to_owned(),
798        })?;
799    let neighbors_start = code_end.saturating_add(NODE_METADATA_BYTES);
800    let neighbors_end = neighbors_start
801        .checked_add(usize::from(layout.max_degree).saturating_mul(4))
802        .ok_or(GraphNodeError::ArithmeticOverflow)?;
803    let neighbors = block
804        .get(neighbors_start..neighbors_end)
805        .ok_or(GraphNodeError::Truncated {
806            minimum: neighbors_end,
807            actual: block.len(),
808        })?;
809    let padding = block
810        .get(neighbors_end..)
811        .ok_or(GraphNodeError::Truncated {
812            minimum: neighbors_end,
813            actual: block.len(),
814        })?;
815    Ok(GraphNodeBlock {
816        codes,
817        factors: Bit4Factors::from_persisted(scale, norm, correction),
818        degree,
819        flags,
820        reserved,
821        neighbors,
822        padding,
823    })
824}
825
826fn read_u16(bytes: &[u8], offset: usize, field: &str) -> Result<u16, GraphNodeError> {
827    let end = offset
828        .checked_add(2)
829        .ok_or(GraphNodeError::ArithmeticOverflow)?;
830    let raw: [u8; 2] = bytes
831        .get(offset..end)
832        .ok_or(GraphNodeError::Truncated {
833            minimum: end,
834            actual: bytes.len(),
835        })?
836        .try_into()
837        .map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
838    Ok(u16::from_le_bytes(raw))
839}
840
841fn read_u32(bytes: &[u8], offset: usize, field: &str) -> Result<u32, GraphNodeError> {
842    let end = offset
843        .checked_add(4)
844        .ok_or(GraphNodeError::ArithmeticOverflow)?;
845    let raw: [u8; 4] = bytes
846        .get(offset..end)
847        .ok_or(GraphNodeError::Truncated {
848            minimum: end,
849            actual: bytes.len(),
850        })?
851        .try_into()
852        .map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
853    Ok(u32::from_le_bytes(raw))
854}
855
856fn read_f32(bytes: &[u8], offset: usize, field: &str) -> Result<f32, GraphNodeError> {
857    read_u32(bytes, offset, field).map(f32::from_bits)
858}
859
860fn read_u64(bytes: &[u8], offset: usize, field: &str) -> Result<u64, GraphNodeError> {
861    let end = offset
862        .checked_add(8)
863        .ok_or(GraphNodeError::ArithmeticOverflow)?;
864    let raw: [u8; 8] = bytes
865        .get(offset..end)
866        .ok_or(GraphNodeError::Truncated {
867            minimum: end,
868            actual: bytes.len(),
869        })?
870        .try_into()
871        .map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
872    Ok(u64::from_le_bytes(raw))
873}
874
875fn encode_node(
876    layout: GraphNodeLayout,
877    node_id: u32,
878    node: &GraphNodeBlockInput<'_>,
879    node_count: u32,
880    output: &mut Vec<u8>,
881) -> Result<(), GraphNodeError> {
882    if node.codes.len() != layout.code_bytes() {
883        return Err(GraphNodeError::InvalidNode {
884            node_id,
885            detail: format!(
886                "code bytes {}, expected {}",
887                node.codes.len(),
888                layout.code_bytes()
889            ),
890        });
891    }
892    validate_code_padding(layout, node_id, node.codes)?;
893    if node
894        .factors
895        .persisted_fields()
896        .iter()
897        .any(|field| !field.is_finite())
898    {
899        return Err(GraphNodeError::InvalidNode {
900            node_id,
901            detail: "factor record contains a non-finite value".to_owned(),
902        });
903    }
904    if node.flags & !ALLOWED_NODE_FLAGS != 0 {
905        return Err(GraphNodeError::InvalidNode {
906            node_id,
907            detail: format!("reserved flag bits are non-zero: {:#04x}", node.flags),
908        });
909    }
910    if node.neighbors.len() > usize::from(layout.max_degree) {
911        return Err(GraphNodeError::InvalidNode {
912            node_id,
913            detail: format!(
914                "degree {}, maximum {}",
915                node.neighbors.len(),
916                layout.max_degree
917            ),
918        });
919    }
920    if let Some(invalid) = node
921        .neighbors
922        .iter()
923        .copied()
924        .find(|neighbor| *neighbor >= node_count)
925    {
926        return Err(GraphNodeError::InvalidNode {
927            node_id,
928            detail: format!("neighbour id {invalid} is outside row count {node_count}"),
929        });
930    }
931
932    let start = output.len();
933    output.extend_from_slice(node.codes);
934    for field in node.factors.persisted_fields() {
935        output.extend_from_slice(&field.to_bits().to_le_bytes());
936    }
937    let degree = u8::try_from(node.neighbors.len()).map_err(|_| GraphNodeError::InvalidNode {
938        node_id,
939        detail: "degree exceeds u8".to_owned(),
940    })?;
941    output.push(degree);
942    output.push(node.flags);
943    output.extend_from_slice(&0_u16.to_le_bytes());
944    for neighbor in node.neighbors {
945        output.extend_from_slice(&neighbor.to_le_bytes());
946    }
947    for _ in node.neighbors.len()..usize::from(layout.max_degree) {
948        output.extend_from_slice(&UNUSED_NEIGHBOR_ID.to_le_bytes());
949    }
950    let end = start
951        .checked_add(layout.stride as usize)
952        .ok_or(GraphNodeError::ArithmeticOverflow)?;
953    output.resize(end, 0);
954    Ok(())
955}
956
957fn validate_code_padding(
958    layout: GraphNodeLayout,
959    node_id: u32,
960    codes: &[u8],
961) -> Result<(), GraphNodeError> {
962    let logical_bytes = (layout.dims as usize).div_ceil(2);
963    if !layout.dims.is_multiple_of(2)
964        && codes
965            .get(logical_bytes.saturating_sub(1))
966            .is_some_and(|byte| byte & 0x0f != 0)
967    {
968        return Err(GraphNodeError::InvalidNode {
969            node_id,
970            detail: "unused low nibble is non-zero".to_owned(),
971        });
972    }
973    if codes
974        .get(logical_bytes..)
975        .is_some_and(|padding| padding.iter().any(|byte| *byte != 0))
976    {
977        return Err(GraphNodeError::InvalidNode {
978            node_id,
979            detail: "padded Bit4 dimensions are non-zero".to_owned(),
980        });
981    }
982    Ok(())
983}
984
985fn encode_trailer(
986    layout: GraphNodeLayout,
987    node_count: u32,
988    refinement_passes: RefinementPasses,
989    output: &mut Vec<u8>,
990) {
991    let start = output.len();
992    output.extend_from_slice(&NODE_BLOCK_MAGIC);
993    output.extend_from_slice(&NODE_BLOCK_VERSION.to_le_bytes());
994    output.extend_from_slice(&refinement_passes.bits().to_le_bytes());
995    output.extend_from_slice(&layout.dims.to_le_bytes());
996    output.extend_from_slice(&layout.padded_dims.to_le_bytes());
997    output.push(layout.max_degree);
998    output.extend_from_slice(&[0_u8; 3]);
999    output.extend_from_slice(&layout.stride.to_le_bytes());
1000    output.extend_from_slice(&node_count.to_le_bytes());
1001    let checksum = xxh3_64(output);
1002    output.extend_from_slice(&checksum.to_le_bytes());
1003    output.resize(start.saturating_add(NODE_BLOCK_TRAILER_LEN), 0);
1004}
1005
1006fn round_up_cache_line(value: usize) -> Result<usize, GraphNodeError> {
1007    value
1008        .checked_add(CACHE_LINE_BYTES.saturating_sub(1))
1009        .map(|rounded| rounded / CACHE_LINE_BYTES * CACHE_LINE_BYTES)
1010        .ok_or(GraphNodeError::ArithmeticOverflow)
1011}
1012
1013#[cfg(test)]
1014#[allow(clippy::expect_used)]
1015mod tests {
1016    use super::{
1017        GraphNodeBlockBuild, GraphNodeBlockInput, GraphNodeError, GraphNodeLayout,
1018        decode_node_blocks, encode_node_blocks,
1019    };
1020    use crate::quant::Bit4Factors;
1021
1022    #[test]
1023    fn cached_descriptor_rejects_a_different_mapping_length() {
1024        let layout = GraphNodeLayout::new(128, 128, 0).expect("cache fixture layout");
1025        let codes = [0_u8; 64];
1026        let node = GraphNodeBlockInput {
1027            codes: &codes,
1028            factors: Bit4Factors::from_persisted(1.0, 0.0, 0.0),
1029            flags: 1,
1030            neighbors: &[],
1031        };
1032        let encoded = encode_node_blocks(GraphNodeBlockBuild {
1033            layout,
1034            nodes: &[node],
1035        })
1036        .expect("cache fixture encodes");
1037        let graph = decode_node_blocks(encoded.as_bytes()).expect("cache fixture decodes");
1038        let descriptor = graph.validated_descriptor();
1039
1040        let shortened = encoded
1041            .as_bytes()
1042            .get(..encoded.as_bytes().len().saturating_sub(1))
1043            .expect("fixture has a trailer byte to remove");
1044        assert!(matches!(
1045            descriptor.bind(shortened),
1046            Err(GraphNodeError::InvalidHeader(detail))
1047                if detail.contains("cached graph region length")
1048        ));
1049        let rebound = descriptor
1050            .bind(encoded.as_bytes())
1051            .expect("identical mapping length binds");
1052        assert_eq!(rebound.block(0).expect("node zero").flags(), 1);
1053
1054        let outside = encoded
1055            .block_offset(1)
1056            .expect_err("one-node graph rejects row one");
1057        assert_eq!(
1058            outside.to_string(),
1059            "graph node id 1 is outside dense row count 1"
1060        );
1061        assert_eq!(
1062            GraphNodeError::ArithmeticOverflow.to_string(),
1063            "graph node block arithmetic overflowed"
1064        );
1065    }
1066}