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;
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    /// A controlled encode/decode operation was interrupted.
482    Control,
483    /// The region is shorter than its bounded fixed trailer or block geometry.
484    Truncated {
485        /// Minimum bytes required by the declared format.
486        minimum: usize,
487        /// Bytes supplied by the caller.
488        actual: usize,
489    },
490    /// Fixed magic, version, flags, lengths, or reserved trailer bytes are invalid.
491    InvalidHeader(String),
492    /// Layout geometry cannot describe a bounded fixed-stride region.
493    Layout(String),
494    /// A node id is outside the dense row-id range.
495    NodeIdOutOfRange {
496        /// Requested graph node / segment row id.
497        node_id: u32,
498        /// Number of nodes present in the region.
499        node_count: u32,
500    },
501    /// One node's production input violates the persisted contract.
502    InvalidNode {
503        /// Dense graph node / segment row id.
504        node_id: u32,
505        /// Contract violation.
506        detail: String,
507    },
508    /// Checked offset or length arithmetic overflowed.
509    ArithmeticOverflow,
510}
511
512impl std::fmt::Display for GraphNodeError {
513    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
514        match self {
515            Self::Control => formatter.write_str("graph node block operation was interrupted"),
516            Self::Truncated { minimum, actual } => write!(
517                formatter,
518                "graph node region is truncated: need {minimum} bytes, got {actual}"
519            ),
520            Self::InvalidHeader(detail) => {
521                write!(formatter, "graph node region header is invalid: {detail}")
522            }
523            Self::Layout(detail) => write!(formatter, "graph node layout is invalid: {detail}"),
524            Self::NodeIdOutOfRange {
525                node_id,
526                node_count,
527            } => write!(
528                formatter,
529                "graph node id {node_id} is outside dense row count {node_count}"
530            ),
531            Self::InvalidNode { node_id, detail } => {
532                write!(formatter, "graph node {node_id} is invalid: {detail}")
533            }
534            Self::ArithmeticOverflow => {
535                formatter.write_str("graph node block arithmetic overflowed")
536            }
537        }
538    }
539}
540
541impl std::error::Error for GraphNodeError {}
542
543/// Encodes node blocks in dense row-id order followed by a 128-byte trailer.
544pub fn encode_node_blocks(
545    build: GraphNodeBlockBuild<'_>,
546) -> Result<EncodedNodeBlocks, GraphNodeError> {
547    encode_node_blocks_controlled(build, RefinementPasses::default(), &mut |_| true)
548}
549
550/// Encodes node blocks with an authenticated additive refinement-pass record.
551pub fn encode_node_blocks_with_refinement_passes(
552    build: GraphNodeBlockBuild<'_>,
553    refinement_passes: RefinementPasses,
554) -> Result<EncodedNodeBlocks, GraphNodeError> {
555    encode_node_blocks_controlled(build, refinement_passes, &mut |_| true)
556}
557
558pub(crate) fn encode_node_blocks_controlled(
559    build: GraphNodeBlockBuild<'_>,
560    refinement_passes: RefinementPasses,
561    control: &mut impl FnMut(u64) -> bool,
562) -> Result<EncodedNodeBlocks, GraphNodeError> {
563    let node_count = u32::try_from(build.nodes.len())
564        .map_err(|_| GraphNodeError::Layout("node count exceeds u32".to_owned()))?;
565    let block_bytes = usize::try_from(build.layout.stride)
566        .map_err(|_| GraphNodeError::ArithmeticOverflow)?
567        .checked_mul(build.nodes.len())
568        .ok_or(GraphNodeError::ArithmeticOverflow)?;
569    let total_bytes = block_bytes
570        .checked_add(NODE_BLOCK_TRAILER_LEN)
571        .ok_or(GraphNodeError::ArithmeticOverflow)?;
572    let mut output = Vec::new();
573    output
574        .try_reserve_exact(total_bytes)
575        .map_err(|_| GraphNodeError::ArithmeticOverflow)?;
576    for (position, node) in build.nodes.iter().enumerate() {
577        if !control(1) {
578            return Err(GraphNodeError::Control);
579        }
580        let node_id = u32::try_from(position).map_err(|_| GraphNodeError::ArithmeticOverflow)?;
581        encode_node(
582            build.layout,
583            node_id,
584            node,
585            node_count,
586            &mut output,
587            control,
588        )?;
589    }
590    encode_trailer(
591        build.layout,
592        node_count,
593        refinement_passes,
594        &mut output,
595        control,
596    )?;
597    Ok(EncodedNodeBlocks {
598        bytes: output,
599        layout: build.layout,
600        node_count,
601    })
602}
603
604/// Validates and borrows a complete fixed-stride graph node-block region.
605pub fn decode_node_blocks(bytes: &[u8]) -> Result<GraphNodeBlocks<'_>, GraphNodeError> {
606    decode_node_blocks_controlled(bytes, &mut |_| true)
607}
608
609pub(crate) fn decode_node_blocks_controlled<'a>(
610    bytes: &'a [u8],
611    control: &mut impl FnMut(u64) -> bool,
612) -> Result<GraphNodeBlocks<'a>, GraphNodeError> {
613    if bytes.len() < NODE_BLOCK_TRAILER_LEN {
614        return Err(GraphNodeError::Truncated {
615            minimum: NODE_BLOCK_TRAILER_LEN,
616            actual: bytes.len(),
617        });
618    }
619    let trailer_start = bytes
620        .len()
621        .checked_sub(NODE_BLOCK_TRAILER_LEN)
622        .ok_or(GraphNodeError::ArithmeticOverflow)?;
623    let trailer = bytes
624        .get(trailer_start..)
625        .ok_or(GraphNodeError::Truncated {
626            minimum: NODE_BLOCK_TRAILER_LEN,
627            actual: bytes.len(),
628        })?;
629    if trailer.get(..8) != Some(NODE_BLOCK_MAGIC.as_slice()) {
630        return Err(GraphNodeError::InvalidHeader(
631            "magic does not match ZEGRNB01".to_owned(),
632        ));
633    }
634    let version = read_u16(trailer, 8, "version")?;
635    if version != NODE_BLOCK_VERSION {
636        return Err(GraphNodeError::InvalidHeader(format!(
637            "version {version}, expected {NODE_BLOCK_VERSION}"
638        )));
639    }
640    let refinement_bits = read_u16(trailer, 10, "refinement passes")?;
641    let refinement_passes = RefinementPasses::from_bits(refinement_bits).ok_or_else(|| {
642        GraphNodeError::InvalidHeader(format!(
643            "unknown refinement-pass bits {refinement_bits:#06x}"
644        ))
645    })?;
646    let dims = read_u32(trailer, 12, "logical dimensions")?;
647    let padded_dims = read_u32(trailer, 16, "padded dimensions")?;
648    let max_degree = *trailer.get(20).ok_or(GraphNodeError::Truncated {
649        minimum: 21,
650        actual: trailer.len(),
651    })?;
652    if trailer
653        .get(21..24)
654        .is_none_or(|reserved| reserved.iter().any(|byte| *byte != 0))
655    {
656        return Err(GraphNodeError::InvalidHeader(
657            "degree padding is non-zero".to_owned(),
658        ));
659    }
660    let stored_stride = read_u32(trailer, 24, "stride")?;
661    let node_count = read_u32(trailer, 28, "node count")?;
662    let stored_checksum = read_u64(trailer, 32, "region checksum")?;
663    let checksummed_end = trailer_start
664        .checked_add(32)
665        .ok_or(GraphNodeError::ArithmeticOverflow)?;
666    let checksummed = bytes
667        .get(..checksummed_end)
668        .ok_or(GraphNodeError::Truncated {
669            minimum: checksummed_end,
670            actual: bytes.len(),
671        })?;
672    let mut hasher = Xxh3::new();
673    for chunk in checksummed.chunks(256) {
674        if !control(chunk.len() as u64) {
675            return Err(GraphNodeError::Control);
676        }
677        hasher.update(chunk);
678    }
679    let actual_checksum = hasher.digest();
680    if actual_checksum != stored_checksum {
681        return Err(GraphNodeError::InvalidHeader(format!(
682            "xxh3-64 expected {stored_checksum:#018x}, computed {actual_checksum:#018x}"
683        )));
684    }
685    if trailer
686        .get(40..)
687        .is_none_or(|reserved| reserved.iter().any(|byte| *byte != 0))
688    {
689        return Err(GraphNodeError::InvalidHeader(
690            "reserved trailer bytes are non-zero".to_owned(),
691        ));
692    }
693    let layout = GraphNodeLayout::new(dims, padded_dims, max_degree)?;
694    if stored_stride != layout.stride {
695        return Err(GraphNodeError::InvalidHeader(format!(
696            "stored stride {stored_stride}, formula yields {}",
697            layout.stride
698        )));
699    }
700    let block_bytes = usize::try_from(layout.stride)
701        .map_err(|_| GraphNodeError::ArithmeticOverflow)?
702        .checked_mul(node_count as usize)
703        .ok_or(GraphNodeError::ArithmeticOverflow)?;
704    if trailer_start != block_bytes {
705        let minimum = block_bytes
706            .checked_add(NODE_BLOCK_TRAILER_LEN)
707            .ok_or(GraphNodeError::ArithmeticOverflow)?;
708        if bytes.len() < minimum {
709            return Err(GraphNodeError::Truncated {
710                minimum,
711                actual: bytes.len(),
712            });
713        }
714        return Err(GraphNodeError::InvalidHeader(format!(
715            "region length {}, expected {minimum}",
716            bytes.len()
717        )));
718    }
719    for node_id in 0..node_count {
720        if !control(1) {
721            return Err(GraphNodeError::Control);
722        }
723        let start = usize::try_from(layout.block_offset(node_id)?)
724            .map_err(|_| GraphNodeError::ArithmeticOverflow)?;
725        let end = start
726            .checked_add(layout.stride as usize)
727            .ok_or(GraphNodeError::ArithmeticOverflow)?;
728        let block = bytes.get(start..end).ok_or(GraphNodeError::Truncated {
729            minimum: end,
730            actual: bytes.len(),
731        })?;
732        validate_decoded_node(layout, node_id, node_count, block, control)?;
733    }
734    Ok(GraphNodeBlocks {
735        bytes,
736        layout,
737        node_count,
738        block_bytes,
739        refinement_passes,
740    })
741}
742
743fn validate_decoded_node(
744    layout: GraphNodeLayout,
745    node_id: u32,
746    node_count: u32,
747    block: &[u8],
748    control: &mut impl FnMut(u64) -> bool,
749) -> Result<(), GraphNodeError> {
750    let view = decode_block_view(layout, node_id, block)?;
751    validate_code_padding_controlled(layout, node_id, view.codes, control)?;
752    if view.flags & !ALLOWED_NODE_FLAGS != 0 {
753        return Err(GraphNodeError::InvalidNode {
754            node_id,
755            detail: format!("reserved flag bits are non-zero: {:#04x}", view.flags),
756        });
757    }
758    if view.reserved != [0, 0] {
759        return Err(GraphNodeError::InvalidNode {
760            node_id,
761            detail: "reserved bytes are non-zero".to_owned(),
762        });
763    }
764    for chunk in view.padding.chunks(256) {
765        if !control(chunk.len() as u64) {
766            return Err(GraphNodeError::Control);
767        }
768        if chunk.iter().any(|byte| *byte != 0) {
769            return Err(GraphNodeError::InvalidNode {
770                node_id,
771                detail: "cache-line padding is non-zero".to_owned(),
772            });
773        }
774    }
775    let degree = usize::from(view.degree);
776    if degree > usize::from(layout.max_degree) {
777        return Err(GraphNodeError::InvalidNode {
778            node_id,
779            detail: format!("degree {degree}, maximum {}", layout.max_degree),
780        });
781    }
782    for (slot, neighbor) in view.neighbors_padded().enumerate() {
783        if !control(1) {
784            return Err(GraphNodeError::Control);
785        }
786        if slot < degree {
787            if neighbor >= node_count {
788                return Err(GraphNodeError::InvalidNode {
789                    node_id,
790                    detail: format!(
791                        "active neighbour slot {slot} id {neighbor} is outside row count {node_count}"
792                    ),
793                });
794            }
795            if neighbor == node_id {
796                return Err(GraphNodeError::InvalidNode {
797                    node_id,
798                    detail: format!("active neighbour slot {slot} refers to its owner"),
799                });
800            }
801            for previous in view.neighbors_padded().take(slot) {
802                if !control(1) {
803                    return Err(GraphNodeError::Control);
804                }
805                if previous == neighbor {
806                    return Err(GraphNodeError::InvalidNode {
807                        node_id,
808                        detail: format!("active neighbour slot {slot} repeats id {neighbor}"),
809                    });
810                }
811            }
812        } else if neighbor != UNUSED_NEIGHBOR_ID {
813            return Err(GraphNodeError::InvalidNode {
814                node_id,
815                detail: format!("unused neighbour slot {slot} holds {neighbor}, expected u32::MAX"),
816            });
817        }
818    }
819    Ok(())
820}
821
822fn decode_block_view<'a>(
823    layout: GraphNodeLayout,
824    node_id: u32,
825    block: &'a [u8],
826) -> Result<GraphNodeBlock<'a>, GraphNodeError> {
827    if block.len() != layout.stride as usize {
828        return Err(GraphNodeError::InvalidNode {
829            node_id,
830            detail: format!("block bytes {}, expected {}", block.len(), layout.stride),
831        });
832    }
833    let code_end = layout.code_bytes();
834    let codes = block.get(..code_end).ok_or(GraphNodeError::Truncated {
835        minimum: code_end,
836        actual: block.len(),
837    })?;
838    let scale = read_f32(block, code_end, "scale")?;
839    let norm = read_f32(block, code_end.saturating_add(4), "normalized norm")?;
840    let correction = read_f32(block, code_end.saturating_add(8), "normalized correction")?;
841    if [scale, norm, correction]
842        .iter()
843        .any(|value| !value.is_finite())
844    {
845        return Err(GraphNodeError::InvalidNode {
846            node_id,
847            detail: "factor record contains a non-finite value".to_owned(),
848        });
849    }
850    let degree_offset = code_end.saturating_add(12);
851    let degree = *block.get(degree_offset).ok_or(GraphNodeError::Truncated {
852        minimum: degree_offset.saturating_add(1),
853        actual: block.len(),
854    })?;
855    let flags = *block
856        .get(degree_offset.saturating_add(1))
857        .ok_or(GraphNodeError::Truncated {
858            minimum: degree_offset.saturating_add(2),
859            actual: block.len(),
860        })?;
861    let reserved: [u8; 2] = block
862        .get(degree_offset.saturating_add(2)..degree_offset.saturating_add(4))
863        .ok_or(GraphNodeError::Truncated {
864            minimum: degree_offset.saturating_add(4),
865            actual: block.len(),
866        })?
867        .try_into()
868        .map_err(|_| GraphNodeError::InvalidNode {
869            node_id,
870            detail: "reserved field width is invalid".to_owned(),
871        })?;
872    let neighbors_start = code_end.saturating_add(NODE_METADATA_BYTES);
873    let neighbors_end = neighbors_start
874        .checked_add(usize::from(layout.max_degree).saturating_mul(4))
875        .ok_or(GraphNodeError::ArithmeticOverflow)?;
876    let neighbors = block
877        .get(neighbors_start..neighbors_end)
878        .ok_or(GraphNodeError::Truncated {
879            minimum: neighbors_end,
880            actual: block.len(),
881        })?;
882    let padding = block
883        .get(neighbors_end..)
884        .ok_or(GraphNodeError::Truncated {
885            minimum: neighbors_end,
886            actual: block.len(),
887        })?;
888    Ok(GraphNodeBlock {
889        codes,
890        factors: Bit4Factors::from_persisted(scale, norm, correction),
891        degree,
892        flags,
893        reserved,
894        neighbors,
895        padding,
896    })
897}
898
899fn read_u16(bytes: &[u8], offset: usize, field: &str) -> Result<u16, GraphNodeError> {
900    let end = offset
901        .checked_add(2)
902        .ok_or(GraphNodeError::ArithmeticOverflow)?;
903    let raw: [u8; 2] = bytes
904        .get(offset..end)
905        .ok_or(GraphNodeError::Truncated {
906            minimum: end,
907            actual: bytes.len(),
908        })?
909        .try_into()
910        .map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
911    Ok(u16::from_le_bytes(raw))
912}
913
914fn read_u32(bytes: &[u8], offset: usize, field: &str) -> Result<u32, GraphNodeError> {
915    let end = offset
916        .checked_add(4)
917        .ok_or(GraphNodeError::ArithmeticOverflow)?;
918    let raw: [u8; 4] = bytes
919        .get(offset..end)
920        .ok_or(GraphNodeError::Truncated {
921            minimum: end,
922            actual: bytes.len(),
923        })?
924        .try_into()
925        .map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
926    Ok(u32::from_le_bytes(raw))
927}
928
929fn read_f32(bytes: &[u8], offset: usize, field: &str) -> Result<f32, GraphNodeError> {
930    read_u32(bytes, offset, field).map(f32::from_bits)
931}
932
933fn read_u64(bytes: &[u8], offset: usize, field: &str) -> Result<u64, GraphNodeError> {
934    let end = offset
935        .checked_add(8)
936        .ok_or(GraphNodeError::ArithmeticOverflow)?;
937    let raw: [u8; 8] = bytes
938        .get(offset..end)
939        .ok_or(GraphNodeError::Truncated {
940            minimum: end,
941            actual: bytes.len(),
942        })?
943        .try_into()
944        .map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
945    Ok(u64::from_le_bytes(raw))
946}
947
948fn encode_node(
949    layout: GraphNodeLayout,
950    node_id: u32,
951    node: &GraphNodeBlockInput<'_>,
952    node_count: u32,
953    output: &mut Vec<u8>,
954    control: &mut impl FnMut(u64) -> bool,
955) -> Result<(), GraphNodeError> {
956    if node.codes.len() != layout.code_bytes() {
957        return Err(GraphNodeError::InvalidNode {
958            node_id,
959            detail: format!(
960                "code bytes {}, expected {}",
961                node.codes.len(),
962                layout.code_bytes()
963            ),
964        });
965    }
966    validate_code_padding_controlled(layout, node_id, node.codes, control)?;
967    if node
968        .factors
969        .persisted_fields()
970        .iter()
971        .any(|field| !field.is_finite())
972    {
973        return Err(GraphNodeError::InvalidNode {
974            node_id,
975            detail: "factor record contains a non-finite value".to_owned(),
976        });
977    }
978    if node.flags & !ALLOWED_NODE_FLAGS != 0 {
979        return Err(GraphNodeError::InvalidNode {
980            node_id,
981            detail: format!("reserved flag bits are non-zero: {:#04x}", node.flags),
982        });
983    }
984    if node.neighbors.len() > usize::from(layout.max_degree) {
985        return Err(GraphNodeError::InvalidNode {
986            node_id,
987            detail: format!(
988                "degree {}, maximum {}",
989                node.neighbors.len(),
990                layout.max_degree
991            ),
992        });
993    }
994    for (slot, neighbor) in node.neighbors.iter().copied().enumerate() {
995        if !control(1) {
996            return Err(GraphNodeError::Control);
997        }
998        if neighbor >= node_count {
999            return Err(GraphNodeError::InvalidNode {
1000                node_id,
1001                detail: format!("neighbour id {neighbor} is outside row count {node_count}"),
1002            });
1003        }
1004        if neighbor == node_id {
1005            return Err(GraphNodeError::InvalidNode {
1006                node_id,
1007                detail: format!("active neighbour slot {slot} refers to its owner"),
1008            });
1009        }
1010        for previous in node.neighbors.iter().take(slot) {
1011            if !control(1) {
1012                return Err(GraphNodeError::Control);
1013            }
1014            if *previous == neighbor {
1015                return Err(GraphNodeError::InvalidNode {
1016                    node_id,
1017                    detail: format!("active neighbour slot {slot} repeats id {neighbor}"),
1018                });
1019            }
1020        }
1021    }
1022
1023    let start = output.len();
1024    for chunk in node.codes.chunks(256) {
1025        if !control(chunk.len() as u64) {
1026            return Err(GraphNodeError::Control);
1027        }
1028        output.extend_from_slice(chunk);
1029    }
1030    for field in node.factors.persisted_fields() {
1031        output.extend_from_slice(&field.to_bits().to_le_bytes());
1032    }
1033    let degree = u8::try_from(node.neighbors.len()).map_err(|_| GraphNodeError::InvalidNode {
1034        node_id,
1035        detail: "degree exceeds u8".to_owned(),
1036    })?;
1037    output.push(degree);
1038    output.push(node.flags);
1039    output.extend_from_slice(&0_u16.to_le_bytes());
1040    for neighbor in node.neighbors {
1041        if !control(1) {
1042            return Err(GraphNodeError::Control);
1043        }
1044        output.extend_from_slice(&neighbor.to_le_bytes());
1045    }
1046    for _ in node.neighbors.len()..usize::from(layout.max_degree) {
1047        output.extend_from_slice(&UNUSED_NEIGHBOR_ID.to_le_bytes());
1048    }
1049    let end = start
1050        .checked_add(layout.stride as usize)
1051        .ok_or(GraphNodeError::ArithmeticOverflow)?;
1052    output.resize(end, 0);
1053    Ok(())
1054}
1055
1056fn validate_code_padding_controlled(
1057    layout: GraphNodeLayout,
1058    node_id: u32,
1059    codes: &[u8],
1060    control: &mut impl FnMut(u64) -> bool,
1061) -> Result<(), GraphNodeError> {
1062    let logical_bytes = (layout.dims as usize).div_ceil(2);
1063    if !layout.dims.is_multiple_of(2)
1064        && codes
1065            .get(logical_bytes.saturating_sub(1))
1066            .is_some_and(|byte| byte & 0x0f != 0)
1067    {
1068        return Err(GraphNodeError::InvalidNode {
1069            node_id,
1070            detail: "unused low nibble is non-zero".to_owned(),
1071        });
1072    }
1073    if let Some(padding) = codes.get(logical_bytes..) {
1074        for chunk in padding.chunks(256) {
1075            if !control(chunk.len() as u64) {
1076                return Err(GraphNodeError::Control);
1077            }
1078            if chunk.iter().any(|byte| *byte != 0) {
1079                return Err(GraphNodeError::InvalidNode {
1080                    node_id,
1081                    detail: "padded Bit4 dimensions are non-zero".to_owned(),
1082                });
1083            }
1084        }
1085    }
1086    Ok(())
1087}
1088
1089fn encode_trailer(
1090    layout: GraphNodeLayout,
1091    node_count: u32,
1092    refinement_passes: RefinementPasses,
1093    output: &mut Vec<u8>,
1094    control: &mut impl FnMut(u64) -> bool,
1095) -> Result<(), GraphNodeError> {
1096    let start = output.len();
1097    output.extend_from_slice(&NODE_BLOCK_MAGIC);
1098    output.extend_from_slice(&NODE_BLOCK_VERSION.to_le_bytes());
1099    output.extend_from_slice(&refinement_passes.bits().to_le_bytes());
1100    output.extend_from_slice(&layout.dims.to_le_bytes());
1101    output.extend_from_slice(&layout.padded_dims.to_le_bytes());
1102    output.push(layout.max_degree);
1103    output.extend_from_slice(&[0_u8; 3]);
1104    output.extend_from_slice(&layout.stride.to_le_bytes());
1105    output.extend_from_slice(&node_count.to_le_bytes());
1106    let mut hasher = Xxh3::new();
1107    for chunk in output.chunks(256) {
1108        if !control(chunk.len() as u64) {
1109            return Err(GraphNodeError::Control);
1110        }
1111        hasher.update(chunk);
1112    }
1113    let checksum = hasher.digest();
1114    output.extend_from_slice(&checksum.to_le_bytes());
1115    output.resize(start.saturating_add(NODE_BLOCK_TRAILER_LEN), 0);
1116    Ok(())
1117}
1118
1119fn round_up_cache_line(value: usize) -> Result<usize, GraphNodeError> {
1120    value
1121        .checked_add(CACHE_LINE_BYTES.saturating_sub(1))
1122        .map(|rounded| rounded / CACHE_LINE_BYTES * CACHE_LINE_BYTES)
1123        .ok_or(GraphNodeError::ArithmeticOverflow)
1124}
1125
1126#[cfg(test)]
1127#[allow(clippy::expect_used)]
1128mod tests {
1129    use super::{
1130        GraphNodeBlockBuild, GraphNodeBlockInput, GraphNodeError, GraphNodeLayout,
1131        decode_node_blocks, encode_node_blocks,
1132    };
1133    use crate::quant::Bit4Factors;
1134
1135    #[test]
1136    fn cached_descriptor_rejects_a_different_mapping_length() {
1137        let layout = GraphNodeLayout::new(128, 128, 0).expect("cache fixture layout");
1138        let codes = [0_u8; 64];
1139        let node = GraphNodeBlockInput {
1140            codes: &codes,
1141            factors: Bit4Factors::from_persisted(1.0, 0.0, 0.0),
1142            flags: 1,
1143            neighbors: &[],
1144        };
1145        let encoded = encode_node_blocks(GraphNodeBlockBuild {
1146            layout,
1147            nodes: &[node],
1148        })
1149        .expect("cache fixture encodes");
1150        let graph = decode_node_blocks(encoded.as_bytes()).expect("cache fixture decodes");
1151        let descriptor = graph.validated_descriptor();
1152
1153        let shortened = encoded
1154            .as_bytes()
1155            .get(..encoded.as_bytes().len().saturating_sub(1))
1156            .expect("fixture has a trailer byte to remove");
1157        assert!(matches!(
1158            descriptor.bind(shortened),
1159            Err(GraphNodeError::InvalidHeader(detail))
1160                if detail.contains("cached graph region length")
1161        ));
1162        let rebound = descriptor
1163            .bind(encoded.as_bytes())
1164            .expect("identical mapping length binds");
1165        assert_eq!(rebound.block(0).expect("node zero").flags(), 1);
1166
1167        let outside = encoded
1168            .block_offset(1)
1169            .expect_err("one-node graph rejects row one");
1170        assert_eq!(
1171            outside.to_string(),
1172            "graph node id 1 is outside dense row count 1"
1173        );
1174        assert_eq!(
1175            GraphNodeError::ArithmeticOverflow.to_string(),
1176            "graph node block arithmetic overflowed"
1177        );
1178    }
1179}