1use crate::graph::refine::RefinementPasses;
4use crate::kernels::Bit4Row;
5use crate::quant::Bit4Factors;
6use xxhash_rust::xxh3::Xxh3;
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 Control,
483 Truncated {
485 minimum: usize,
487 actual: usize,
489 },
490 InvalidHeader(String),
492 Layout(String),
494 NodeIdOutOfRange {
496 node_id: u32,
498 node_count: u32,
500 },
501 InvalidNode {
503 node_id: u32,
505 detail: String,
507 },
508 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
543pub fn encode_node_blocks(
545 build: GraphNodeBlockBuild<'_>,
546) -> Result<EncodedNodeBlocks, GraphNodeError> {
547 encode_node_blocks_controlled(build, RefinementPasses::default(), &mut |_| true)
548}
549
550pub 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
604pub 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}