haematite-cli 0.21.7

//! Stored embedding declarations, companion headers, model tags and keys.
//!
//! A `w\0` value in decimal is a field as it was before declarations: its
//! vectors are supplied by the caller and stored bare on the branch, with that
//! length. A field declared embedded has its vectors made by the job and kept
//! in the branch's companion namespace, never on the branch (Q1); each value
//! there is headed by the version of the record it was made from and the tag
//! of the model that made it.

use std::num::NonZeroUsize;

use haematite::tree::Hash;
use serde::{Deserialize, Serialize};

use super::Failure;
use super::vectors::{decode, parse_dimension};
use crate::serve::protocol::{Metric, Version};

/// Each branch's companion is the namespace `@haem/embed/{branch}`.
pub const COMPANION_PREFIX: &str = "@haem/embed/";
/// The only header format: format, generation, sequence, tag.
pub const FORMAT: u8 = 1;
pub const TAG_BYTES: usize = 16;
pub const HEADER_BYTES: usize = 1 + 8 + 8 + TAG_BYTES;

/// A model as a declaration names it, and as its worker's Ready names it.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelName {
    pub name: String,
    pub revision: String,
    pub manifest_sha256: String,
}

/// The kinds a field may be declared as; dense first (design item "Dense
/// vectors first").
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum VectorKind {
    Dense,
}

/// A field declared embedded: which text field its vectors are made from,
/// by which model, of what length, and the measure a search uses.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Declaration {
    pub kind: VectorKind,
    pub source: String,
    pub model: ModelName,
    pub dimensions: NonZeroUsize,
    pub distance: Metric,
}

/// What a field's `w\0` value says it is.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Field {
    /// Written before declarations, or supplied by the caller: bare vectors
    /// on the branch, of this length.
    Supplied { dimensions: usize },
    /// Made by the job and kept in the companion.
    Embedded(Declaration),
}

/// The field a `w\0` value describes: decimal is supplied, as every value
/// written before declarations is; anything else is a declaration or
/// refused by name.
///
/// # Errors
/// `vector_dimension_corrupt` for decimal that is not a canonical positive
/// length; `vector_declaration_corrupt` for a declaration that does not
/// decode.
pub fn read_field(bytes: &[u8]) -> Result<Field, Failure> {
    if bytes.first().is_some_and(u8::is_ascii_digit) {
        return parse_dimension(bytes).map(|dimensions| Field::Supplied { dimensions });
    }
    serde_json::from_slice(bytes)
        .map(Field::Embedded)
        .map_err(|error| Failure::new("vector_declaration_corrupt", error.to_string()))
}

/// The `w\0` value of `declaration`.
///
/// # Errors
/// Whatever encoding it answers.
pub fn encode_declaration(declaration: &Declaration) -> Result<Vec<u8>, Failure> {
    serde_json::to_vec(declaration).map_err(Failure::storage)
}

/// The first 16 bytes of the store's hash of the model's name, revision and
/// manifest, as JSON so no two models share the bytes hashed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelTag(pub [u8; TAG_BYTES]);

/// The tag of `model`.
///
/// # Errors
/// Whatever encoding the name answers.
pub fn model_tag(model: &ModelName) -> Result<ModelTag, Failure> {
    let named = serde_json::to_vec(model).map_err(Failure::storage)?;
    let hash = Hash::of(&named);
    let tag = hash
        .as_bytes()
        .first_chunk::<TAG_BYTES>()
        .ok_or_else(|| Failure::new("embed_tag_short", "the store's hash is under 16 bytes"))?;
    Ok(ModelTag(*tag))
}

/// What an embedded vector was made from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Provenance {
    /// The record's revision (`m\0` `written_at`) its text was read at.
    pub made_from: Version,
    pub tag: ModelTag,
}

/// A companion value: the header, then the members as little-endian f32s,
/// as a bare vector stores them.
#[must_use]
pub fn encode_embedded(provenance: &Provenance, vector: &[f32]) -> Vec<u8> {
    let capacity = vector.len().saturating_mul(4).saturating_add(HEADER_BYTES);
    let mut bytes = Vec::with_capacity(capacity);
    bytes.push(FORMAT);
    bytes.extend_from_slice(&provenance.made_from.generation.to_be_bytes());
    bytes.extend_from_slice(&provenance.made_from.sequence.to_be_bytes());
    bytes.extend_from_slice(&provenance.tag.0);
    for member in vector {
        bytes.extend_from_slice(&member.to_le_bytes());
    }
    bytes
}

