tidaldb/tidal/src/storage/vector/brute/mod.rs

//! Brute-force (exact) vector index.
//!
//! [`BruteForceIndex`] performs linear-scan L2 search over all stored vectors.
//! It is the correctness baseline: every other index implementation must return
//! the same top-k results (within quantization tolerance). It is also used for
//! small datasets where the O(n) scan cost is acceptable.

use std::{collections::HashMap, io::Write, path::Path, sync::RwLock};

use super::{
    VectorError, VectorId, VectorIndex, VectorIndexConfig, VectorSearchResult, read_f32_le,
    read_u32_le, read_u64_le, validate_dimensions,
};

// ---------------------------------------------------------------------------
// Binary format constants
// ---------------------------------------------------------------------------

/// Magic bytes identifying a brute-force vector index file.
const MAGIC: &[u8; 4] = b"BFVI";

/// Current binary format version.
const FORMAT_VERSION: u8 = 0x01;

// ---------------------------------------------------------------------------
// Distance functions
// ---------------------------------------------------------------------------

/// Compute the squared Euclidean (L2) distance between two vectors.
///
/// This avoids the `sqrt` call -- the squared distance preserves ranking order
/// and is sufficient for nearest-neighbor selection.
///
/// # Panics (debug only)
///
/// Debug-asserts that `a` and `b` have the same length.
pub fn l2_distance_sq(a: &[f32], b: &[f32]) -> f32 {
    debug_assert_eq!(a.len(), b.len());
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| {
            let d = x - y;
            d * d
        })
        .sum()
}

// ---------------------------------------------------------------------------
// BruteForceIndex
// ---------------------------------------------------------------------------

/// Exact nearest-neighbor index using linear scan.
///
/// Stores all vectors in a `HashMap` behind a `RwLock`. Search computes L2
/// squared distance against every stored vector and returns the top-k by
/// ascending distance.
///
/// This implementation has no tombstones -- `delete()` is a true removal.
/// `len()` and `len_live()` always return the same value.
///
/// # Thread safety
///
/// All reads take a shared lock; all writes take an exclusive lock. This is
/// acceptable because brute-force is not a hot-path production index -- it is
/// used for correctness baselines, small datasets, and tests.
#[derive(Debug)]
pub struct BruteForceIndex {
    vectors: RwLock<HashMap<VectorId, Vec<f32>>>,
    config: VectorIndexConfig,
}

impl BruteForceIndex {
    /// Create a new, empty brute-force index with the given configuration.
    #[must_use]
    pub fn new(config: VectorIndexConfig) -> Self {
        Self {
            vectors: RwLock::new(HashMap::new()),
            config,
        }
    }

    /// Validate that a vector's dimensionality matches the index configuration.
    const fn validate_dimensions(&self, vec: &[f32]) -> Result<(), VectorError> {
        validate_dimensions(self.config.dimensions, vec.len())
    }
}

impl VectorIndex for BruteForceIndex {
    fn insert(&self, id: VectorId, embedding: &[f32]) -> Result<(), VectorError> {
        self.validate_dimensions(embedding)?;
        self.vectors
            .write()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on write: {e}")))?
            .insert(id, embedding.to_vec());
        Ok(())
    }

