summa-core 2.0.0

//! Quantized sparse values addressed by logical document/ordinal, inside BMP.
//! The directory is compact and the payload stays evictable. Neither contains
//! physical BMP IDs, so physical reorder never rewrites forward values.

use super::logical_address::LogicalUnit;
use crate::directories::OwnedBytes;
use crate::{Error, Result};

mod codec;
pub(crate) use codec::ForwardVector;
#[cfg(any(feature = "native", feature = "wasm", test))]
pub(crate) use codec::RowWriter;

const ROW_BYTES: usize = 16;
pub(super) const TRAILER_BYTES: usize = 16;
const STORAGE_DISABLED: u32 = 1;

#[derive(Clone)]
pub(crate) struct BmpForward {
    payload: OwnedBytes,
    rows: OwnedBytes,
    #[cfg_attr(not(any(feature = "native", feature = "wasm", test)), allow(dead_code))]
    dims: u32,
}

fn corrupt(message: &str) -> Error {
    Error::Corruption(format!("BMP forward values: {message}"))
}

impl BmpForward {
    pub(crate) fn parse_optional(
        bytes: OwnedBytes,
        count: u32,
        docs: u32,
        dims: u32,
    ) -> Result<Option<Self>> {
        let trailer = bytes
            .len()
            .checked_sub(TRAILER_BYTES)
            .ok_or_else(|| corrupt("missing storage trailer"))?;
        let flags = u32::from_le_bytes(bytes[trailer + 4..trailer + 8].try_into().unwrap());
        if flags == STORAGE_DISABLED {
            if bytes.len() != TRAILER_BYTES || bytes[..4] != [0; 4] || bytes[8..] != [0; 8] {
                return Err(corrupt("disabled storage has a payload or vector count"));
            }
            return Ok(None);
        }
        Self::parse(bytes, count, docs, dims).map(Some)
    }

    pub(crate) fn parse(bytes: OwnedBytes, count: u32, docs: u32, dims: u32) -> Result<Self> {
        let trailer = bytes
            .len()
            .checked_sub(TRAILER_BYTES)
            .ok_or_else(|| corrupt("missing trailer"))?;
        let data = bytes.as_slice();
        let declared = u32::from_le_bytes(data[trailer..trailer + 4].try_into().unwrap());
        let payload_len =
            usize::try_from(u64::from_le_bytes(data[trailer + 8..].try_into().unwrap()))
                .map_err(|_| corrupt("payload exceeds address space"))?;
        let flags = u32::from_le_bytes(data[trailer + 4..trailer + 8].try_into().unwrap());
        if declared != count
            || flags != 0
            || (count as usize)
                .checked_mul(ROW_BYTES)
                .and_then(|n| n.checked_add(payload_len))
                != Some(trailer)
        {
            return Err(corrupt("invalid directory length or count"));
        }
        let result = Self {
            payload: bytes.slice(0..payload_len),
            rows: bytes.slice(payload_len..trailer),
            dims,
        };
        let mut previous_key = None;
        let mut previous_offset = 0;
        for i in 0..count {
            let key = result.key(i);
            let offset = result.offset(i);
            if key.doc >= docs
                || previous_key.is_some_and(|p| p >= key)
                || result.rows.as_slice()[i as usize * ROW_BYTES + 6..i as usize * ROW_BYTES + 8]
                    != [0; 2]
                || (i == 0 && offset != 0)
                || (i > 0 && offset <= previous_offset)
                || offset >= payload_len as u64
            {
                return Err(corrupt("invalid logical key or vector offset"));
            }
            previous_key = Some(key);
            previous_offset = offset;
        }
        if count == 0 && payload_len != 0 {
            return Err(corrupt("empty directory has a payload"));
        }
        #[cfg(feature = "native")]
        result.advise(libc::MADV_RANDOM);
        Ok(result)
    }

    pub(crate) fn len(&self) -> u32 {
        (self.rows.len() / ROW_BYTES) as u32
    }

    pub(crate) fn key(&self, index: u32) -> LogicalUnit {
        let row = &self.rows.as_slice()[index as usize * ROW_BYTES..][..ROW_BYTES];
        LogicalUnit {
            doc: u32::from_le_bytes(row[..4].try_into().unwrap()),
            ordinal: u16::from_le_bytes(row[4..6].try_into().unwrap()),
        }
    }

