haematite-cli 0.21.8

//! V1: the embedding path's stored shapes.

use std::error::Error;
use std::num::NonZeroUsize;

use super::{
    Declaration, FORMAT, Field, HEADER_BYTES, ModelName, ModelTag, Position, Provenance, Toward,
    VectorKind, companion_name, decode_embedded, encode_declaration, encode_embedded,
    encode_position, model_tag, position_key, read_field, read_position, vector_key,
};
use crate::serve::protocol::{Metric, Version};

type TestResult = Result<(), Box<dyn Error>>;

fn model(manifest: &str) -> ModelName {
    ModelName {
        name: "stand-in".to_owned(),
        revision: "1".to_owned(),
        manifest_sha256: manifest.to_owned(),
    }
}

const fn provenance() -> Provenance {
    Provenance {
        made_from: Version {
            generation: 7,
            sequence: 9,
        },
        tag: ModelTag([5; 16]),
    }
}

fn code<T>(outcome: Result<T, super::Failure>) -> Result<String, Box<dyn Error>> {
    Ok(outcome
        .err()
        .ok_or("a value was read that should be refused")?
        .code)
}

/// Fails while a value written before declarations is read as anything but
/// a supplied field of its length; a length that is not canonical is still
/// refused by name.
#[test]
fn a_decimal_field_is_supplied_with_its_length() -> TestResult {
    assert_eq!(read_field(b"768")?, Field::Supplied { dimensions: 768 });
    assert_eq!(code(read_field(b"0768"))?, "vector_dimension_corrupt");
    Ok(())
}

/// A declaration is read back as written, and its JSON names the kind and
/// the measure as the plan gives them.
#[test]
fn a_declaration_round_trips() -> TestResult {
    let declaration = Declaration {
        kind: VectorKind::Dense,
        source: "body".to_owned(),
        model: model("abc"),
        dimensions: NonZeroUsize::new(2).ok_or("two is not zero")?,
        distance: Metric::Cosine,
    };
    let bytes = encode_declaration(&declaration)?;
    let text = String::from_utf8(bytes.clone())?;
    assert!(text.contains(r#""kind":"dense""#), "{text}");
    assert!(text.contains(r#""distance":"cosine""#), "{text}");
    assert_eq!(read_field(&bytes)?, Field::Embedded(declaration));
    assert_eq!(
        code(read_field(br#"{"kind":"dense"}"#))?,
        "vector_declaration_corrupt"
    );
    Ok(())
}

/// Fails while the header's layout moves: one format byte, the generation
/// and sequence big-endian, the tag, then the members as a bare vector
/// stores them.
#[test]
fn the_header_is_laid_out_as_ruled() {
    let bytes = encode_embedded(&provenance(), &[1.5]);
    let mut expected = vec![FORMAT];
    expected.extend_from_slice(&7_u64.to_be_bytes());
    expected.extend_from_slice(&9_u64.to_be_bytes());
    expected.extend_from_slice(&[5; 16]);
    expected.extend_from_slice(&1.5_f32.to_le_bytes());
    assert_eq!(bytes.len(), HEADER_BYTES + 4);
    assert_eq!(bytes, expected);
}

/// A value is read back with the provenance and members it was written with.
#[test]
fn an_embedded_value_round_trips() -> TestResult {
    let bytes = encode_embedded(&provenance(), &[1.5, -2.0]);
    let (read, vector) = decode_embedded(&bytes, 2)?;
    assert_eq!(read, provenance());
    assert_eq!(vector, [1.5, -2.0]);
    Ok(())
}

/// Fails while a value is read past what it is: shorter than its header, of
/// another format, of another length, or a bare vector, each is
/// `vector_corrupt`.
#[test]
fn a_short_or_foreign_value_is_corrupt() -> TestResult {
    let bytes = encode_embedded(&provenance(), &[1.5, -2.0]);
    let short = bytes.get(..HEADER_BYTES - 1).ok_or("no header")?;
    assert_eq!(code(decode_embedded(short, 0))?, "vector_corrupt");
    let mut foreign = bytes.clone();
    if let Some(format) = foreign.first_mut() {
        *format = FORMAT + 1;
    }
    assert_eq!(code(decode_embedded(&foreign, 2))?, "vector_corrupt");
    assert_eq!(code(decode_embedded(&bytes, 3))?, "vector_corrupt");
    let bare: Vec<u8> = [1.5_f32, -2.0]
        .iter()
        .flat_map(|member| member.to_le_bytes())
        .collect();
    assert_eq!(code(decode_embedded(&bare, 2))?, "vector_corrupt");
    Ok(())
}

/// Fails while the tag leaves part of the model out: models that differ only
/// in their manifest are told apart, and one model always has one tag.
#[test]
fn models_differing_only_in_manifest_have_different_tags() -> TestResult {
    assert_ne!(model_tag(&model("abc"))?, model_tag(&model("abd"))?);
    assert_eq!(model_tag(&model("abc"))?, model_tag(&model("abc"))?);
    Ok(())
}

/// The position is read back as written; the companion and its keys are
/// named as the plan gives them.
#[test]
fn the_position_and_the_companion_keys() -> TestResult {
    let position = Position {
        branch: "main".to_owned(),
        through: None,
        toward: Some(Toward {
            version: provenance().made_from,
            after: "n-1".to_owned(),
        }),
    };
    assert_eq!(read_position(&encode_position(&position)?)?, position);
    assert_eq!(code(read_position(b"7"))?, "embed_position_corrupt");
    assert_eq!(companion_name("main"), "@haem/embed/main");
    assert_eq!(position_key("notes", "body_vec"), b"e\0notes\0body_vec");
    assert_eq!(
        vector_key("notes", "body_vec", "n-1"),
        b"v\0notes\0body_vec\0n-1"
    );
    Ok(())
}

/// One fixed model: the stand-in's name and revision, and a manifest of 64
/// zeros.
fn pinned_model() -> ModelName {
    ModelName {
        name: "stand-in".to_owned(),
        revision: "1".to_owned(),
        manifest_sha256: "0".repeat(64),
    }
}

/// The stored tag of `pinned_model`, as the first run computed it on Dean's
/// laptop (5 October 2026).
const PINNED_TAG: [u8; 16] = [
    105, 255, 15, 224, 125, 238, 195, 77, 97, 52, 156, 9, 122, 202, 170, 138,
];

/// Fails if the tag of a stored model changes: a field reordered in
/// `ModelName`, or the store's hash changing under it, would make every
/// stored vector another model's.
#[test]
fn the_tag_of_a_fixed_model_is_pinned() -> TestResult {
    assert_eq!(model_tag(&pinned_model())?, ModelTag(PINNED_TAG));
    Ok(())
}

/// Fails if one whole stored value changes: the header and members of a
/// vector made from version 7.9 by `pinned_model`, byte for byte.
#[test]
fn one_encoded_value_is_pinned() -> TestResult {
    let provenance = Provenance {
        made_from: provenance().made_from,
        tag: model_tag(&pinned_model())?,
    };
    let mut expected = vec![1, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 9];
    expected.extend_from_slice(&PINNED_TAG);
    expected.extend_from_slice(&[0, 0, 0xc0, 0x3f, 0, 0, 0, 0xc0]);
    assert_eq!(encode_embedded(&provenance, &[1.5, -2.0]), expected);
    Ok(())
}