use crate::graph::refine::RefinementPasses;
use crate::kernels::Bit4Row;
use crate::quant::Bit4Factors;
use xxhash_rust::xxh3::xxh3_64;
pub const CACHE_LINE_BYTES: usize = 128;
pub const NODE_BLOCK_TRAILER_LEN: usize = CACHE_LINE_BYTES;
pub const UNUSED_NEIGHBOR_ID: u32 = u32::MAX;
const NODE_BLOCK_MAGIC: [u8; 8] = *b"ZEGRNB01";
const NODE_BLOCK_VERSION: u16 = 1;
const NODE_METADATA_BYTES: usize = 16;
const ALLOWED_NODE_FLAGS: u8 = 0b0000_0011;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct GraphNodeLayout {
dims: u32,
padded_dims: u32,
max_degree: u8,
stride: u32,
}
impl GraphNodeLayout {
pub fn new(dims: u32, padded_dims: u32, max_degree: u8) -> Result<Self, GraphNodeError> {
if dims == 0 {
return Err(GraphNodeError::Layout(
"logical dimensions must not be zero".to_owned(),
));
}
if padded_dims < dims {
return Err(GraphNodeError::Layout(format!(
"padded dimensions {padded_dims} are below logical dimensions {dims}"
)));
}
if !padded_dims.is_multiple_of(CACHE_LINE_BYTES as u32) {
return Err(GraphNodeError::Layout(format!(
"padded dimensions {padded_dims} are not a multiple of {CACHE_LINE_BYTES}"
)));
}
let code_bytes = usize::try_from(padded_dims)
.map_err(|_| GraphNodeError::Layout("padded dimensions exceed usize".to_owned()))?
.div_ceil(2);
let neighbor_bytes = usize::from(max_degree)
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| GraphNodeError::Layout("neighbour bytes overflow".to_owned()))?;
let unaligned = code_bytes
.checked_add(NODE_METADATA_BYTES)
.and_then(|value| value.checked_add(neighbor_bytes))
.ok_or_else(|| GraphNodeError::Layout("node block size overflow".to_owned()))?;
let stride = round_up_cache_line(unaligned)?;
Ok(Self {
dims,
padded_dims,
max_degree,
stride: u32::try_from(stride)
.map_err(|_| GraphNodeError::Layout("node stride exceeds u32".to_owned()))?,
})
}
#[must_use]
pub const fn dims(self) -> u32 {
self.dims
}
#[must_use]
pub const fn padded_dims(self) -> u32 {
self.padded_dims
}
#[must_use]
pub const fn max_degree(self) -> u8 {
self.max_degree
}
#[must_use]
pub const fn code_bytes(self) -> usize {
(self.padded_dims as usize).div_ceil(2)
}
#[must_use]
pub const fn stride(self) -> u32 {
self.stride
}
pub fn block_offset(self, node_id: u32) -> Result<u32, GraphNodeError> {
node_id
.checked_mul(self.stride)
.ok_or(GraphNodeError::ArithmeticOverflow)
}
}
#[derive(Clone, Copy, Debug)]
pub struct GraphNodeBlockInput<'a> {
pub codes: &'a [u8],
pub factors: Bit4Factors,
pub flags: u8,
pub neighbors: &'a [u32],
}
#[derive(Clone, Copy, Debug)]
pub struct GraphNodeBlockBuild<'a> {
pub layout: GraphNodeLayout,
pub nodes: &'a [GraphNodeBlockInput<'a>],
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct EncodedNodeBlocks {
bytes: Vec<u8>,
layout: GraphNodeLayout,
node_count: u32,
}
#[derive(Clone, Copy, Debug)]
pub struct GraphNodeBlocks<'a> {
bytes: &'a [u8],
layout: GraphNodeLayout,
node_count: u32,
block_bytes: usize,
refinement_passes: RefinementPasses,
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct ValidatedGraphNodeBlocks {
layout: GraphNodeLayout,
node_count: u32,
block_bytes: usize,
region_bytes: usize,
refinement_passes: RefinementPasses,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) struct CheckedNodeId {
raw: u32,
offset: usize,
}
impl CheckedNodeId {
#[must_use]
pub(super) const fn raw(self) -> u32 {
self.raw
}
}
impl GraphNodeBlocks<'_> {
#[must_use]
pub const fn layout(&self) -> GraphNodeLayout {
self.layout
}
#[must_use]
pub const fn node_count(&self) -> u32 {
self.node_count
}
#[must_use]
pub const fn refinement_passes(&self) -> RefinementPasses {
self.refinement_passes
}
pub(crate) const fn validated_descriptor(&self) -> ValidatedGraphNodeBlocks {
ValidatedGraphNodeBlocks {
layout: self.layout,
node_count: self.node_count,
block_bytes: self.block_bytes,
region_bytes: self.bytes.len(),
refinement_passes: self.refinement_passes,
}
}
pub fn block(&self, node_id: u32) -> Result<GraphNodeBlock<'_>, GraphNodeError> {
self.block_checked(self.checked_node_id(node_id)?)
}
pub(super) fn checked_node_id(&self, node_id: u32) -> Result<CheckedNodeId, GraphNodeError> {
if node_id >= self.node_count {
return Err(GraphNodeError::NodeIdOutOfRange {
node_id,
node_count: self.node_count,
});
}
let offset = usize::try_from(self.layout.block_offset(node_id)?)
.map_err(|_| GraphNodeError::ArithmeticOverflow)?;
Ok(CheckedNodeId {
raw: node_id,
offset,
})
}
pub(super) fn block_checked(
&self,
node_id: CheckedNodeId,
) -> Result<GraphNodeBlock<'_>, GraphNodeError> {
let end = node_id
.offset
.checked_add(self.layout.stride as usize)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let block = self
.bytes
.get(node_id.offset..end.min(self.block_bytes))
.ok_or_else(|| GraphNodeError::InvalidNode {
node_id: node_id.raw,
detail: "validated block range is unavailable".to_owned(),
})?;
decode_block_view(self.layout, node_id.raw, block)
}
pub(super) fn score_row_checked(
&self,
node_id: CheckedNodeId,
) -> Result<(Bit4Row<'_>, Bit4Factors), GraphNodeError> {
let code_bytes = self.layout.code_bytes();
let row_bytes = code_bytes
.checked_add(12)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let end = node_id
.offset
.checked_add(row_bytes)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let score_row =
self.bytes
.get(node_id.offset..end)
.ok_or_else(|| GraphNodeError::InvalidNode {
node_id: node_id.raw,
detail: "validated score row is unavailable".to_owned(),
})?;
let Some((codes, factor_bytes)) = score_row.split_at_checked(code_bytes) else {
return Err(GraphNodeError::InvalidNode {
node_id: node_id.raw,
detail: "validated score row width is invalid".to_owned(),
});
};
let factors: [u8; 12] =
factor_bytes
.try_into()
.map_err(|_| GraphNodeError::InvalidNode {
node_id: node_id.raw,
detail: "validated factor width is invalid".to_owned(),
})?;
let [s0, s1, s2, s3, n0, n1, n2, n3, c0, c1, c2, c3] = factors;
Ok((
Bit4Row::from_validated_bytes(codes),
Bit4Factors::from_persisted(
f32::from_le_bytes([s0, s1, s2, s3]),
f32::from_le_bytes([n0, n1, n2, n3]),
f32::from_le_bytes([c0, c1, c2, c3]),
),
))
}
pub(super) fn adjacency_checked(
&self,
node_id: CheckedNodeId,
) -> Result<(u8, NeighborIds<'_>), GraphNodeError> {
let degree_offset = self
.layout
.code_bytes()
.checked_add(12)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let adjacency_bytes = usize::from(self.layout.max_degree)
.checked_mul(std::mem::size_of::<u32>())
.and_then(|value| value.checked_add(4))
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let start = node_id
.offset
.checked_add(degree_offset)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let end = start
.checked_add(adjacency_bytes)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let adjacency = self
.bytes
.get(start..end)
.ok_or_else(|| GraphNodeError::InvalidNode {
node_id: node_id.raw,
detail: "validated adjacency is unavailable".to_owned(),
})?;
let [degree, _flags, _reserved0, _reserved1, neighbors @ ..] = adjacency else {
return Err(GraphNodeError::InvalidNode {
node_id: node_id.raw,
detail: "validated adjacency width is invalid".to_owned(),
});
};
Ok((
*degree,
NeighborIds {
chunks: neighbors.chunks_exact(std::mem::size_of::<u32>()),
},
))
}
pub(super) fn prefetch_line0(&self, node_id: CheckedNodeId) {
if let Some(bytes) = self.bytes.get(node_id.offset..) {
prefetch_address(bytes.as_ptr());
}
}
pub(super) fn prefetch_head_block(&self, node_id: CheckedNodeId) {
let Some(bytes) = self.bytes.get(node_id.offset..) else {
return;
};
prefetch_address(bytes.as_ptr());
if self.layout.stride as usize > CACHE_LINE_BYTES
&& let Some(second_line) = bytes.get(CACHE_LINE_BYTES..)
{
prefetch_address(second_line.as_ptr());
}
}
}
impl ValidatedGraphNodeBlocks {
pub(crate) fn bind<'a>(self, bytes: &'a [u8]) -> Result<GraphNodeBlocks<'a>, GraphNodeError> {
if bytes.len() != self.region_bytes {
return Err(GraphNodeError::InvalidHeader(format!(
"cached graph region length {}, current length {}",
self.region_bytes,
bytes.len()
)));
}
Ok(GraphNodeBlocks {
bytes,
layout: self.layout,
node_count: self.node_count,
block_bytes: self.block_bytes,
refinement_passes: self.refinement_passes,
})
}
}
#[inline]
fn prefetch_address(address: *const u8) {
#[cfg(target_arch = "aarch64")]
unsafe {
std::arch::asm!(
"prfm pldl1keep, [{address}]",
address = in(reg) address,
options(readonly, nostack)
);
}
#[cfg(not(target_arch = "aarch64"))]
let _ = address;
}
#[derive(Clone, Copy, Debug)]
pub struct GraphNodeBlock<'a> {
codes: &'a [u8],
factors: Bit4Factors,
degree: u8,
flags: u8,
reserved: [u8; 2],
neighbors: &'a [u8],
padding: &'a [u8],
}
impl GraphNodeBlock<'_> {
#[must_use]
pub const fn codes(&self) -> &[u8] {
self.codes
}
#[must_use]
pub const fn factors(&self) -> Bit4Factors {
self.factors
}
#[must_use]
pub const fn degree(&self) -> u8 {
self.degree
}
#[must_use]
pub const fn flags(&self) -> u8 {
self.flags
}
#[must_use]
pub const fn reserved_bytes(&self) -> [u8; 2] {
self.reserved
}
#[must_use]
pub fn neighbors_padded(&self) -> NeighborIds<'_> {
NeighborIds {
chunks: self.neighbors.chunks_exact(std::mem::size_of::<u32>()),
}
}
#[must_use]
pub const fn padding_bytes(&self) -> &[u8] {
self.padding
}
}
#[derive(Clone, Debug)]
pub struct NeighborIds<'a> {
chunks: std::slice::ChunksExact<'a, u8>,
}
impl Iterator for NeighborIds<'_> {
type Item = u32;
fn next(&mut self) -> Option<Self::Item> {
let raw: [u8; 4] = self.chunks.next()?.try_into().ok()?;
Some(u32::from_le_bytes(raw))
}
fn size_hint(&self) -> (usize, Option<usize>) {
let length = self.chunks.len();
(length, Some(length))
}
}
impl ExactSizeIterator for NeighborIds<'_> {}
impl EncodedNodeBlocks {
#[must_use]
pub fn as_bytes(&self) -> &[u8] {
&self.bytes
}
#[must_use]
pub fn into_bytes(self) -> Vec<u8> {
self.bytes
}
pub fn block_offset(&self, node_id: u32) -> Result<u32, GraphNodeError> {
if node_id >= self.node_count {
return Err(GraphNodeError::NodeIdOutOfRange {
node_id,
node_count: self.node_count,
});
}
self.layout.block_offset(node_id)
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum GraphNodeError {
Truncated {
minimum: usize,
actual: usize,
},
InvalidHeader(String),
Layout(String),
NodeIdOutOfRange {
node_id: u32,
node_count: u32,
},
InvalidNode {
node_id: u32,
detail: String,
},
ArithmeticOverflow,
}
impl std::fmt::Display for GraphNodeError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Truncated { minimum, actual } => write!(
formatter,
"graph node region is truncated: need {minimum} bytes, got {actual}"
),
Self::InvalidHeader(detail) => {
write!(formatter, "graph node region header is invalid: {detail}")
}
Self::Layout(detail) => write!(formatter, "graph node layout is invalid: {detail}"),
Self::NodeIdOutOfRange {
node_id,
node_count,
} => write!(
formatter,
"graph node id {node_id} is outside dense row count {node_count}"
),
Self::InvalidNode { node_id, detail } => {
write!(formatter, "graph node {node_id} is invalid: {detail}")
}
Self::ArithmeticOverflow => {
formatter.write_str("graph node block arithmetic overflowed")
}
}
}
}
impl std::error::Error for GraphNodeError {}
pub fn encode_node_blocks(
build: GraphNodeBlockBuild<'_>,
) -> Result<EncodedNodeBlocks, GraphNodeError> {
encode_node_blocks_with_refinement_passes(build, RefinementPasses::default())
}
pub fn encode_node_blocks_with_refinement_passes(
build: GraphNodeBlockBuild<'_>,
refinement_passes: RefinementPasses,
) -> Result<EncodedNodeBlocks, GraphNodeError> {
let node_count = u32::try_from(build.nodes.len())
.map_err(|_| GraphNodeError::Layout("node count exceeds u32".to_owned()))?;
let block_bytes = usize::try_from(build.layout.stride)
.map_err(|_| GraphNodeError::ArithmeticOverflow)?
.checked_mul(build.nodes.len())
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let total_bytes = block_bytes
.checked_add(NODE_BLOCK_TRAILER_LEN)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let mut output = Vec::with_capacity(total_bytes);
for (position, node) in build.nodes.iter().enumerate() {
let node_id = u32::try_from(position).map_err(|_| GraphNodeError::ArithmeticOverflow)?;
encode_node(build.layout, node_id, node, node_count, &mut output)?;
}
encode_trailer(build.layout, node_count, refinement_passes, &mut output);
Ok(EncodedNodeBlocks {
bytes: output,
layout: build.layout,
node_count,
})
}
pub fn decode_node_blocks(bytes: &[u8]) -> Result<GraphNodeBlocks<'_>, GraphNodeError> {
if bytes.len() < NODE_BLOCK_TRAILER_LEN {
return Err(GraphNodeError::Truncated {
minimum: NODE_BLOCK_TRAILER_LEN,
actual: bytes.len(),
});
}
let trailer_start = bytes
.len()
.checked_sub(NODE_BLOCK_TRAILER_LEN)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let trailer = bytes
.get(trailer_start..)
.ok_or(GraphNodeError::Truncated {
minimum: NODE_BLOCK_TRAILER_LEN,
actual: bytes.len(),
})?;
if trailer.get(..8) != Some(NODE_BLOCK_MAGIC.as_slice()) {
return Err(GraphNodeError::InvalidHeader(
"magic does not match ZEGRNB01".to_owned(),
));
}
let version = read_u16(trailer, 8, "version")?;
if version != NODE_BLOCK_VERSION {
return Err(GraphNodeError::InvalidHeader(format!(
"version {version}, expected {NODE_BLOCK_VERSION}"
)));
}
let refinement_bits = read_u16(trailer, 10, "refinement passes")?;
let refinement_passes = RefinementPasses::from_bits(refinement_bits).ok_or_else(|| {
GraphNodeError::InvalidHeader(format!(
"unknown refinement-pass bits {refinement_bits:#06x}"
))
})?;
let dims = read_u32(trailer, 12, "logical dimensions")?;
let padded_dims = read_u32(trailer, 16, "padded dimensions")?;
let max_degree = *trailer.get(20).ok_or(GraphNodeError::Truncated {
minimum: 21,
actual: trailer.len(),
})?;
if trailer
.get(21..24)
.is_none_or(|reserved| reserved.iter().any(|byte| *byte != 0))
{
return Err(GraphNodeError::InvalidHeader(
"degree padding is non-zero".to_owned(),
));
}
let stored_stride = read_u32(trailer, 24, "stride")?;
let node_count = read_u32(trailer, 28, "node count")?;
let stored_checksum = read_u64(trailer, 32, "region checksum")?;
let checksummed_end = trailer_start
.checked_add(32)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let checksummed = bytes
.get(..checksummed_end)
.ok_or(GraphNodeError::Truncated {
minimum: checksummed_end,
actual: bytes.len(),
})?;
let actual_checksum = xxh3_64(checksummed);
if actual_checksum != stored_checksum {
return Err(GraphNodeError::InvalidHeader(format!(
"xxh3-64 expected {stored_checksum:#018x}, computed {actual_checksum:#018x}"
)));
}
if trailer
.get(40..)
.is_none_or(|reserved| reserved.iter().any(|byte| *byte != 0))
{
return Err(GraphNodeError::InvalidHeader(
"reserved trailer bytes are non-zero".to_owned(),
));
}
let layout = GraphNodeLayout::new(dims, padded_dims, max_degree)?;
if stored_stride != layout.stride {
return Err(GraphNodeError::InvalidHeader(format!(
"stored stride {stored_stride}, formula yields {}",
layout.stride
)));
}
let block_bytes = usize::try_from(layout.stride)
.map_err(|_| GraphNodeError::ArithmeticOverflow)?
.checked_mul(node_count as usize)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
if trailer_start != block_bytes {
let minimum = block_bytes
.checked_add(NODE_BLOCK_TRAILER_LEN)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
if bytes.len() < minimum {
return Err(GraphNodeError::Truncated {
minimum,
actual: bytes.len(),
});
}
return Err(GraphNodeError::InvalidHeader(format!(
"region length {}, expected {minimum}",
bytes.len()
)));
}
for node_id in 0..node_count {
let start = usize::try_from(layout.block_offset(node_id)?)
.map_err(|_| GraphNodeError::ArithmeticOverflow)?;
let end = start
.checked_add(layout.stride as usize)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let block = bytes.get(start..end).ok_or(GraphNodeError::Truncated {
minimum: end,
actual: bytes.len(),
})?;
validate_decoded_node(layout, node_id, node_count, block)?;
}
Ok(GraphNodeBlocks {
bytes,
layout,
node_count,
block_bytes,
refinement_passes,
})
}
fn validate_decoded_node(
layout: GraphNodeLayout,
node_id: u32,
node_count: u32,
block: &[u8],
) -> Result<(), GraphNodeError> {
let view = decode_block_view(layout, node_id, block)?;
validate_code_padding(layout, node_id, view.codes)?;
if view.flags & !ALLOWED_NODE_FLAGS != 0 {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!("reserved flag bits are non-zero: {:#04x}", view.flags),
});
}
if view.reserved != [0, 0] {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: "reserved bytes are non-zero".to_owned(),
});
}
if view.padding.iter().any(|byte| *byte != 0) {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: "cache-line padding is non-zero".to_owned(),
});
}
let degree = usize::from(view.degree);
if degree > usize::from(layout.max_degree) {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!("degree {degree}, maximum {}", layout.max_degree),
});
}
for (slot, neighbor) in view.neighbors_padded().enumerate() {
if slot < degree {
if neighbor >= node_count {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!(
"active neighbour slot {slot} id {neighbor} is outside row count {node_count}"
),
});
}
} else if neighbor != UNUSED_NEIGHBOR_ID {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!("unused neighbour slot {slot} holds {neighbor}, expected u32::MAX"),
});
}
}
Ok(())
}
fn decode_block_view<'a>(
layout: GraphNodeLayout,
node_id: u32,
block: &'a [u8],
) -> Result<GraphNodeBlock<'a>, GraphNodeError> {
if block.len() != layout.stride as usize {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!("block bytes {}, expected {}", block.len(), layout.stride),
});
}
let code_end = layout.code_bytes();
let codes = block.get(..code_end).ok_or(GraphNodeError::Truncated {
minimum: code_end,
actual: block.len(),
})?;
let scale = read_f32(block, code_end, "scale")?;
let norm = read_f32(block, code_end.saturating_add(4), "normalized norm")?;
let correction = read_f32(block, code_end.saturating_add(8), "normalized correction")?;
if [scale, norm, correction]
.iter()
.any(|value| !value.is_finite())
{
return Err(GraphNodeError::InvalidNode {
node_id,
detail: "factor record contains a non-finite value".to_owned(),
});
}
let degree_offset = code_end.saturating_add(12);
let degree = *block.get(degree_offset).ok_or(GraphNodeError::Truncated {
minimum: degree_offset.saturating_add(1),
actual: block.len(),
})?;
let flags = *block
.get(degree_offset.saturating_add(1))
.ok_or(GraphNodeError::Truncated {
minimum: degree_offset.saturating_add(2),
actual: block.len(),
})?;
let reserved: [u8; 2] = block
.get(degree_offset.saturating_add(2)..degree_offset.saturating_add(4))
.ok_or(GraphNodeError::Truncated {
minimum: degree_offset.saturating_add(4),
actual: block.len(),
})?
.try_into()
.map_err(|_| GraphNodeError::InvalidNode {
node_id,
detail: "reserved field width is invalid".to_owned(),
})?;
let neighbors_start = code_end.saturating_add(NODE_METADATA_BYTES);
let neighbors_end = neighbors_start
.checked_add(usize::from(layout.max_degree).saturating_mul(4))
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let neighbors = block
.get(neighbors_start..neighbors_end)
.ok_or(GraphNodeError::Truncated {
minimum: neighbors_end,
actual: block.len(),
})?;
let padding = block
.get(neighbors_end..)
.ok_or(GraphNodeError::Truncated {
minimum: neighbors_end,
actual: block.len(),
})?;
Ok(GraphNodeBlock {
codes,
factors: Bit4Factors::from_persisted(scale, norm, correction),
degree,
flags,
reserved,
neighbors,
padding,
})
}
fn read_u16(bytes: &[u8], offset: usize, field: &str) -> Result<u16, GraphNodeError> {
let end = offset
.checked_add(2)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let raw: [u8; 2] = bytes
.get(offset..end)
.ok_or(GraphNodeError::Truncated {
minimum: end,
actual: bytes.len(),
})?
.try_into()
.map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
Ok(u16::from_le_bytes(raw))
}
fn read_u32(bytes: &[u8], offset: usize, field: &str) -> Result<u32, GraphNodeError> {
let end = offset
.checked_add(4)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let raw: [u8; 4] = bytes
.get(offset..end)
.ok_or(GraphNodeError::Truncated {
minimum: end,
actual: bytes.len(),
})?
.try_into()
.map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
Ok(u32::from_le_bytes(raw))
}
fn read_f32(bytes: &[u8], offset: usize, field: &str) -> Result<f32, GraphNodeError> {
read_u32(bytes, offset, field).map(f32::from_bits)
}
fn read_u64(bytes: &[u8], offset: usize, field: &str) -> Result<u64, GraphNodeError> {
let end = offset
.checked_add(8)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
let raw: [u8; 8] = bytes
.get(offset..end)
.ok_or(GraphNodeError::Truncated {
minimum: end,
actual: bytes.len(),
})?
.try_into()
.map_err(|_| GraphNodeError::InvalidHeader(format!("invalid {field} width")))?;
Ok(u64::from_le_bytes(raw))
}
fn encode_node(
layout: GraphNodeLayout,
node_id: u32,
node: &GraphNodeBlockInput<'_>,
node_count: u32,
output: &mut Vec<u8>,
) -> Result<(), GraphNodeError> {
if node.codes.len() != layout.code_bytes() {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!(
"code bytes {}, expected {}",
node.codes.len(),
layout.code_bytes()
),
});
}
validate_code_padding(layout, node_id, node.codes)?;
if node
.factors
.persisted_fields()
.iter()
.any(|field| !field.is_finite())
{
return Err(GraphNodeError::InvalidNode {
node_id,
detail: "factor record contains a non-finite value".to_owned(),
});
}
if node.flags & !ALLOWED_NODE_FLAGS != 0 {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!("reserved flag bits are non-zero: {:#04x}", node.flags),
});
}
if node.neighbors.len() > usize::from(layout.max_degree) {
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!(
"degree {}, maximum {}",
node.neighbors.len(),
layout.max_degree
),
});
}
if let Some(invalid) = node
.neighbors
.iter()
.copied()
.find(|neighbor| *neighbor >= node_count)
{
return Err(GraphNodeError::InvalidNode {
node_id,
detail: format!("neighbour id {invalid} is outside row count {node_count}"),
});
}
let start = output.len();
output.extend_from_slice(node.codes);
for field in node.factors.persisted_fields() {
output.extend_from_slice(&field.to_bits().to_le_bytes());
}
let degree = u8::try_from(node.neighbors.len()).map_err(|_| GraphNodeError::InvalidNode {
node_id,
detail: "degree exceeds u8".to_owned(),
})?;
output.push(degree);
output.push(node.flags);
output.extend_from_slice(&0_u16.to_le_bytes());
for neighbor in node.neighbors {
output.extend_from_slice(&neighbor.to_le_bytes());
}
for _ in node.neighbors.len()..usize::from(layout.max_degree) {
output.extend_from_slice(&UNUSED_NEIGHBOR_ID.to_le_bytes());
}
let end = start
.checked_add(layout.stride as usize)
.ok_or(GraphNodeError::ArithmeticOverflow)?;
output.resize(end, 0);
Ok(())
}
fn validate_code_padding(
layout: GraphNodeLayout,
node_id: u32,
codes: &[u8],
) -> Result<(), GraphNodeError> {
let logical_bytes = (layout.dims as usize).div_ceil(2);
if !layout.dims.is_multiple_of(2)
&& codes
.get(logical_bytes.saturating_sub(1))
.is_some_and(|byte| byte & 0x0f != 0)
{
return Err(GraphNodeError::InvalidNode {
node_id,
detail: "unused low nibble is non-zero".to_owned(),
});
}
if codes
.get(logical_bytes..)
.is_some_and(|padding| padding.iter().any(|byte| *byte != 0))
{
return Err(GraphNodeError::InvalidNode {
node_id,
detail: "padded Bit4 dimensions are non-zero".to_owned(),
});
}
Ok(())
}
fn encode_trailer(
layout: GraphNodeLayout,
node_count: u32,
refinement_passes: RefinementPasses,
output: &mut Vec<u8>,
) {
let start = output.len();
output.extend_from_slice(&NODE_BLOCK_MAGIC);
output.extend_from_slice(&NODE_BLOCK_VERSION.to_le_bytes());
output.extend_from_slice(&refinement_passes.bits().to_le_bytes());
output.extend_from_slice(&layout.dims.to_le_bytes());
output.extend_from_slice(&layout.padded_dims.to_le_bytes());
output.push(layout.max_degree);
output.extend_from_slice(&[0_u8; 3]);
output.extend_from_slice(&layout.stride.to_le_bytes());
output.extend_from_slice(&node_count.to_le_bytes());
let checksum = xxh3_64(output);
output.extend_from_slice(&checksum.to_le_bytes());
output.resize(start.saturating_add(NODE_BLOCK_TRAILER_LEN), 0);
}
fn round_up_cache_line(value: usize) -> Result<usize, GraphNodeError> {
value
.checked_add(CACHE_LINE_BYTES.saturating_sub(1))
.map(|rounded| rounded / CACHE_LINE_BYTES * CACHE_LINE_BYTES)
.ok_or(GraphNodeError::ArithmeticOverflow)
}
#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
use super::{
GraphNodeBlockBuild, GraphNodeBlockInput, GraphNodeError, GraphNodeLayout,
decode_node_blocks, encode_node_blocks,
};
use crate::quant::Bit4Factors;
#[test]
fn cached_descriptor_rejects_a_different_mapping_length() {
let layout = GraphNodeLayout::new(128, 128, 0).expect("cache fixture layout");
let codes = [0_u8; 64];
let node = GraphNodeBlockInput {
codes: &codes,
factors: Bit4Factors::from_persisted(1.0, 0.0, 0.0),
flags: 1,
neighbors: &[],
};
let encoded = encode_node_blocks(GraphNodeBlockBuild {
layout,
nodes: &[node],
})
.expect("cache fixture encodes");
let graph = decode_node_blocks(encoded.as_bytes()).expect("cache fixture decodes");
let descriptor = graph.validated_descriptor();
let shortened = encoded
.as_bytes()
.get(..encoded.as_bytes().len().saturating_sub(1))
.expect("fixture has a trailer byte to remove");
assert!(matches!(
descriptor.bind(shortened),
Err(GraphNodeError::InvalidHeader(detail))
if detail.contains("cached graph region length")
));
let rebound = descriptor
.bind(encoded.as_bytes())
.expect("identical mapping length binds");
assert_eq!(rebound.block(0).expect("node zero").flags(), 1);
let outside = encoded
.block_offset(1)
.expect_err("one-node graph rejects row one");
assert_eq!(
outside.to_string(),
"graph node id 1 is outside dense row count 1"
);
assert_eq!(
GraphNodeError::ArithmeticOverflow.to_string(),
"graph node block arithmetic overflowed"
);
}
}