Skip to main content

nodedb_vector/hnsw/graph/index/
neighbors.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Neighbor-list access.
4//!
5//! After a checkpoint restore the lists live in one flat, zero-copy store; a
6//! mutation materializes them back into per-node `Vec`s. Both readers below
7//! consult the flat store first so callers never have to know which
8//! representation is currently in use.
9
10use super::state::HnswIndex;
11
12impl HnswIndex {
13    /// Get neighbors of a node at a specific layer.
14    /// Uses flat zero-copy storage if available, otherwise per-node Vec.
15    #[inline]
16    pub(crate) fn neighbors_at(&self, node_id: u32, layer: usize) -> &[u32] {
17        if let Some(ref flat) = self.flat_neighbors {
18            return flat.neighbors_at(node_id, layer);
19        }
20        let node = &self.nodes[node_id as usize];
21        if layer < node.neighbors.len() {
22            &node.neighbors[layer]
23        } else {
24            &[]
25        }
26    }
27
28    /// Number of layers a node participates in.
29    #[inline]
30    pub(crate) fn node_num_layers(&self, node_id: u32) -> usize {
31        if let Some(ref flat) = self.flat_neighbors {
32            return flat.num_layers(node_id);
33        }
34        self.nodes[node_id as usize].neighbors.len()
35    }
36
37    /// Ensure mutable per-node neighbor Vecs are available.
38    /// Materializes flat storage back to per-node Vecs if needed.
39    pub(crate) fn ensure_mutable_neighbors(&mut self) {
40        if let Some(flat) = self.flat_neighbors.take() {
41            let nested = flat.to_nested(self.nodes.len());
42            for (i, layers) in nested.into_iter().enumerate() {
43                self.nodes[i].neighbors = layers;
44            }
45        }
46    }
47
48    /// Export all neighbor lists for snapshot transfer.
49    pub fn export_neighbors(&self) -> Vec<Vec<Vec<u32>>> {
50        self.nodes.iter().map(|n| n.neighbors.clone()).collect()
51    }
52}