uqa-storage 0.5.3

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Evaluated IVF candidates preserve canonical mutation semantics and caller resource limits.

use super::state::{StoredVector, VectorKey};
use super::{IVFIndex, IVFMetadataSnapshot, IVFMutation, IVFState};
use crate::{read_control::StorageReadControl, StorageBackendError, VectorIndex};
use uqa_core::memory::MemoryError;

fn trained() -> IVFIndex {
    let mut index = IVFIndex::with_params(3, 2, 2, 2);
    for doc in 1..=10 {
        index
            .add(doc, vec![doc as f32, 11.0 - doc as f32, 0.5])
            .unwrap();
    }
    index.train().unwrap();
    index
}

#[test]
fn controlled_capture_preserves_training_metadata_and_defers_stale_retraining_to_search() {
    let mut untrained = IVFIndex::with_params(3, 2, 1, 4);
    untrained.add(1, vec![1.0, 0.0, 0.0]).unwrap();
    let mut stale = trained();
    for document in 1..=3 {
        stale.delete(document).unwrap();
    }
    assert_eq!(stale.state(), IVFState::Stale);
    for source in [untrained, trained(), stale] {
        source.set_nprobe(1);
        let before = source.metadata_snapshot();
        let control = StorageReadControl::with_limit(1 << 20);
        let snapshot = source.snapshot_controlled(&control).unwrap();
        assert_eq!(snapshot.metadata_snapshot(), before);
        assert_eq!(snapshot.nprobe(), 1);
        assert_eq!(
            *snapshot.inverted_lists.lock(),
            *source.inverted_lists.lock()
        );
        let expected = source.detached_clone();
        assert_eq!(
            snapshot.search_knn(&[1.0, 0.0, 0.0], 2).unwrap(),
            expected.search_knn(&[1.0, 0.0, 0.0], 2).unwrap()
        );
        assert_eq!(snapshot.metadata_snapshot(), expected.metadata_snapshot());
        assert_eq!(source.metadata_snapshot(), before);
        drop(source);
        assert_eq!(control.memory().used(), snapshot.reserved_bytes());
        drop(snapshot);
        assert_eq!(control.memory().used(), 0);
    }
}

#[test]
fn candidates_preserve_centroids_counters_and_the_source_generation() {
    let source = trained();
    let before = source.metadata_snapshot();
    let replacement = [vec![0.0, 0.0, 1.0], vec![0.0, 1.0, 0.0]];
    for mutation in [
        IVFMutation::Replace {
            document: 2,
            vectors: &replacement,
        },
        IVFMutation::Replace {
            document: 2,
            vectors: &[],
        },
        IVFMutation::Delete(2),
        IVFMutation::Delete(100),
        IVFMutation::Clear,
        IVFMutation::Train,
    ] {
        let control = StorageReadControl::with_limit(1 << 20);
        let prepared = source.prepare_metadata(mutation, &control).unwrap();
        let mut expected = source.detached_clone();
        match mutation {
            IVFMutation::Replace { document, vectors } => {
                expected.add_many(document, vectors.to_vec()).unwrap();
            }
            IVFMutation::Delete(document) => expected.delete(document).unwrap(),
            IVFMutation::Clear => expected.clear().unwrap(),
            IVFMutation::Train => expected.train().unwrap(),
        }
        if expected.state() == IVFState::Stale {
            expected.train().unwrap();
        }
        assert_eq!(*prepared, expected.metadata_snapshot());
        assert_eq!(source.metadata_snapshot(), before);
        assert_eq!(control.memory().used(), prepared.reserved_bytes());
        drop(prepared);
        assert_eq!(control.memory().used(), 0);
    }
}

#[test]
fn failed_candidate_returns_its_allowance_and_preserves_the_source() {
    let source = trained();
    let before = source.metadata_snapshot();
    let control = StorageReadControl::with_limit(1);
    let allocation = allocation_counter::measure(|| {
        assert!(matches!(
            source.prepare_metadata(IVFMutation::Train, &control),
            Err(StorageBackendError::Memory(MemoryError::Limit { .. }))
        ));
    });
    assert_eq!(
        allocation.bytes_total, 0,
        "the allowance must reject the candidate before its corpus is cloned"
    );
    assert_eq!(control.memory().used(), 0);
    let empty = source
        .prepare_metadata(IVFMutation::Clear, &control)
        .unwrap();
    assert_eq!(empty.vector_count, 0);
    assert_eq!(control.memory().used(), 0);
    let control = StorageReadControl::with_limit(1 << 20);
    let invalid = [vec![f32::NAN; 3]];
    assert!(source
        .prepare_metadata(
            IVFMutation::Replace {
                document: 1,
                vectors: &invalid
            },
            &control
        )
        .is_err());
    assert_eq!(control.memory().used(), 0);
    control.cancellation().cancel();
    assert!(matches!(
        source.prepare_metadata(IVFMutation::Train, &control),
        Err(StorageBackendError::Cancelled(_))
    ));
    assert_eq!(control.memory().used(), 0);
    assert_eq!(source.metadata_snapshot(), before);
}