    fn offset(&self, index: u32) -> u64 {
        if index == self.len() {
            return self.payload.len() as u64;
        }
        let start = index as usize * ROW_BYTES + 8;
        u64::from_le_bytes(self.rows.as_slice()[start..start + 8].try_into().unwrap())
    }

    fn lower_bound(&self, target: LogicalUnit) -> u32 {
        let mut low = 0;
        let mut high = self.len();
        while low < high {
            let mid = low + (high - low) / 2;
            if self.key(mid) < target {
                low = mid + 1;
            } else {
                high = mid;
            }
        }
        low
    }

    pub(crate) fn find(&self, target: LogicalUnit) -> Option<u32> {
        let index = self.lower_bound(target);
        (index < self.len() && self.key(index) == target).then_some(index)
    }

    pub(crate) fn for_document(&self, doc: u32) -> impl Iterator<Item = (u16, u32)> + '_ {
        let start = self.lower_bound(LogicalUnit { doc, ordinal: 0 });
        (start..self.len())
            .map(|i| (self.key(i), i))
            .take_while(move |(key, _)| key.doc == doc)
            .map(|(key, i)| (key.ordinal, i))
    }

    fn vector_range(&self, index: u32) -> Result<std::ops::Range<usize>> {
        if index >= self.len() {
            return Err(corrupt("vector index out of bounds"));
        }
        Ok(self.offset(index) as usize..self.offset(index + 1) as usize)
    }

    /// Directory-only admission: do not fault or validate payload before its
    /// bytes have been charged to the candidate scoring request.
    pub(crate) fn vector_byte_len(&self, index: u32) -> Result<u64> {
        Ok(self.vector_range(index)?.len() as u64)
    }

    /// Borrow encoded extents after directory lookup. Contents are checked only
    /// by the explicit vector() integrity entry point; this does not prefetch.
    pub(crate) fn encoded_vector(&self, index: u32) -> Result<&[u8]> {
        Ok(&self.payload.as_slice()[self.vector_range(index)?])
    }

    /// Query scoring trusts writer-produced payload contents. Admitted extents
    /// retain safe slice bounds; explicit integrity checks use vector().
    pub(crate) fn vector_for_scoring(&self, index: u32) -> Result<ForwardVector<'_>> {
        Ok(ForwardVector(self.encoded_vector(index)?))
    }

    /// Validate only the selected vector's payload, before returning its view.
    #[cfg(any(feature = "native", feature = "wasm", test))]
    pub(crate) fn vector(&self, index: u32) -> Result<ForwardVector<'_>> {
        let vector = self.vector_for_scoring(index)?;
        vector.validate(self.dims)?;
        Ok(vector)
    }

    pub(crate) fn encoded_bytes(&self) -> usize {
        self.payload.len() + self.rows.len() + TRAILER_BYTES
    }
    #[cfg(feature = "native")]
    pub(crate) fn payload_bytes(&self) -> usize {
        self.payload.len()
    }

    #[cfg(any(feature = "native", feature = "wasm", test))]
    pub(crate) fn validate_payload(
        &self,
        check_cancel: &(impl Fn() -> Result<()> + Sync),
    ) -> Result<ValidatedForward<'_>> {
        for i in 0..self.len() {
            check_cancel()?;
            self.vector(i)?;
        }
        Ok(ValidatedForward(self))
    }

    #[cfg(feature = "native")]
    pub(crate) fn advise(&self, advice: i32) {
        self.rows.madvise(advice);
        self.payload.madvise(advice);
    }
}

#[cfg(any(feature = "native", feature = "wasm", test))]
pub(crate) fn write_disabled(writer: &mut dyn std::io::Write) -> std::io::Result<u64> {
    use byteorder::{LittleEndian, WriteBytesExt};
    writer.write_u32::<LittleEndian>(0)?;
    writer.write_u32::<LittleEndian>(STORAGE_DISABLED)?;
    writer.write_u64::<LittleEndian>(0)?;
    Ok(TRAILER_BYTES as u64)
}

