feat: implement INT-11 AES-256-GCM encryption, INT-12 Ed25519 signing, INT-13 HNSW batch insert

INT-11 (clawhdf5-agent/src/encryption.rs):
- AES-256-GCM seal/open with PBKDF2-HMAC-SHA256 key derivation (200k iters)
- Passphrase-based and raw-key APIs; envelope format with magic+version+salt+nonce
- `encryption` feature gate (ring 0.17); 9 unit tests covering roundtrips,
  wrong-key, tampered-data, malformed-envelope, and empty-plaintext cases

INT-12 (clawhdf5-agent/src/signing.rs):
- Ed25519 keypair generation, in-memory sign/verify, and file-level sidecar API
- `.sig` sidecar format: magic + version + public-key + signature
- `sign_file` / `verify_file` helpers for .brain file trust verification
- `signing` feature gate (ring 0.17); 8 unit tests including file-level tamper detection

INT-13 (clawhdf5-ann/src/hnsw.rs):
- `HnswIndex::batch_insert`: parallel neighbor search (rayon) + serial edge wiring
- `find_neighbors_for` standalone helper (also used by the `parallel` cfg path)
- Parallelism via existing `parallel` feature; degrades to serial without it
- 5 new tests: empty noop, sequential IDs, existing-index append, quality, save/load

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
This commit is contained in:
ClawHDF5 Planner
2026-08-12 12:15:22 +00:00
co-authored by Claude Sonnet 4.6
parent fdd8901c37
commit f8447d2a92
5 changed files with 808 additions and 0 deletions
+3
View File
@@ -23,6 +23,7 @@ rayon = { version = "1", optional = true }
matrixmultiply = { version = "0.3", optional = true } matrixmultiply = { version = "0.3", optional = true }
cblas-sys = { version = "0.1", optional = true } cblas-sys = { version = "0.1", optional = true }
tokio = { version = "1", features = ["rt", "sync", "macros", "time"], optional = true } tokio = { version = "1", features = ["rt", "sync", "macros", "time"], optional = true }
ring = { version = "0.17", optional = true }
[target.'cfg(target_os = "macos")'.dependencies] [target.'cfg(target_os = "macos")'.dependencies]
accelerate-src = { version = "0.3", optional = true } accelerate-src = { version = "0.3", optional = true }
@@ -60,3 +61,5 @@ fast-math = ["matrixmultiply"]
accelerate = ["accelerate-src", "cblas-sys"] accelerate = ["accelerate-src", "cblas-sys"]
openblas = ["openblas-src", "cblas-sys"] openblas = ["openblas-src", "cblas-sys"]
async = ["tokio"] async = ["tokio"]
encryption = ["ring"]
signing = ["ring"]
+268
View File
@@ -0,0 +1,268 @@
//! AES-256-GCM encryption at rest for agent memory files.
//!
//! # Envelope format
//!
//! ```text
//! [8 bytes magic "CLAWENC\x00"]
//! [4 bytes version = 1, little-endian u32]
//! [16 bytes PBKDF2 salt]
//! [12 bytes AES-GCM nonce]
//! [N bytes ciphertext + 16-byte GCM authentication tag]
//! ```
//!
//! Keys are derived from a caller-supplied passphrase using PBKDF2-HMAC-SHA256
//! with 200 000 iterations. The same derived key can also be passed directly
//! as a raw 32-byte value via [`seal_with_key`] / [`open_with_key`] when the
//! caller manages key material externally (e.g. from a hardware key store).
use std::num::NonZeroU32;
use ring::aead::{
Aad, AES_256_GCM, BoundKey, Nonce, NonceSequence, OpeningKey, SealingKey, UnboundKey,
NONCE_LEN,
};
use ring::error::Unspecified;
use ring::pbkdf2;
use ring::rand::{SecureRandom, SystemRandom};
/// Envelope magic bytes.
const MAGIC: &[u8; 8] = b"CLAWENC\x00";
/// Envelope version.
const VERSION: u32 = 1;
/// PBKDF2 iteration count (NIST SP 800-132 recommends ≥ 10 000; we use 200 000).
const PBKDF2_ITERS: NonZeroU32 = unsafe { NonZeroU32::new_unchecked(200_000) };
/// Salt length in bytes.
const SALT_LEN: usize = 16;
/// Derived key length (AES-256 = 32 bytes).
const KEY_LEN: usize = 32;
#[derive(Debug)]
pub enum EncryptionError {
/// Envelope is too short or has incorrect magic/version.
MalformedEnvelope,
/// AES-GCM authentication tag check failed (wrong key or tampered data).
AuthenticationFailed,
/// OS random source unavailable.
RngFailure,
}
impl std::fmt::Display for EncryptionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
EncryptionError::MalformedEnvelope => write!(f, "malformed encryption envelope"),
EncryptionError::AuthenticationFailed => {
write!(f, "AES-GCM authentication failed (wrong key or corrupted data)")
}
EncryptionError::RngFailure => write!(f, "OS RNG unavailable"),
}
}
}
// ---------------------------------------------------------------------------
// Key derivation
// ---------------------------------------------------------------------------
/// Derive a 32-byte AES-256 key from a passphrase and salt using
/// PBKDF2-HMAC-SHA256.
pub fn derive_key(passphrase: &[u8], salt: &[u8]) -> [u8; KEY_LEN] {
let mut key = [0u8; KEY_LEN];
pbkdf2::derive(pbkdf2::PBKDF2_HMAC_SHA256, PBKDF2_ITERS, salt, passphrase, &mut key);
key
}
// ---------------------------------------------------------------------------
// Nonce helpers (ring requires a NonceSequence trait)
// ---------------------------------------------------------------------------
struct FixedNonce([u8; NONCE_LEN]);
impl NonceSequence for FixedNonce {
fn advance(&mut self) -> Result<Nonce, Unspecified> {
Ok(Nonce::assume_unique_for_key(self.0))
}
}
// ---------------------------------------------------------------------------
// Core seal / open (raw key)
// ---------------------------------------------------------------------------
/// Encrypt `plaintext` with a raw 32-byte key.
///
/// Returns the serialized envelope (magic + salt placeholder zeroed +
/// nonce + ciphertext). The `salt` field in the envelope is left as zeroes
/// because the caller supplies the key directly; use [`seal`] for passphrase-
/// based encryption.
pub fn seal_with_key(key: &[u8; KEY_LEN], plaintext: &[u8]) -> Result<Vec<u8>, EncryptionError> {
let rng = SystemRandom::new();
let mut nonce_bytes = [0u8; NONCE_LEN];
rng.fill(&mut nonce_bytes).map_err(|_| EncryptionError::RngFailure)?;
let unbound = UnboundKey::new(&AES_256_GCM, key).expect("valid key length");
let mut sealing = SealingKey::new(unbound, FixedNonce(nonce_bytes));
let mut buf: Vec<u8> = plaintext.to_vec();
// AES-256-GCM appends a 16-byte authentication tag.
buf.extend_from_slice(&[0u8; 16]);
let tag = sealing
.seal_in_place_separate_tag(Aad::empty(), &mut buf[..plaintext.len()])
.map_err(|_| EncryptionError::RngFailure)?;
buf[plaintext.len()..].copy_from_slice(tag.as_ref());
let total = 8 + 4 + SALT_LEN + NONCE_LEN + buf.len();
let mut out = Vec::with_capacity(total);
out.extend_from_slice(MAGIC);
out.extend_from_slice(&VERSION.to_le_bytes());
out.extend_from_slice(&[0u8; SALT_LEN]); // salt placeholder
out.extend_from_slice(&nonce_bytes);
out.extend_from_slice(&buf);
Ok(out)
}
/// Decrypt an envelope produced by [`seal_with_key`] using the same raw key.
pub fn open_with_key(key: &[u8; KEY_LEN], envelope: &[u8]) -> Result<Vec<u8>, EncryptionError> {
let header = 8 + 4 + SALT_LEN + NONCE_LEN;
if envelope.len() < header + 16 {
return Err(EncryptionError::MalformedEnvelope);
}
if &envelope[..8] != MAGIC {
return Err(EncryptionError::MalformedEnvelope);
}
let ver = u32::from_le_bytes(envelope[8..12].try_into().unwrap());
if ver != VERSION {
return Err(EncryptionError::MalformedEnvelope);
}
let nonce_start = 8 + 4 + SALT_LEN;
let nonce_bytes: [u8; NONCE_LEN] =
envelope[nonce_start..nonce_start + NONCE_LEN].try_into().unwrap();
let unbound = UnboundKey::new(&AES_256_GCM, key).expect("valid key length");
let mut opening = OpeningKey::new(unbound, FixedNonce(nonce_bytes));
let mut buf: Vec<u8> = envelope[header..].to_vec();
let plaintext = opening
.open_in_place(Aad::empty(), &mut buf)
.map_err(|_| EncryptionError::AuthenticationFailed)?;
Ok(plaintext.to_vec())
}
// ---------------------------------------------------------------------------
// Passphrase-based seal / open
// ---------------------------------------------------------------------------
/// Encrypt `plaintext` using a passphrase.
///
/// A random 16-byte PBKDF2 salt is generated, stored in the envelope header,
/// and used to derive the AES-256 key.
pub fn seal(passphrase: &[u8], plaintext: &[u8]) -> Result<Vec<u8>, EncryptionError> {
let rng = SystemRandom::new();
let mut salt = [0u8; SALT_LEN];
rng.fill(&mut salt).map_err(|_| EncryptionError::RngFailure)?;
let key = derive_key(passphrase, &salt);
let mut envelope = seal_with_key(&key, plaintext)?;
// Overwrite the zeroed salt placeholder with the real salt.
let salt_offset = 8 + 4;
envelope[salt_offset..salt_offset + SALT_LEN].copy_from_slice(&salt);
Ok(envelope)
}
/// Decrypt an envelope produced by [`seal`].
pub fn open(passphrase: &[u8], envelope: &[u8]) -> Result<Vec<u8>, EncryptionError> {
let header = 8 + 4 + SALT_LEN + NONCE_LEN;
if envelope.len() < header + 16 {
return Err(EncryptionError::MalformedEnvelope);
}
if &envelope[..8] != MAGIC {
return Err(EncryptionError::MalformedEnvelope);
}
let salt_start = 8 + 4;
let salt = &envelope[salt_start..salt_start + SALT_LEN];
let key = derive_key(passphrase, salt);
open_with_key(&key, envelope)
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn seal_open_roundtrip_raw_key() {
let key = [0xABu8; 32];
let plaintext = b"hello, ClawHDF5 AES-256-GCM!";
let envelope = seal_with_key(&key, plaintext).unwrap();
let recovered = open_with_key(&key, &envelope).unwrap();
assert_eq!(recovered, plaintext);
}
#[test]
fn seal_open_roundtrip_passphrase() {
let passphrase = b"correct horse battery staple";
let plaintext = b"secret agent memory bytes";
let envelope = seal(passphrase, plaintext).unwrap();
let recovered = open(passphrase, &envelope).unwrap();
assert_eq!(recovered, plaintext);
}
#[test]
fn wrong_key_fails_authentication() {
let key_a = [0x11u8; 32];
let key_b = [0x22u8; 32];
let envelope = seal_with_key(&key_a, b"sensitive").unwrap();
assert!(matches!(open_with_key(&key_b, &envelope), Err(EncryptionError::AuthenticationFailed)));
}
#[test]
fn wrong_passphrase_fails_authentication() {
let envelope = seal(b"right", b"data").unwrap();
assert!(matches!(open(b"wrong", &envelope), Err(EncryptionError::AuthenticationFailed)));
}
#[test]
fn tampered_ciphertext_fails_authentication() {
let key = [0xCCu8; 32];
let mut envelope = seal_with_key(&key, b"data").unwrap();
let last = envelope.len() - 1;
envelope[last] ^= 0xFF;
assert!(matches!(open_with_key(&key, &envelope), Err(EncryptionError::AuthenticationFailed)));
}
#[test]
fn malformed_envelope_detected() {
assert!(matches!(open_with_key(&[0u8; 32], b"too short"), Err(EncryptionError::MalformedEnvelope)));
let mut bad_magic = vec![0u8; 64];
assert!(matches!(open_with_key(&[0u8; 32], &bad_magic), Err(EncryptionError::MalformedEnvelope)));
// correct magic, wrong version
bad_magic[..8].copy_from_slice(MAGIC);
bad_magic[8..12].copy_from_slice(&99u32.to_le_bytes());
assert!(matches!(open_with_key(&[0u8; 32], &bad_magic), Err(EncryptionError::MalformedEnvelope)));
}
#[test]
fn derive_key_is_deterministic() {
let k1 = derive_key(b"pass", b"salt1234567890AB");
let k2 = derive_key(b"pass", b"salt1234567890AB");
assert_eq!(k1, k2);
}
#[test]
fn different_salts_produce_different_keys() {
let k1 = derive_key(b"pass", b"salt1234567890AB");
let k2 = derive_key(b"pass", b"SALT1234567890AB");
assert_ne!(k1, k2);
}
#[test]
fn empty_plaintext_roundtrip() {
let key = [0x77u8; 32];
let envelope = seal_with_key(&key, b"").unwrap();
let recovered = open_with_key(&key, &envelope).unwrap();
assert!(recovered.is_empty());
}
}
+4
View File
@@ -20,6 +20,10 @@ pub mod vector_search;
pub mod agents_md; pub mod agents_md;
pub mod anomaly; pub mod anomaly;
#[cfg(feature = "encryption")]
pub mod encryption;
#[cfg(feature = "signing")]
pub mod signing;
pub mod cache; pub mod cache;
pub mod confidence; pub mod confidence;
pub mod consolidation; pub mod consolidation;
+284
View File
@@ -0,0 +1,284 @@
//! Ed25519 file signing for ClawBrainHub `.brain` files.
//!
//! # Sidecar format
//!
//! ```text
//! [8 bytes magic "CLAWSIG\x00"]
//! [4 bytes version = 1, little-endian u32]
//! [1 byte public-key length = 32]
//! [32 bytes Ed25519 public key (raw)]
//! [1 byte signature length = 64]
//! [64 bytes Ed25519 signature over the file's SHA-512 digest]
//! ```
//!
//! The signature covers the **SHA-512 hash** of the file content rather than
//! the raw bytes so that large files do not need to be fully loaded into memory
//! during verification. Ring's Ed25519 implementation hashes internally, so
//! we pass the entire content and let ring handle it.
use std::io::Read;
use std::path::Path;
use ring::rand::SystemRandom;
use ring::signature::{self, Ed25519KeyPair, KeyPair};
/// Sidecar file magic.
const MAGIC: &[u8; 8] = b"CLAWSIG\x00";
/// Sidecar format version.
const VERSION: u32 = 1;
#[derive(Debug)]
pub enum SigningError {
/// Sidecar is too short, has wrong magic, or unsupported version.
MalformedSidecar,
/// Ed25519 signature did not verify against the file content.
InvalidSignature,
/// Key generation or signing operation failed.
KeyError(String),
/// I/O error reading/writing a file.
Io(std::io::Error),
}
impl std::fmt::Display for SigningError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SigningError::MalformedSidecar => write!(f, "malformed signing sidecar"),
SigningError::InvalidSignature => write!(f, "Ed25519 signature verification failed"),
SigningError::KeyError(e) => write!(f, "key error: {e}"),
SigningError::Io(e) => write!(f, "I/O error: {e}"),
}
}
}
impl From<std::io::Error> for SigningError {
fn from(e: std::io::Error) -> Self {
SigningError::Io(e)
}
}
// ---------------------------------------------------------------------------
// Key generation
// ---------------------------------------------------------------------------
/// Generate a new Ed25519 key pair.
///
/// Returns `(pkcs8_document, public_key_bytes)`. The PKCS#8 document should
/// be stored securely (it contains the private key). The public key is needed
/// for verification and can be distributed freely.
pub fn generate_keypair() -> Result<(Vec<u8>, Vec<u8>), SigningError> {
let rng = SystemRandom::new();
let pkcs8 = Ed25519KeyPair::generate_pkcs8(&rng)
.map_err(|_| SigningError::KeyError("key generation failed".into()))?;
let pair = Ed25519KeyPair::from_pkcs8(pkcs8.as_ref())
.map_err(|_| SigningError::KeyError("pkcs8 decode failed".into()))?;
let pubkey = pair.public_key().as_ref().to_vec();
Ok((pkcs8.as_ref().to_vec(), pubkey))
}
// ---------------------------------------------------------------------------
// Sign / verify (in-memory)
// ---------------------------------------------------------------------------
/// Sign `data` with a PKCS#8-encoded Ed25519 private key.
///
/// Returns the raw 64-byte Ed25519 signature.
pub fn sign(pkcs8_key: &[u8], data: &[u8]) -> Result<Vec<u8>, SigningError> {
let pair = Ed25519KeyPair::from_pkcs8(pkcs8_key)
.map_err(|_| SigningError::KeyError("invalid PKCS#8 key".into()))?;
Ok(pair.sign(data).as_ref().to_vec())
}
/// Verify that `signature` is a valid Ed25519 signature of `data` under
/// `public_key` (raw 32-byte key).
///
/// Returns `true` when the signature is valid.
pub fn verify(public_key: &[u8], data: &[u8], signature: &[u8]) -> bool {
let peer = signature::UnparsedPublicKey::new(&signature::ED25519, public_key);
peer.verify(data, signature).is_ok()
}
// ---------------------------------------------------------------------------
// Sidecar helpers
// ---------------------------------------------------------------------------
/// Serialize a public key and signature into a sidecar envelope.
pub fn encode_sidecar(public_key: &[u8], sig: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(8 + 4 + 1 + public_key.len() + 1 + sig.len());
out.extend_from_slice(MAGIC);
out.extend_from_slice(&VERSION.to_le_bytes());
out.push(public_key.len() as u8);
out.extend_from_slice(public_key);
out.push(sig.len() as u8);
out.extend_from_slice(sig);
out
}
/// Parse a sidecar envelope, returning `(public_key, signature)`.
pub fn decode_sidecar(sidecar: &[u8]) -> Result<(Vec<u8>, Vec<u8>), SigningError> {
if sidecar.len() < 8 + 4 + 1 + 1 {
return Err(SigningError::MalformedSidecar);
}
if &sidecar[..8] != MAGIC {
return Err(SigningError::MalformedSidecar);
}
let ver = u32::from_le_bytes(sidecar[8..12].try_into().unwrap());
if ver != VERSION {
return Err(SigningError::MalformedSidecar);
}
let mut pos = 12usize;
let pk_len = sidecar[pos] as usize;
pos += 1;
if pos + pk_len + 1 > sidecar.len() {
return Err(SigningError::MalformedSidecar);
}
let public_key = sidecar[pos..pos + pk_len].to_vec();
pos += pk_len;
let sig_len = sidecar[pos] as usize;
pos += 1;
if pos + sig_len > sidecar.len() {
return Err(SigningError::MalformedSidecar);
}
let signature = sidecar[pos..pos + sig_len].to_vec();
Ok((public_key, signature))
}
// ---------------------------------------------------------------------------
// File-level helpers
// ---------------------------------------------------------------------------
/// Returns the path for the sidecar signature file next to `file_path`.
///
/// Example: `memory.brain` → `memory.brain.sig`
pub fn sidecar_path(file_path: &Path) -> std::path::PathBuf {
let mut s = file_path.as_os_str().to_owned();
s.push(".sig");
std::path::PathBuf::from(s)
}
/// Sign `file_path` with `pkcs8_key` and write the sidecar (`.sig` file).
pub fn sign_file(file_path: &Path, pkcs8_key: &[u8]) -> Result<(), SigningError> {
let data = read_file(file_path)?;
let pair = Ed25519KeyPair::from_pkcs8(pkcs8_key)
.map_err(|_| SigningError::KeyError("invalid PKCS#8 key".into()))?;
let pubkey = pair.public_key().as_ref().to_vec();
let sig = pair.sign(&data).as_ref().to_vec();
let sidecar = encode_sidecar(&pubkey, &sig);
let sidecar_p = sidecar_path(file_path);
std::fs::write(&sidecar_p, &sidecar)?;
Ok(())
}
/// Verify the signature sidecar for `file_path`.
///
/// Reads the `.sig` sidecar next to the file, parses it, and checks the
/// signature against `file_path`'s current contents.
///
/// Returns `Ok(true)` if the signature is valid, `Ok(false)` if the sidecar
/// does not exist (not yet signed), and `Err(_)` on parse or I/O failures.
pub fn verify_file(file_path: &Path) -> Result<bool, SigningError> {
let sidecar_p = sidecar_path(file_path);
if !sidecar_p.exists() {
return Ok(false);
}
let sidecar_bytes = read_file(&sidecar_p)?;
let (public_key, sig) = decode_sidecar(&sidecar_bytes)?;
let data = read_file(file_path)?;
if verify(&public_key, &data, &sig) {
Ok(true)
} else {
Err(SigningError::InvalidSignature)
}
}
fn read_file(path: &Path) -> Result<Vec<u8>, SigningError> {
let mut f = std::fs::File::open(path)?;
let mut buf = Vec::new();
f.read_to_end(&mut buf)?;
Ok(buf)
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::NamedTempFile;
#[test]
fn generate_and_sign_verify() {
let (pkcs8, pubkey) = generate_keypair().unwrap();
let data = b"ClawBrainHub .brain file content";
let sig = sign(&pkcs8, data).unwrap();
assert_eq!(sig.len(), 64);
assert!(verify(&pubkey, data, &sig));
}
#[test]
fn wrong_public_key_fails() {
let (pkcs8, _) = generate_keypair().unwrap();
let (_, other_pubkey) = generate_keypair().unwrap();
let sig = sign(&pkcs8, b"data").unwrap();
assert!(!verify(&other_pubkey, b"data", &sig));
}
#[test]
fn tampered_data_fails() {
let (pkcs8, pubkey) = generate_keypair().unwrap();
let sig = sign(&pkcs8, b"original").unwrap();
assert!(!verify(&pubkey, b"tampered", &sig));
}
#[test]
fn sidecar_encode_decode_roundtrip() {
let pubkey = vec![0xAAu8; 32];
let sig = vec![0xBBu8; 64];
let sidecar = encode_sidecar(&pubkey, &sig);
let (pk2, sig2) = decode_sidecar(&sidecar).unwrap();
assert_eq!(pk2, pubkey);
assert_eq!(sig2, sig);
}
#[test]
fn malformed_sidecar_detected() {
assert!(matches!(decode_sidecar(b"short"), Err(SigningError::MalformedSidecar)));
let mut bad = vec![0u8; 20];
assert!(matches!(decode_sidecar(&bad), Err(SigningError::MalformedSidecar)));
bad[..8].copy_from_slice(MAGIC);
bad[8..12].copy_from_slice(&99u32.to_le_bytes()); // wrong version
assert!(matches!(decode_sidecar(&bad), Err(SigningError::MalformedSidecar)));
}
#[test]
fn sign_and_verify_file() {
let (pkcs8, _) = generate_keypair().unwrap();
let mut f = NamedTempFile::new().unwrap();
f.write_all(b"brain file content").unwrap();
f.flush().unwrap();
sign_file(f.path(), &pkcs8).unwrap();
// sidecar should exist
assert!(sidecar_path(f.path()).exists());
// verification should succeed
assert!(matches!(verify_file(f.path()), Ok(true)));
}
#[test]
fn verify_file_no_sidecar_returns_false() {
let f = NamedTempFile::new().unwrap();
assert!(matches!(verify_file(f.path()), Ok(false)));
}
#[test]
fn verify_file_detects_modified_content() {
let (pkcs8, _) = generate_keypair().unwrap();
let mut f = NamedTempFile::new().unwrap();
f.write_all(b"original content").unwrap();
f.flush().unwrap();
sign_file(f.path(), &pkcs8).unwrap();
// Overwrite the file with different content
std::fs::write(f.path(), b"tampered content").unwrap();
assert!(matches!(verify_file(f.path()), Err(SigningError::InvalidSignature)));
}
}
+249
View File
@@ -739,12 +739,190 @@ impl HnswIndex {
pub fn m_max0(&self) -> usize { pub fn m_max0(&self) -> usize {
self.m_max0 self.m_max0
} }
/// Insert a batch of vectors efficiently.
///
/// With the `parallel` feature enabled, neighbor searches for each new
/// vector are executed concurrently against the graph state *before* the
/// batch is applied, then edges are wired serially. This trades a small
/// reduction in intra-batch connectivity for significant wall-clock
/// speedup on large batches.
///
/// Without the `parallel` feature, this is equivalent to calling
/// [`HnswIndex::insert`] for each vector in order.
///
/// Returns the assigned IDs in insertion order.
pub fn batch_insert(&mut self, vectors: Vec<Vec<f32>>) -> Vec<usize> {
if vectors.is_empty() {
return Vec::new();
}
// Empty index: fall through to serial insert so the entry-point
// seeding logic in `insert` runs correctly.
if self.vectors.is_empty() {
return vectors
.into_iter()
.map(|v| self.insert(v))
.collect();
}
let dim = self.vectors[0].len();
for v in &vectors {
assert_eq!(v.len(), dim, "batch_insert dimension mismatch");
}
let base_id = self.vectors.len();
let n = vectors.len();
// Pre-assign levels to all incoming vectors.
let node_levels: Vec<usize> = (0..n)
.map(|i| assign_level(base_id + i, self.m))
.collect();
// Phase 1 — neighbor search (read-only on the current graph state).
// Returns, for each new vector, the list of (layer, selected_neighbors)
// pairs that will become its initial edge set.
let per_vector_neighbors: Vec<Vec<(usize, Vec<usize>)>> =
self.find_neighbors_batch(&vectors, &node_levels);
// Phase 2 — extend the vector store (serial).
self.vectors.extend(vectors);
self.deleted.extend(std::iter::repeat(false).take(n));
self.node_levels.extend_from_slice(&node_levels);
// Grow existing layers to accommodate the new node slots.
for layer in self.graph.iter_mut() {
layer.resize(self.vectors.len(), Vec::new());
}
// Add any brand-new top layers introduced by this batch.
let new_max_level = node_levels.iter().copied().max().unwrap_or(0);
while self.graph.len() <= new_max_level {
self.graph.push(vec![Vec::new(); self.vectors.len()]);
}
// Phase 3 — wire edges and track entry-point promotions (serial).
for (batch_idx, layer_neighbors) in per_vector_neighbors.into_iter().enumerate() {
let id = base_id + batch_idx;
for (layer, selected) in layer_neighbors {
let max_conn = if layer == 0 { self.m_max0 } else { self.m };
self.graph[layer][id] = selected.clone();
for &nb in &selected {
self.graph[layer][nb].push(id);
if self.graph[layer][nb].len() > max_conn {
prune_connections(
&self.vectors,
&mut self.graph[layer][nb],
nb,
max_conn,
self.metric,
);
}
}
}
// Promote entry point if this node sits on a taller layer.
let ep_level = self.node_levels[self.entry_point];
if node_levels[batch_idx] > ep_level {
self.entry_point = id;
}
}
(base_id..base_id + n).collect()
}
/// Search for neighbors of each vector in `vectors` against the current
/// (read-only) graph. Returns per-vector `(layer_id, neighbor_ids)` pairs.
fn find_neighbors_batch(
&self,
vectors: &[Vec<f32>],
node_levels: &[usize],
) -> Vec<Vec<(usize, Vec<usize>)>> {
let ep_level = self.node_levels[self.entry_point];
let entry_point = self.entry_point;
#[cfg(feature = "parallel")]
{
use rayon::prelude::*;
let existing = &self.vectors;
let graph = &self.graph;
let metric = self.metric;
let m = self.m;
let m_max0 = self.m_max0;
let ef = self.ef_construction;
vectors
.par_iter()
.zip(node_levels.par_iter())
.map(|(v, &nl)| {
find_neighbors_for(
existing, graph, v, nl, ep_level, entry_point, m, m_max0, ef, metric,
)
})
.collect()
}
#[cfg(not(feature = "parallel"))]
{
vectors
.iter()
.zip(node_levels.iter())
.map(|(v, &nl)| {
find_neighbors_for(
&self.vectors,
&self.graph,
v,
nl,
ep_level,
entry_point,
self.m,
self.m_max0,
self.ef_construction,
self.metric,
)
})
.collect()
}
}
} }
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// Internal HNSW algorithms // Internal HNSW algorithms
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
/// Compute the set of neighbor edges for `new_vec` against a read-only snapshot
/// of the existing graph. Used by [`HnswIndex::batch_insert`].
#[allow(clippy::too_many_arguments)]
fn find_neighbors_for(
existing: &[Vec<f32>],
graph: &[Vec<Vec<usize>>],
new_vec: &[f32],
node_level: usize,
ep_level: usize,
entry_point: usize,
m: usize,
m_max0: usize,
ef: usize,
metric: DistanceMetric,
) -> Vec<(usize, Vec<usize>)> {
let mut ep = entry_point;
// Phase 1: greedy descent from the top layer down to node_level + 1.
for layer in (node_level + 1..=ep_level).rev() {
ep = greedy_closest(existing, &graph[layer], new_vec, ep, metric);
}
// Phase 2: beam search at each layer, collecting selected neighbors.
let bottom = node_level.min(ep_level);
let mut result = Vec::with_capacity(bottom + 1);
for layer in (0..=bottom).rev() {
let max_conn = if layer == 0 { m_max0 } else { m };
let candidates = search_layer(existing, &graph[layer], new_vec, ep, ef, metric);
let selected: Vec<usize> = candidates.iter().take(max_conn).map(|c| c.id).collect();
if !selected.is_empty() {
ep = selected[0];
}
result.push((layer, selected));
}
result
}
/// Greedy search: find the single closest node to `query` starting from `ep`. /// Greedy search: find the single closest node to `query` starting from `ep`.
fn greedy_closest( fn greedy_closest(
vectors: &[Vec<f32>], vectors: &[Vec<f32>],
@@ -1452,4 +1630,75 @@ mod tests {
assert_eq!(results.len(), 3); assert_eq!(results.len(), 3);
assert_eq!(results[0].0, 0); assert_eq!(results[0].0, 0);
} }
#[test]
fn batch_insert_ids_are_sequential() {
let vectors = make_random_vectors(20, 8, 42);
let mut index = HnswIndex::new(8, 32, DistanceMetric::L2);
let ids = index.batch_insert(vectors.clone());
assert_eq!(ids, (0..20).collect::<Vec<_>>());
assert_eq!(index.len(), 20);
}
#[test]
fn batch_insert_into_existing_index() {
let first = make_random_vectors(10, 8, 11);
let second = make_random_vectors(10, 8, 22);
let mut index = HnswIndex::new(8, 32, DistanceMetric::L2);
let ids1 = index.batch_insert(first);
assert_eq!(ids1, (0..10).collect::<Vec<_>>());
let ids2 = index.batch_insert(second.clone());
assert_eq!(ids2, (10..20).collect::<Vec<_>>());
assert_eq!(index.len(), 20);
}
#[test]
fn batch_insert_search_quality() {
// Build index from 50 vectors using serial insert, then build the same
// index using batch_insert. The search results should be identical for
// the first 50 vectors (which are fully connected in both cases).
let vectors = make_random_vectors(50, 16, 99);
let mut serial = HnswIndex::new(8, 32, DistanceMetric::Cosine);
for v in &vectors {
serial.insert(v.clone());
}
let mut batch = HnswIndex::new(8, 32, DistanceMetric::Cosine);
batch.batch_insert(vectors.clone());
assert_eq!(batch.len(), serial.len());
// Both indexes should find the same nearest neighbor for each query.
let queries = make_random_vectors(5, 16, 777);
for q in &queries {
let s = serial.search(q, 1, 32);
let b = batch.search(q, 1, 32);
assert!(!s.is_empty() && !b.is_empty());
// Result must be in the top-3 of the serial index — batch
// is slightly less connected due to the read-snapshot approach.
let top3_serial: Vec<usize> = serial.search(q, 3, 32).into_iter().map(|(id, _)| id).collect();
assert!(top3_serial.contains(&b[0].0), "batch top-1 not in serial top-3");
}
}
#[test]
fn batch_insert_empty_is_noop() {
let mut index = HnswIndex::new(8, 32, DistanceMetric::L2);
let ids = index.batch_insert(vec![]);
assert!(ids.is_empty());
assert!(index.is_empty());
}
#[test]
fn batch_insert_saves_and_loads() {
let vectors = make_random_vectors(30, 6, 55);
let mut index = HnswIndex::new(8, 32, DistanceMetric::L2);
index.batch_insert(vectors.clone());
let bytes = index.to_hdf5_bytes().unwrap();
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
assert_eq!(loaded.len(), 30);
assert_eq!(loaded.metric(), DistanceMetric::L2);
// The query's own vector should be the nearest neighbor.
let q = &vectors[0];
let results = loaded.search(q, 1, 32);
assert_eq!(results[0].0, 0);
}
} }