#[test]
fn prepared_deletion_retrains_only_after_crossing_the_stale_threshold() {
    let mut source = trained();
    source.delete(1).unwrap();
    source.delete(2).unwrap();
    let before = source.metadata_snapshot();
    assert_eq!(before.deletes_since_train, 2);
    let control = StorageReadControl::with_limit(1 << 20);
    let empty = source
        .prepare_metadata(
            IVFMutation::Replace {
                document: 3,
                vectors: &[],
            },
            &control,
        )
        .unwrap();
    assert_eq!(empty.trained_size, 10);
    assert_eq!(empty.deletes_since_train, 2);
    assert_eq!(empty.centroids, before.centroids);
    let deleted = source
        .prepare_metadata(IVFMutation::Delete(3), &control)
        .unwrap();
    assert_eq!(deleted.trained_size, 7);
    assert_eq!(deleted.deletes_since_train, 0);
    assert_eq!(deleted.vector_count, 7);
    assert_eq!(source.metadata_snapshot(), before);
}

#[test]
fn restoration_rejects_missing_tensor_ordinals_and_accepts_unsorted_complete_tensors() {
    let mut index = IVFIndex::with_params(2, 2, 2, 100);
    index
        .add_many(1, vec![vec![1.0, 0.0], vec![0.0, 1.0]])
        .unwrap();
    for ordinals in [[0, 2], [1, 2], [0, 0], [1, 0]] {
        let vectors = ordinals
            .into_iter()
            .map(|ordinal| (1, ordinal, vec![1.0, 0.0]))
            .collect();
        let restored = IVFIndex::from_persistence(2, 2, 2, 100, vectors, index.metadata_snapshot());
        assert_eq!(restored.is_ok(), ordinals == [1, 0]);
    }
}

#[test]
fn controlled_canonical_build_preserves_incremental_training_and_explicit_rebuild() {
    let vectors = vec![
        (1, 0, vec![1.0, 0.0]),
        (1, 1, vec![0.0, 1.0]),
        (2, 0, vec![0.5, 0.5]),
        (3, 0, vec![-1.0, 0.0]),
    ];
    for threshold in [2, 100] {
        let params = crate::vector_index::IVFIndexParams {
            nlist: 2,
            nprobe: 2,
            train_threshold: threshold,
        };
        let mut expected = IVFIndex::with_params(2, 2, 2, threshold);
        for document in vectors.chunk_by(|left, right| left.0 == right.0) {
            expected
                .add_many(
                    document[0].0,
                    document
                        .iter()
                        .map(|(_, _, vector)| vector.clone())
                        .collect(),
                )
                .unwrap();
        }
        let control = StorageReadControl::with_limit(1 << 20);
        let captured = IVFIndex::from_canonical_controlled(2, params, &vectors, &control).unwrap();
        assert_eq!(captured.metadata_snapshot(), expected.metadata_snapshot());
        assert!(control.memory().used() > 0);
        expected.initialize().unwrap();
        let rebuilt = IVFIndex::prepare_canonical(2, params, &vectors, &control).unwrap();
        assert_eq!(*rebuilt, expected.metadata_snapshot());
        drop((captured, rebuilt));
        assert_eq!(control.memory().used(), 0);
        assert!(matches!(
            IVFIndex::from_canonical_controlled(
                2,
                params,
                &vectors,
                &StorageReadControl::with_limit(0)
            ),
            Err(StorageBackendError::Memory(_))
        ));
        control.cancellation().cancel();
        assert!(matches!(
            IVFIndex::from_canonical_controlled(2, params, &vectors, &control),
            Err(StorageBackendError::Cancelled(_))
        ));
        assert_eq!(control.memory().used(), 0);
    }
}