/// A background scan validates once before BP's repeated infallible passes.
#[cfg(any(feature = "native", feature = "wasm", test))]
pub(crate) struct ValidatedForward<'a>(&'a BmpForward);
#[cfg(any(feature = "native", feature = "wasm", test))]
impl ValidatedForward<'_> {
    pub(crate) fn vector(&self, index: u32) -> ForwardVector<'_> {
        ForwardVector(
            &self.0.payload.as_slice()
                [self.0.offset(index) as usize..self.0.offset(index + 1) as usize],
        )
    }
}

#[cfg(any(feature = "native", feature = "wasm", test))]
pub(crate) fn write_directory(
    writer: &mut dyn std::io::Write,
    rows: impl IntoIterator<Item = std::io::Result<(LogicalUnit, u64)>>,
    count: u32,
    payload_bytes: u64,
) -> std::io::Result<u64> {
    use byteorder::{LittleEndian, WriteBytesExt};
    let mut written = 0u32;
    let mut previous = None;
    let mut previous_offset = 0;
    for row in rows {
        let (key, offset) = row?;
        if previous.is_some_and(|p| p >= key)
            || (written == 0 && offset != 0)
            || (written > 0 && offset <= previous_offset)
            || offset >= payload_bytes
        {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "invalid BMP forward directory",
            ));
        }
        writer.write_u32::<LittleEndian>(key.doc)?;
        writer.write_u16::<LittleEndian>(key.ordinal)?;
        writer.write_u16::<LittleEndian>(0)?;
        writer.write_u64::<LittleEndian>(offset)?;
        written = written
            .checked_add(1)
            .ok_or_else(|| std::io::Error::other("BMP forward count overflow"))?;
        previous = Some(key);
        previous_offset = offset;
    }
    if written != count {
        return Err(std::io::Error::other("BMP forward count mismatch"));
    }
    writer.write_u32::<LittleEndian>(count)?;
    writer.write_u32::<LittleEndian>(0)?;
    writer.write_u64::<LittleEndian>(payload_bytes)?;
    Ok(u64::from(count) * ROW_BYTES as u64 + TRAILER_BYTES as u64)
}

#[cfg(feature = "native")]
mod rewrite;
#[cfg(feature = "native")]
pub(crate) use rewrite::{validate_copy_sources, write_forward_sources};

#[cfg(all(test, feature = "native"))]
mod tests;

/// Copy retained forward payloads and rewrite only their logical directory.
#[cfg(feature = "native")]
pub(crate) fn write_compacted_forward(
    bmp: &super::BmpIndex,
    rows: &super::row_map::RowMap,
    writer: &mut super::OffsetWriter,
    cancellation: Option<&std::sync::atomic::AtomicBool>,
) -> Result<()> {
    use std::io::Write;
    let Some(forward) = bmp.forward() else {
        return write_disabled(writer).map(|_| ()).map_err(Error::from);
    };
    let mut count = 0u32;
    let start = writer.offset();
    let mut i = 0;
    while i < forward.len() {
        if cancellation.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Acquire)) {
            return Err(Error::IndexClosed);
        }
        if rows.get(forward.key(i).doc).is_none() {
            i += 1;
            continue;
        }
        let from = i;
        // Bound both label scanning and writes even for a fully live field.
        let end = i.saturating_add(4096).min(forward.len());
        while i < end && rows.get(forward.key(i).doc).is_some() {
            i += 1;
        }
        for chunk in forward.payload.as_slice()
            [forward.offset(from) as usize..forward.offset(i) as usize]
            .chunks(4 * 1024 * 1024)
        {
            if cancellation.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Acquire)) {
                return Err(Error::IndexClosed);
            }
            writer.write_all(chunk)?;
        }
        count += i - from;
    }
    let payload_len = writer.offset() - start;
    let mut offset = 0u64;
    let records = (0..forward.len()).filter_map(|i| {
        if i.is_multiple_of(4096)
            && cancellation.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Acquire))
        {
            return Some(Err(std::io::Error::new(
                std::io::ErrorKind::Interrupted,
                "BMP forward compaction cancelled",
            )));
        }
        let key = forward.key(i);
        rows.get(key.doc).map(|doc| {
            let at = offset;
            offset += forward.offset(i + 1) - forward.offset(i);
            Ok((
                LogicalUnit {
                    doc,
                    ordinal: key.ordinal,
                },
                at,
            ))
        })
    });
    write_directory(writer, records, count, payload_len)?;
    Ok(())
}