urna_runtime/ann/codec/
mod.rs

1//! On-disk encoding/decoding for the HNSW section (`0x07`, `encoding=raw`).
2//!
3//! Payload version 2 bitpacks the graph with the shared `intpack` codec:
4//! after the fixed header, three length-prefixed `intpack` columns hold the
5//! per-node level, the per-layer neighbour counts, and every neighbour id
6//! (in build order). intpack is order-preserving, so the decoded graph is
7//! identical to the built one and recall is unchanged; the win is the
8//! neighbour ids dropping from a flat u32 to ~`ceil(log2 n)` bits each.
9//!
10//! Version 1 (flat u32 per id) is still accepted on read so existing files
11//! keep opening. The section is optional and excluded from content_hash.
12//!
13//! Section checksum (8 bytes of SHA-256 over the physical bytes) is
14//! computed by the writer at file-build time, not here.
15
16use urna_format::encoding::{pack_u64s, unpack_u64s};
17use urna_format::error::UrnaError;
18
19use super::{HNSW_PAYLOAD_VERSION, HnswIndex, Node};
20use crate::error::RuntimeError;
21use crate::materialize::PackedVectors;
22use urna_format::bytes::le_u32;
23
24impl HnswIndex {
25    /// Encode the index to bytes for embedding in section `0x07` (v2).
26    pub fn to_bytes(&self) -> Vec<u8> {
27        let mut out = Vec::with_capacity(7 * 4 + self.n * 4);
28        out.extend_from_slice(&HNSW_PAYLOAD_VERSION.to_le_bytes());
29        out.extend_from_slice(&(self.m as u32).to_le_bytes());
30        out.extend_from_slice(&(self.m_max0 as u32).to_le_bytes());
31        out.extend_from_slice(&(self.ef_construction as u32).to_le_bytes());
32        out.extend_from_slice(&self.entry_point.to_le_bytes());
33        out.extend_from_slice(&self.max_level.to_le_bytes());
34        out.extend_from_slice(&(self.n as u32).to_le_bytes());
35
36        let mut levels: Vec<u64> = Vec::with_capacity(self.nodes.len());
37        let mut counts: Vec<u64> = Vec::new();
38        let mut ids: Vec<u64> = Vec::new();
39        for node in &self.nodes {
40            levels.push(node.level as u64);
41            for layer in 0..=node.level {
42                let nbrs = &node.neighbors[layer as usize];
43                counts.push(nbrs.len() as u64);
44                ids.extend(nbrs.iter().map(|&id| id as u64));
45            }
46        }
47        for col in [&levels, &counts, &ids] {
48            let blob = pack_u64s(col);
49            out.extend_from_slice(&(blob.len() as u32).to_le_bytes());
50            out.extend_from_slice(&blob);
51        }
52        out
53    }
54
55    /// Parse an HNSW payload (v1 or v2). The vectors are reconstructed
56    /// from the embeddings section by the caller; this constructor is for
57    /// the on-disk graph only. Call `attach_vectors` before search.
58    pub fn from_bytes(bytes: &[u8], n: usize, dim: usize) -> Result<Self, RuntimeError> {
59        let mut cur = ByteCursor::new(bytes);
60        let version = cur.u32()?;
61        let m = cur.u32()? as usize;
62        let m_max0 = cur.u32()? as usize;
63        let ef_construction = cur.u32()? as usize;
64        let entry_point = cur.u32()?;
65        let max_level = cur.u32()?;
66        let n_nodes = cur.u32()? as usize;
67        if n_nodes != n {
68            return Err(malformed(format!(
69                "node count {} != n_embeddings {}",
70                n_nodes, n
71            )));
72        }
73        let nodes = match version {
74            1 => decode_nodes_v1(&mut cur, n_nodes)?,
75            2 => decode_nodes_v2(&mut cur, n_nodes)?,
76            other => {
77                return Err(RuntimeError::Format(UrnaError::UnsupportedSectionVersion {
78                    section_id: urna_format::layout::SECTION_HNSW_INDEX,
79                    version: other,
80                }));
81            }
82        };
83        // hostile-input guard: `entry_point` and `max_level` are parsed from
84        // untrusted bytes and later index `nodes` / drive the layer descent at
85        // search time. Neighbour ids are range-checked in decode_nodes_v{1,2};
86        // validate the two scalar fields here so a crafted payload can never
87        // index out of range (an OOB slice = guaranteed panic / DoS).
88        if max_level as u64 > MAX_LEVEL {
89            return Err(malformed("hnsw: max_level exceeds cap"));
90        }
91        if n_nodes > 0 && entry_point as usize >= n_nodes {
92            return Err(malformed("hnsw: entry_point out of range"));
93        }
94        Ok(Self {
95            m,
96            m_max0,
97            ef_construction,
98            entry_point,
99            max_level,
100            nodes,
101            store: PackedVectors::empty(),
102            dim,
103            n,
104            ef_search: ef_construction,
105        })
106    }
107}
108
109fn malformed(reason: impl Into<String>) -> RuntimeError {
110    RuntimeError::Format(UrnaError::MalformedSectionPayload {
111        section_id: urna_format::layout::SECTION_HNSW_INDEX,
112        reason: reason.into(),
113    })
114}
115
116// hostile-input caps. a real hnsw never approaches these (levels are
117// geometric ~log_M(n); degree <= m_max0 ~ 2*M), so they never reject a
118// valid file. they bound the decoded-but-untrusted level/degree VALUES so a
119// crafted payload cannot overflow the count sums or force a giant
120// allocation before the column-length cross-checks run.
121const MAX_LEVEL: u64 = 64;
122const MAX_DEGREE: u64 = 4096;
123
124fn check_cap(values: &[u64], cap: u64, what: &str) -> Result<(), RuntimeError> {
125    if values.iter().any(|&v| v > cap) {
126        return Err(malformed(format!("hnsw: {} exceeds cap {}", what, cap)));
127    }
128    Ok(())
129}
130
131/// v1: flat `u32` per id, inline per node.
132fn decode_nodes_v1(cur: &mut ByteCursor, n_nodes: usize) -> Result<Vec<Node>, RuntimeError> {
133    let mut nodes = Vec::with_capacity(n_nodes.min(1 << 16));
134    for _ in 0..n_nodes {
135        let level = cur.u32()?;
136        if level as u64 > MAX_LEVEL {
137            return Err(malformed("hnsw v1: level exceeds cap"));
138        }
139        let mut neighbors = Vec::with_capacity((level as usize) + 1);
140        for _ in 0..=level {
141            let k = cur.u32()? as usize;
142            if k as u64 > MAX_DEGREE {
143                return Err(malformed("hnsw v1: degree exceeds cap"));
144            }
145            let mut ids = Vec::with_capacity(k);
146            for _ in 0..k {
147                let id = cur.u32()?;
148                if id as usize >= n_nodes {
149                    return Err(malformed("hnsw v1: neighbour id out of range"));
150                }
151                ids.push(id);
152            }
153            neighbors.push(ids);
154        }
155        nodes.push(Node { level, neighbors });
156    }
157    Ok(nodes)
158}
159
160/// v2: three `intpack` columns (levels, per-layer counts, neighbour ids).
161fn decode_nodes_v2(cur: &mut ByteCursor, n_nodes: usize) -> Result<Vec<Node>, RuntimeError> {
162    let levels = cur.intpack_column()?;
163    if levels.len() != n_nodes {
164        return Err(malformed("hnsw v2: level count mismatch"));
165    }
166    check_cap(&levels, MAX_LEVEL, "level")?;
167    // safe: each level <= MAX_LEVEL and levels.len() is bounded by the
168    // physically present column, so the sum cannot overflow usize.
169    let total_counts: usize = levels.iter().map(|l| *l as usize + 1).sum();
170    let counts = cur.intpack_column()?;
171    if counts.len() != total_counts {
172        return Err(malformed("hnsw v2: layer-count column mismatch"));
173    }
174    check_cap(&counts, MAX_DEGREE, "degree")?;
175    let total_ids: usize = counts.iter().map(|c| *c as usize).sum();
176    let ids = cur.intpack_column()?;
177    if ids.len() != total_ids {
178        return Err(malformed("hnsw v2: neighbour-id column mismatch"));
179    }
180    // range-check every neighbour id before it is used to index `nodes` /
181    // the vector store at search time (a hostile id = OOB panic / DoS).
182    if ids.iter().any(|&id| id as usize >= n_nodes) {
183        return Err(malformed("hnsw v2: neighbour id out of range"));
184    }
185    let mut nodes = Vec::with_capacity(n_nodes.min(1 << 16));
186    let (mut ci, mut ii) = (0usize, 0usize);
187    for &lvl in &levels {
188        let level = lvl as u32;
189        let mut neighbors = Vec::with_capacity(level as usize + 1);
190        for _ in 0..=level {
191            let k = counts[ci] as usize;
192            ci += 1;
193            let mut layer = Vec::with_capacity(k);
194            for _ in 0..k {
195                layer.push(ids[ii] as u32);
196                ii += 1;
197            }
198            neighbors.push(layer);
199        }
200        nodes.push(Node { level, neighbors });
201    }
202    Ok(nodes)
203}
204
205/// Light cursor for parsing the on-disk HNSW payload.
206struct ByteCursor<'a> {
207    buf: &'a [u8],
208    pos: usize,
209}
210impl<'a> ByteCursor<'a> {
211    fn new(buf: &'a [u8]) -> Self {
212        Self { buf, pos: 0 }
213    }
214    fn u32(&mut self) -> Result<u32, RuntimeError> {
215        if 4 > self.buf.len() - self.pos {
216            return Err(malformed("hnsw: unexpected EOF"));
217        }
218        let v = le_u32(&self.buf[self.pos..self.pos + 4])?;
219        self.pos += 4;
220        Ok(v)
221    }
222    /// read a length-prefixed `intpack` blob and decode it to `u64`s.
223    fn intpack_column(&mut self) -> Result<Vec<u64>, RuntimeError> {
224        let len = self.u32()? as usize;
225        if len > self.buf.len() - self.pos {
226            return Err(malformed("hnsw: truncated intpack column"));
227        }
228        let blob = &self.buf[self.pos..self.pos + len];
229        self.pos += len;
230        unpack_u64s(blob).map_err(RuntimeError::Format)
231    }
232}
233
234#[cfg(test)]
235mod tests;
urna_runtime/ann/codec/mod.rs

urna_runtime/ann/codec/
mod.rs
