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(crate) fn resident_bytes(&self) -> usize {
467 self.bytes.capacity()
468 }
469
470 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#[derive(Clone, Debug, Eq, PartialEq)]
484pub enum GraphNodeError {
485 Control,
487 Truncated {
489 minimum: usize,
491 actual: usize,
493 },
494 InvalidHeader(String),
496 Layout(String),
498 NodeIdOutOfRange {
500 node_id: u32,
502 node_count: u32,
504 },
505 InvalidNode {
507 node_id: u32,
509 detail: String,
511 },
512 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
547pub fn encode_node_blocks(
549 build: GraphNodeBlockBuild<'_>,
550) -> Result<EncodedNodeBlocks, GraphNodeError> {
551 encode_node_blocks_controlled(build, RefinementPasses::default(), &mut |_| true)
552}
553
554pub 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
608pub 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}