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