#[test]
fn controlled_training_retains_a_detached_snapshot_and_preserves_stale_source_on_failure() {
    let mut source = trained();
    for document in 1..=3 {
        source.delete(document).unwrap();
    }
    assert_eq!(source.state(), IVFState::Stale);
    let before = source.metadata_snapshot();
    assert!(matches!(
        source.trained_snapshot_controlled(&StorageReadControl::with_limit(0)),
        Err(StorageBackendError::Memory(_))
    ));
    assert_eq!(source.metadata_snapshot(), before);
    let control = StorageReadControl::with_limit(1 << 20);
    let candidate = source.trained_snapshot_controlled(&control).unwrap();
    let expected = source.detached_clone();
    expected.train().unwrap();
    assert_eq!(candidate.metadata_snapshot(), expected.metadata_snapshot());
    assert_eq!(source.metadata_snapshot(), before);
    assert_eq!(control.memory().used(), candidate.reserved_bytes());
    drop(source);
    assert_eq!(
        candidate.search_knn(&[1.0, 0.0, 0.0], 2).unwrap(),
        expected.search_knn(&[1.0, 0.0, 0.0], 2).unwrap()
    );
    drop(candidate);
    assert_eq!(control.memory().used(), 0);
}

#[test]
fn stale_query_training_uses_its_own_workspace_and_preserves_the_source_on_rejection() {
    let mut source = trained();
    for document in 1..=3 {
        source.delete(document).unwrap();
    }
    let before = source.metadata_snapshot();
    assert_eq!(before.state, IVFState::Stale);
    let control = StorageReadControl::with_limit(1 << 20);
    let snapshot = source.snapshot_controlled(&control).unwrap();
    let retained = control.memory().used();
    // Normalization fits, but the per-query training workspace must be separately reserved.
    let held = control
        .memory()
        .reserve(control.memory().limit() - retained - 64)
        .unwrap();
    assert!(matches!(
        snapshot.search_knn_with_control(&[1.0, 0.0, 0.0], 2, &control),
        Err(StorageBackendError::Memory(_))
    ));
    assert_eq!(snapshot.metadata_snapshot(), before);
    assert_eq!(source.metadata_snapshot(), before);
    assert_eq!(control.memory().used(), control.memory().limit() - 64);
    drop(held);
    control.cancellation().cancel();
    assert!(matches!(
        snapshot.search_knn_with_control(&[1.0, 0.0, 0.0], 2, &control),
        Err(StorageBackendError::Cancelled(_))
    ));
    assert_eq!(snapshot.metadata_snapshot(), before);
    assert_eq!(control.memory().used(), retained);
    control.cancellation().reset();
    let expected = source.detached_clone();
    assert_eq!(
        snapshot
            .search_knn_with_control(&[1.0, 0.0, 0.0], 2, &control)
            .unwrap(),
        expected.search_knn(&[1.0, 0.0, 0.0], 2).unwrap()
    );
    assert_eq!(snapshot.metadata_snapshot(), expected.metadata_snapshot());
    assert_eq!(source.metadata_snapshot(), before);
    assert_eq!(control.memory().used(), retained);
    drop(snapshot);
    assert_eq!(control.memory().used(), 0);
}

#[test]
fn restored_generation_releases_reconstruction_scratch_before_preparing_a_write() {
    const DIMENSIONS: usize = 1024;
    const DOCUMENTS: usize = 2338;
    let vectors = (1..=DOCUMENTS)
        .map(|document| {
            let mut vector = vec![0.0; DIMENSIONS];
            vector[document % 100] = 1.0;
            (document as u64, 0, vector)
        })
        .collect();
    let metadata = IVFMetadataSnapshot {
        state: IVFState::Trained,
        centroids: (0..100)
            .map(|coordinate| {
                let mut centroid = vec![0.0; DIMENSIONS];
                centroid[coordinate] = 1.0;
                centroid
            })
            .collect(),
        assignments: (1..=DOCUMENTS)
            .map(|document| (document as u64, 0, document % 100))
            .collect(),
        trained_size: DOCUMENTS,
        deletes_since_train: 0,
        vector_count: DOCUMENTS,
    };
    let control = StorageReadControl::with_limit(64 * 1024 * 1024);
    let restored = IVFIndex::restore_controlled(
        DIMENSIONS as u32,
        crate::vector_index::IVFIndexParams::default(),
        vectors,
        metadata,
        &control,
    )
    .unwrap();
    let retained = control.memory().used();
    assert_eq!(retained, restored.reserved_bytes());
    assert!(retained < control.memory().peak());
    assert!(retained >= DOCUMENTS * DIMENSIONS * size_of::<f32>() * 2);
    let replacement = [vec![1.0; DIMENSIONS]];
    let prepared = restored
        .prepare_metadata(
            IVFMutation::Replace {
                document: DOCUMENTS as u64 + 1,
                vectors: &replacement,
            },
            &control,
        )
        .unwrap();
    assert_eq!(prepared.vector_count, DOCUMENTS + 1);
    assert_eq!(prepared.trained_size, DOCUMENTS);
    assert_eq!(restored.metadata_snapshot().vector_count, DOCUMENTS);
    assert_eq!(
        control.memory().used(),
        retained + prepared.reserved_bytes()
    );
    assert!(control.memory().peak() <= control.memory().limit());
    drop(prepared);
    assert_eq!(control.memory().used(), retained);
    drop(restored);
    assert_eq!(control.memory().used(), 0);
}

