weavatrix-search-vector 0.3.1

use super::format::{
    align8, bytes_for, checksum, decode_f32s, decode_u16s, decode_u32s, decode_u64s, read_u64,
    to_usize,
};
use super::validation::{validate_canonical_offsets, validate_offsets, validate_owned_graph};
use super::{GRAPH_HEADER_LEN, HEADER_LEN, Header, NONE_ENTRY};
use crate::error::SearchError;
use crate::hnsw::{Graph, NodeLinks, RoutingIndex, VectorIndex};
use crate::vector::VectorStore;

pub(super) fn decode_snapshot(bytes: &[u8]) -> Result<VectorIndex, SearchError> {
    let header = Header::parse(bytes)?;
    if header.file_len != bytes.len() {
        return Err(SearchError::CorruptSnapshot(
            "snapshot file length does not match",
        ));
    }
    if checksum(&bytes[HEADER_LEN..]) != header.checksum {
        return Err(SearchError::CorruptSnapshot(
            "snapshot payload checksum does not match",
        ));
    }
    validate_canonical_offsets(&header)?;
    let count = header.count;
    let value_count = count
        .checked_mul(header.config.dimensions)
        .ok_or(SearchError::CapacityOverflow)?;
    let keys = decode_u64s(bytes, header.keys_offset, count)?;
    let values = decode_f32s(bytes, header.vectors_offset, value_count)?;
    let vectors =
        VectorStore::from_values(header.config.dimensions, header.config.metric, keys, values)?;
    let routing = decode_routing(bytes, &header)?;
    let graphs = decode_graphs(bytes, &header)?;
    VectorIndex::from_parts(header.config, vectors, graphs, routing)
}

fn decode_routing(bytes: &[u8], header: &Header) -> Result<RoutingIndex, SearchError> {
    let codes = decode_u16s(bytes, header.routing_codes_offset, header.count)?;
    if codes.windows(2).any(|pair| pair[0] > pair[1]) {
        return Err(SearchError::CorruptSnapshot("routing codes are not sorted"));
    }
    let routing_nodes = decode_u32s(bytes, header.routing_nodes_offset, header.count)?;
    if routing_nodes
        .iter()
        .any(|node| *node as usize >= header.count)
    {
        return Err(SearchError::CorruptSnapshot(
            "routing node is outside the vector range",
        ));
    }
    Ok(RoutingIndex {
        entries: codes.into_iter().zip(routing_nodes).collect(),
    })
}

fn decode_graphs(bytes: &[u8], header: &Header) -> Result<Vec<Graph>, SearchError> {
    let mut cursor = header.graphs_offset;
    let mut graphs = Vec::new();
    graphs
        .try_reserve_exact(header.config.replicas)
        .map_err(|_| SearchError::AllocationFailed)?;
    for _ in 0..header.config.replicas {
        let (graph, next) = decode_graph(bytes, cursor, header.count)?;
        graphs.push(graph);
        cursor = next;
    }
    if cursor != bytes.len() {
        return Err(SearchError::CorruptSnapshot(
            "snapshot graph length does not match file",
        ));
    }
    Ok(graphs)
}

fn decode_graph(
    bytes: &[u8],
    mut cursor: usize,
    count: usize,
) -> Result<(Graph, usize), SearchError> {
    let entry_raw = read_u64(bytes, cursor)?;
    let max_level = to_usize(read_u64(bytes, cursor + 8)?)?;
    let node_count = to_usize(read_u64(bytes, cursor + 16)?)?;
    let total_layers = to_usize(read_u64(bytes, cursor + 24)?)?;
    let neighbor_count = to_usize(read_u64(bytes, cursor + 32)?)?;
    if node_count != count {
        return Err(SearchError::CorruptSnapshot(
            "graph node count does not match vector count",
        ));
    }
    cursor = cursor
        .checked_add(GRAPH_HEADER_LEN)
        .ok_or(SearchError::CapacityOverflow)?;
    let node_layers = decode_u64s(bytes, cursor, count + 1)?;
    cursor = cursor
        .checked_add(bytes_for::<u64>(count + 1)?)
        .ok_or(SearchError::CapacityOverflow)?;
    let layer_neighbors = decode_u64s(bytes, cursor, total_layers + 1)?;
    cursor = cursor
        .checked_add(bytes_for::<u64>(total_layers + 1)?)
        .ok_or(SearchError::CapacityOverflow)?;
    let neighbors = decode_u32s(bytes, cursor, neighbor_count)?;
    cursor = align8(
        cursor
            .checked_add(bytes_for::<u32>(neighbor_count)?)
            .ok_or(SearchError::CapacityOverflow)?,
    )?;
    validate_offsets(&node_layers, total_layers, "graph node-layer")?;
    validate_offsets(&layer_neighbors, neighbor_count, "graph layer-neighbor")?;
    let nodes = materialize_nodes(count, &node_layers, &layer_neighbors, &neighbors)?;
    validate_owned_graph(&nodes, entry_raw, max_level, count)?;
    let entry = if entry_raw == NONE_ENTRY {
        None
    } else {
        Some(to_usize(entry_raw)?)
    };
    Ok((
        Graph {
            entry,
            max_level,
            nodes,
        },
        cursor,
    ))
}

fn materialize_nodes(
    count: usize,
    node_layers: &[u64],
    layer_neighbors: &[u64],
    neighbors: &[u32],
) -> Result<Vec<NodeLinks>, SearchError> {
    let mut nodes = Vec::new();
    nodes
        .try_reserve_exact(count)
        .map_err(|_| SearchError::AllocationFailed)?;
    for node in 0..count {
        let layer_start = to_usize(node_layers[node])?;
        let layer_end = to_usize(node_layers[node + 1])?;
        let mut layers = Vec::new();
        layers
            .try_reserve_exact(layer_end.saturating_sub(layer_start))
            .map_err(|_| SearchError::AllocationFailed)?;
        for layer in layer_start..layer_end {
            let start = to_usize(layer_neighbors[layer])?;
            let end = to_usize(layer_neighbors[layer + 1])?;
            layers.push(neighbors[start..end].to_vec());
        }
        nodes.push(NodeLinks { layers });
    }
    Ok(nodes)
}