/// A companion value's provenance and vector of `dimensions` members.
///
/// # Errors
/// `vector_corrupt` for a value shorter than its header, of another format,
/// or whose members are not `dimensions` finite numbers.
pub fn decode_embedded(bytes: &[u8], dimensions: usize) -> Result<(Provenance, Vec<f32>), Failure> {
    let short = || {
        Failure::new(
            "vector_corrupt",
            "an embedded vector is shorter than its header",
        )
    };
    let (&format, rest) = bytes.split_first().ok_or_else(short)?;
    if format != FORMAT {
        return Err(Failure::new(
            "vector_corrupt",
            format!("an embedded vector has format {format}, not {FORMAT}"),
        ));
    }
    let (generation, rest) = rest.split_first_chunk::<8>().ok_or_else(short)?;
    let (sequence, rest) = rest.split_first_chunk::<8>().ok_or_else(short)?;
    let (tag, members) = rest.split_first_chunk::<TAG_BYTES>().ok_or_else(short)?;
    let provenance = Provenance {
        made_from: Version {
            generation: u64::from_be_bytes(*generation),
            sequence: u64::from_be_bytes(*sequence),
        },
        tag: ModelTag(*tag),
    };
    Ok((provenance, decode(members, dimensions)?))
}

/// The job's field position, committed with its vectors in the companion.
///
/// Naming the branch lets a fork resume from its parent's position.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Position {
    pub branch: String,
    /// Every change to `branch` up to and including this version has its
    /// vector or its removal; `None` before the first walk ends, when the
    /// walk starts from the empty tree.
    pub through: Option<Version>,
    /// The walk under way, if one is.
    pub toward: Option<Toward>,
}

/// A walk of the difference from a position's `through` to `version`, in key
/// order: every record up to and including `after` is finished.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Toward {
    pub version: Version,
    /// A record id: the walk's order is its keys' order, `m\0` then the
    /// collection then the id, so ids compare as bytes.
    pub after: String,
}

/// The position a companion value holds.
///
/// # Errors
/// `embed_position_corrupt` for a value that does not decode.
pub fn read_position(bytes: &[u8]) -> Result<Position, Failure> {
    serde_json::from_slice(bytes)
        .map_err(|error| Failure::new("embed_position_corrupt", error.to_string()))
}

/// The companion value of `position`.
///
/// # Errors
/// Whatever encoding it answers.
pub fn encode_position(position: &Position) -> Result<Vec<u8>, Failure> {
    serde_json::to_vec(position).map_err(Failure::storage)
}

/// The companion namespace of `branch`.
#[must_use]
pub fn companion_name(branch: &str) -> String {
    format!("{COMPANION_PREFIX}{branch}")
}

/// The companion key of a field's position.
#[must_use]
pub fn position_key(collection: &str, field: &str) -> Vec<u8> {
    format!("e\0{collection}\0{field}").into_bytes()
}

/// The companion key of a record's vector, laid out as a bare vector's key on
/// the branch is.
#[must_use]
pub fn vector_key(collection: &str, field: &str, id: &str) -> Vec<u8> {
    format!("v\0{collection}\0{field}\0{id}").into_bytes()
}

/// EMB-E2 (j): the companion key that marks a field's open walk a rewalk.
///
/// A rewalk reads from the empty tree and lists only texts with no vector.
/// The key is not under `e\0`, so a scan of positions never reads it.
#[must_use]
pub fn rewalk_key(collection: &str, field: &str) -> Vec<u8> {
    format!("er\0{collection}\0{field}").into_bytes()
}

/// EMB-E2: the most windows a record may hold, whatever the setting; a stored
/// count past it is corrupt, so clearing a record's windows stays bounded.
pub const WINDOWS_MOST: u32 = 1024;