#[test]
fn restored_generation_keeps_spare_vector_and_centroid_capacity_reserved() {
    let mut vector = Vec::with_capacity(4096);
    vector.extend([1.0, 0.0]);
    let mut centroid = Vec::with_capacity(4096);
    centroid.extend([1.0, 0.0]);
    let retained_payload = (vector.capacity() + centroid.capacity()) * size_of::<f32>();
    let control = StorageReadControl::with_limit(1 << 20);
    let restored = IVFIndex::restore_controlled(
        2,
        crate::vector_index::IVFIndexParams::default(),
        vec![(1, 0, vector)],
        IVFMetadataSnapshot {
            state: IVFState::Trained,
            centroids: vec![centroid],
            assignments: vec![(1, 0, 0)],
            trained_size: 1,
            deletes_since_train: 0,
            vector_count: 1,
        },
        &control,
    )
    .unwrap();
    assert!(restored.reserved_bytes() > retained_payload);
    assert_eq!(control.memory().used(), restored.reserved_bytes());
    let result = restored.search_knn(&[1.0, 0.0], 1).unwrap();
    assert_eq!(result.len(), 1);
    drop(restored);
    assert_eq!(control.memory().used(), 0);
}

#[test]
fn restored_stale_generation_reserves_resident_capacity_after_query_training() {
    let mut source = trained();
    for document in 1..=3 {
        source.delete(document).unwrap();
    }
    let mut metadata = source.metadata_snapshot();
    metadata.centroids.truncate(1);
    for (_, _, cluster) in &mut metadata.assignments {
        *cluster = 0;
    }
    let vectors = source
        .vectors
        .lock()
        .values()
        .map(|vector| {
            (
                vector.doc_id,
                vector.vector_ordinal,
                vector.raw_vector.clone(),
            )
        })
        .collect();
    let control = StorageReadControl::with_limit(1 << 20);
    let restored = IVFIndex::restore_controlled(
        3,
        crate::vector_index::IVFIndexParams {
            nlist: 7,
            nprobe: 7,
            train_threshold: 2,
        },
        vectors,
        metadata,
        &control,
    )
    .unwrap();
    let retained = restored.reserved_bytes();
    restored
        .search_knn_with_control(&[1.0, 0.0, 0.0], 2, &control)
        .unwrap();
    assert_eq!(restored.state(), IVFState::Trained);
    assert_eq!(restored.centroids.lock().len(), 7);
    assert_eq!(control.memory().used(), retained);
    let vectors = restored.vectors.lock();
    let centroids = restored.centroids.lock();
    let lists = restored.inverted_lists.lock();
    let payload = vectors.len() * size_of::<(VectorKey, StoredVector)>()
        + vectors
            .values()
            .map(|vector| {
                (vector.raw_vector.capacity() + vector.vector.capacity()) * size_of::<f32>()
            })
            .sum::<usize>()
        + centroids.capacity() * size_of::<Vec<f32>>()
        + centroids
            .iter()
            .map(|centroid| centroid.capacity() * size_of::<f32>())
            .sum::<usize>()
        + lists.capacity() * size_of::<Vec<(u64, u32)>>()
        + lists
            .iter()
            .map(|list| list.capacity() * size_of::<(u64, u32)>())
            .sum::<usize>();
    assert!(payload <= retained);
    drop((vectors, centroids, lists));
    drop(restored);
    assert_eq!(control.memory().used(), 0);
}