    /// Search for the `k` nearest neighbors by exhaustive linear scan.
    ///
    /// The `ef_search` parameter is accepted for trait compliance but ignored --
    /// brute-force search is exact and has no beam width parameter.
    fn search(
        &self,
        query: &[f32],
        k: usize,
        _ef_search: usize,
    ) -> Result<Vec<VectorSearchResult>, VectorError> {
        self.validate_dimensions(query)?;
        let guard = self
            .vectors
            .read()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on read: {e}")))?;

        let mut results: Vec<VectorSearchResult> = guard
            .iter()
            .map(|(id, vec)| VectorSearchResult {
                id: *id,
                distance: l2_distance_sq(query, vec),
            })
            .collect();

        drop(guard);
        results.sort_by(|a, b| {
            a.distance
                .partial_cmp(&b.distance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results.truncate(k);
        Ok(results)
    }

    /// Filtered search: only compute distance for vectors where `filter(id)` is true.
    ///
    /// The `ef_search` parameter is accepted for trait compliance but ignored.
    fn filtered_search(
        &self,
        query: &[f32],
        k: usize,
        _ef_search: usize,
        filter: &dyn Fn(VectorId) -> bool,
    ) -> Result<Vec<VectorSearchResult>, VectorError> {
        self.validate_dimensions(query)?;
        let guard = self
            .vectors
            .read()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on read: {e}")))?;

        let mut results: Vec<VectorSearchResult> = guard
            .iter()
            .filter(|(id, _)| filter(**id))
            .map(|(id, vec)| VectorSearchResult {
                id: *id,
                distance: l2_distance_sq(query, vec),
            })
            .collect();

        drop(guard);
        results.sort_by(|a, b| {
            a.distance
                .partial_cmp(&b.distance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results.truncate(k);
        Ok(results)
    }

    fn delete(&self, id: VectorId) -> Result<(), VectorError> {
        let removed = self
            .vectors
            .write()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on write: {e}")))?
            .remove(&id);
        if removed.is_none() {
            return Err(VectorError::NotFound { id });
        }
        Ok(())
    }

    /// Reserve additional capacity. This is a no-op for `HashMap`-backed storage --
    /// `HashMap` resizes automatically. The method is provided for trait compliance
    /// and always succeeds.
    fn reserve(&self, _additional: usize) -> Result<(), VectorError> {
        Ok(())
    }

    fn save(&self, path: &Path) -> Result<(), VectorError> {
        let guard = self
            .vectors
            .read()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on read: {e}")))?;

        let mut file = std::fs::File::create(path)?;

        // Header: magic + version + dimensions + count
        file.write_all(MAGIC)?;
        file.write_all(&[FORMAT_VERSION])?;
        let dims = self.config.dimensions as u32;
        file.write_all(&dims.to_le_bytes())?;
        let count = guard.len() as u64;
        file.write_all(&count.to_le_bytes())?;

        // Per-vector: id + floats
        for (id, vec) in &*guard {
            file.write_all(&id.to_le_bytes())?;
            for &val in vec {
                file.write_all(&val.to_le_bytes())?;
            }
        }

        drop(guard);
        file.flush().map_err(std::convert::Into::into)
    }

    fn load(path: &Path, config: &VectorIndexConfig) -> Result<Self, VectorError> {
        let data = std::fs::read(path)?;
        Self::deserialize(&data, config)
    }

    /// For brute-force, `view` delegates to `load` -- there is no mmap mode.
    /// The entire index is read into memory regardless.
    fn view(path: &Path, config: &VectorIndexConfig) -> Result<Self, VectorError> {
        Self::load(path, config)
    }

    fn len(&self) -> usize {
        self.vectors.read().map_or(0, |guard| guard.len())
    }

    fn len_live(&self) -> usize {
        // No tombstones in brute-force -- all entries are live.
        self.len()
    }
}

impl BruteForceIndex {
    /// Deserialize from a byte buffer (shared by `load` and `view`).
    fn deserialize(data: &[u8], config: &VectorIndexConfig) -> Result<Self, VectorError> {
        // Minimum header size: 4 (magic) + 1 (version) + 4 (dims) + 8 (count) = 17
        const HEADER_SIZE: usize = 17;
        if data.len() < HEADER_SIZE {
            return Err(VectorError::CorruptedIndex(
                "file too small for header".into(),
            ));
        }

        // Validate magic
        if &data[..4] != MAGIC {
            return Err(VectorError::CorruptedIndex(format!(
                "invalid magic: expected {:?}, got {:?}",
                MAGIC,
                &data[..4]
            )));
        }

        // Validate version
        if data[4] != FORMAT_VERSION {
            return Err(VectorError::CorruptedIndex(format!(
                "unsupported version: expected {FORMAT_VERSION:#04x}, got {:#04x}",
                data[4]
            )));
        }

        // Read dimensions
        let dims = read_u32_le(data, 5) as usize;
        if dims != config.dimensions {
            return Err(VectorError::CorruptedIndex(format!(
                "dimension mismatch: file has {dims}, config expects {}",
                config.dimensions
            )));
        }

        // Read count.
        //
        // SECURITY: `count` comes straight off disk and is fully
        // corruption/attacker-controlled. Every byte of arithmetic below uses
        // checked operations and is bounded against the *actual* buffer length
        // BEFORE any allocation or loop. Without this, a hostile `count` would
        // (a) overflow `count * bytes_per_vector` in release builds, wrapping to
        // a tiny `expected_size` that slips past the truncation guard, then
        // (b) drive `HashMap::with_capacity(count)` into a multi-exabyte alloc
        // (process abort) and an out-of-bounds read loop (panic). A corrupt
        // *derived* index is recoverable (rebuildable from the entity store), so
        // it must surface as `CorruptedIndex`, never crash the recovery path.
        let count = read_u64_le(data, 9) as usize;

        // Bytes consumed per stored vector: id (8) + floats (dims * 4).
        // Guard `dims` (also from disk-ish via header; already matched against
        // config above, but compute defensively) so the multiply cannot wrap.
        let float_bytes = dims
            .checked_mul(4)
            .ok_or_else(|| VectorError::CorruptedIndex(format!("dimensions too large: {dims}")))?;
        let bytes_per_vector = 8usize.checked_add(float_bytes).ok_or_else(|| {
            VectorError::CorruptedIndex(format!("per-vector size overflow at dims {dims}"))
        })?;
        debug_assert!(
            bytes_per_vector >= 8,
            "bytes_per_vector includes the 8-byte id and cannot be zero"
        );

        // Cap `count` by how many vectors the body could actually contain. This
        // makes the subsequent `with_capacity` and read loop safe regardless of
        // the on-disk value: a damaged header can never request more entries
        // than the buffer can hold.
        let body_len = data.len() - HEADER_SIZE; // safe: data.len() >= HEADER_SIZE checked above
        let max_count = body_len / bytes_per_vector; // bytes_per_vector >= 8, no div-by-zero
        if count > max_count {
            return Err(VectorError::CorruptedIndex(format!(
                "implausible vector count: header claims {count}, but the {body_len}-byte \
                 body holds at most {max_count} vectors of {bytes_per_vector} bytes each"
            )));
        }

        // Validate the exact declared size (now provably non-overflowing because
        // `count <= max_count` and `max_count * bytes_per_vector <= body_len`).
        let expected_size = count
            .checked_mul(bytes_per_vector)
            .and_then(|body| body.checked_add(HEADER_SIZE))
            .ok_or_else(|| {
                VectorError::CorruptedIndex(format!(
                    "index size overflow: count {count} * {bytes_per_vector} + {HEADER_SIZE}"
                ))
            })?;
        if data.len() < expected_size {
            return Err(VectorError::CorruptedIndex(format!(
                "file truncated: expected at least {expected_size} bytes, got {}",
                data.len()
            )));
        }

        // Parse vectors. `count` is now bounded by the buffer, so `with_capacity`
        // cannot over-allocate and the read loop cannot index out of bounds.
        let mut vectors = HashMap::with_capacity(count);
        let mut offset = HEADER_SIZE;
        for _ in 0..count {
            let id = read_u64_le(data, offset);
            offset += 8;

            let mut vec = Vec::with_capacity(dims);
            for _ in 0..dims {
                vec.push(read_f32_le(data, offset));
                offset += 4;
            }

            vectors.insert(id, vec);
        }

        Ok(Self {
            vectors: RwLock::new(vectors),
            config: config.clone(),
        })
    }
}

#[cfg(test)]
mod tests;