/// EMB-E2: the companion prefix of every window of a field.
#[must_use]
pub fn window_prefix(collection: &str, field: &str) -> Vec<u8> {
    format!("vw\0{collection}\0{field}\0").into_bytes()
}

/// EMB-E2: the companion key of window `index` of a record's text.
///
/// The key is the field's window prefix, the id, a NUL and the index
/// big-endian, so a record's windows sort together and in text order. Ids
/// hold no NUL.
#[must_use]
pub fn window_key(collection: &str, field: &str, id: &str, index: u32) -> Vec<u8> {
    let mut key = window_prefix(collection, field);
    key.extend_from_slice(id.as_bytes());
    key.push(0);
    key.extend_from_slice(&index.to_be_bytes());
    key
}

/// EMB-E2: the record id and window index of a key under `prefix`.
///
/// # Errors
/// `vector_corrupt` for a key that is not a window key under `prefix`.
pub fn window_of(prefix: &[u8], key: &[u8]) -> Result<(String, u32), Failure> {
    let corrupt = || Failure::new("vector_corrupt", "a window key names no record and index");
    let rest = key.strip_prefix(prefix).ok_or_else(corrupt)?;
    let (named, index) = rest.split_last_chunk::<4>().ok_or_else(corrupt)?;
    let id = named.strip_suffix(&[0]).ok_or_else(corrupt)?;
    let id = String::from_utf8(id.to_vec()).map_err(|_not_utf8| corrupt())?;
    Ok((id, u32::from_be_bytes(*index)))
}

/// EMB-E2: where a window lies in its record's text, `[start, end)` in
/// bytes, and how many windows the record has.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Span {
    pub count: u32,
    pub start: u64,
    pub end: u64,
}

const SPAN_BYTES: usize = 20;

/// EMB-E2: a window's value: its [`Span`], then its provenance and vector as
/// [`encode_embedded`] lays them out.
#[must_use]
pub fn encode_window(span: Span, provenance: &Provenance, vector: &[f32]) -> Vec<u8> {
    let embedded = encode_embedded(provenance, vector);
    let mut bytes = Vec::with_capacity(SPAN_BYTES.saturating_add(embedded.len()));
    bytes.extend_from_slice(&span.count.to_be_bytes());
    bytes.extend_from_slice(&span.start.to_be_bytes());
    bytes.extend_from_slice(&span.end.to_be_bytes());
    bytes.extend_from_slice(&embedded);
    bytes
}

/// EMB-E2: a window's span, provenance and vector of `dimensions` members.
///
/// # Errors
/// `vector_corrupt` for a value shorter than its span, a count past
/// [`WINDOWS_MOST`], or as [`decode_embedded`].
pub fn decode_window(
    bytes: &[u8],
    dimensions: usize,
) -> Result<(Span, Provenance, Vec<f32>), Failure> {
    let (span, rest) = split_span(bytes)?;
    let (provenance, vector) = decode_embedded(rest, dimensions)?;
    Ok((span, provenance, vector))
}

/// EMB-E2: a window's span alone, without decoding its vector.
///
/// # Errors
/// As [`decode_window`].
pub fn window_span(bytes: &[u8]) -> Result<Span, Failure> {
    split_span(bytes).map(|(span, _)| span)
}

fn split_span(bytes: &[u8]) -> Result<(Span, &[u8]), Failure> {
    let short = || {
        Failure::new(
            "vector_corrupt",
            "an embedded window is shorter than its span",
        )
    };
    let (count, rest) = bytes.split_first_chunk::<4>().ok_or_else(short)?;
    let (start, rest) = rest.split_first_chunk::<8>().ok_or_else(short)?;
    let (end, rest) = rest.split_first_chunk::<8>().ok_or_else(short)?;
    let span = Span {
        count: u32::from_be_bytes(*count),
        start: u64::from_be_bytes(*start),
        end: u64::from_be_bytes(*end),
    };
    if span.count > WINDOWS_MOST || span.end <= span.start {
        return Err(Failure::new(
            "vector_corrupt",
            "an embedded window names no span or too many windows",
        ));
    }
    Ok((span, rest))
}

#[cfg(test)]
#[path = "embedded_tests.rs"]
mod tests;