1use crate::graph::refine::RefinementPasses;
4use crate::kernels::Bit4Row;
5use crate::quant::Bit4Factors;
6use xxhash_rust::xxh3::xxh3_64;
7
8pub const CACHE_LINE_BYTES: usize = 128;
10pub const NODE_BLOCK_TRAILER_LEN: usize = CACHE_LINE_BYTES;
12pub 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#[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 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 #[must_use]
69 pub const fn dims(self) -> u32 {
70 self.dims
71 }
72
73 #[must_use]
75 pub const fn padded_dims(self) -> u32 {
76 self.padded_dims
77 }
78
79 #[must_use]
81 pub const fn max_degree(self) -> u8 {
82 self.max_degree
83 }
84
85 #[must_use]
87 pub const fn code_bytes(self) -> usize {
88 (self.padded_dims as usize).div_ceil(2)
89 }
90
91 #[must_use]
93 pub const fn stride(self) -> u32 {
94 self.stride
95 }
96
97 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#[derive(Clone, Copy, Debug)]
107pub struct GraphNodeBlockInput<'a> {
108 pub codes: &'a [u8],
110 pub factors: Bit4Factors,
112 pub flags: u8,
114 pub neighbors: &'a [u32],
116}
117
118#[derive(Clone, Copy, Debug)]
120pub struct GraphNodeBlockBuild<'a> {
121 pub layout: GraphNodeLayout,
123 pub nodes: &'a [GraphNodeBlockInput<'a>],
125}
126
127#[derive(Clone, Debug, Eq, PartialEq)]
129pub struct EncodedNodeBlocks {
130 bytes: Vec<u8>,
131 layout: GraphNodeLayout,
132 node_count: u32,
133}
134
135#[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#[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#[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 #[must_use]
172 pub const fn layout(&self) -> GraphNodeLayout {
173 self.layout
174 }
175
176 #[must_use]
178 pub const fn node_count(&self) -> u32 {
179 self.node_count
180 }
181
182 #[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 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 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#[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 #[must_use]
388 pub const fn codes(&self) -> &[u8] {
389 self.codes
390 }
391
392 #[must_use]
394 pub const fn factors(&self) -> Bit4Factors {
395 self.factors
396 }
397
398 #[must_use]
400 pub const fn degree(&self) -> u8 {
401 self.degree
402 }
403
404 #[must_use]
406 pub const fn flags(&self) -> u8 {
407 self.flags
408 }
409
410 #[must_use]
412 pub const fn reserved_bytes(&self) -> [u8; 2] {
413 self.reserved
414 }
415
416 #[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 #[must_use]
426 pub const fn padding_bytes(&self) -> &[u8] {
427 self.padding
428 }
429}
430
431#[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 #[must_use]
456 pub fn as_bytes(&self) -> &[u8] {
457 &self.bytes
458 }
459
460 #[must_use]
462 pub fn into_bytes(self) -> Vec<u8> {
463 self.bytes
464 }
465
466 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#[derive(Clone, Debug, Eq, PartialEq)]
480pub enum GraphNodeError {
481 Truncated {
483 minimum: usize,
485 actual: usize,
487 },
488 InvalidHeader(String),
490 Layout(String),
492 NodeIdOutOfRange {
494 node_id: u32,
496 node_count: u32,
498 },
499 InvalidNode {
501 node_id: u32,
503 detail: String,
505 },
506 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
540pub fn encode_node_blocks(
542 build: GraphNodeBlockBuild<'_>,
543) -> Result<EncodedNodeBlocks, GraphNodeError> {
544 encode_node_blocks_with_refinement_passes(build, RefinementPasses::default())
545}
546
547pub 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
574pub 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}