uqa-storage 0.4.5

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

use super::{
    identity::{row, Resolver},
    setup, FIELD, TABLE,
};
use crate::{
    diskann_index::{
        build::{
            DiskANNGenerationOptions, DiskANNMergeOptions, DiskANNPartitionOptions,
            DiskANNTemporaryBudget,
        },
        format::{DiskANNGeneration, PAGE_BYTES},
        pages::{DiskANNPageSource, DiskANNReadLimits},
        DiskANNIndexOptions, DiskANNPersistentOwner, DiskANNQueryMetadata, PQTrainingOptions,
        PersistentDiskANNIndex,
    },
    key_value::{
        conformance::{expect, expect_eq},
        KeyValueDiskANNCanonical, KeyValueVectorIndex,
    },
    read_control::StorageReadControl,
    vector_index::DiskANNIndexParams,
    CatalogFacade, KeyValueCatalog, KeyValueStore, StorageBackendResult, VectorIndex,
};
use std::sync::Arc;

mod statistics;
pub use statistics::verify_diskann_vector_statistics;

mod populations;
pub use populations::{
    seed_diskann_population_upgrade, verify_diskann_population_late_publication,
    verify_diskann_population_lifetimes, verify_diskann_population_upgrade,
    verify_diskann_population_upgrade_reopen, verify_diskann_population_writers,
};
mod retirement;
pub use retirement::{verify_diskann_runtime_retirement, verify_diskann_runtime_retirement_reopen};
mod reclamation;
pub use reclamation::{
    verify_diskann_runtime_reclaimed_reopen, verify_diskann_runtime_reclamation,
};

/// Small deterministic provider acceptance settings, never runtime defaults.
pub fn diskann_runtime_fixture_options(
    dimensions: u32,
) -> StorageBackendResult<DiskANNIndexOptions> {
    let training = PQTrainingOptions {
        max_samples: 8,
        max_centroids: 2,
        max_iterations: 2,
        seed: 42,
    };
    Ok(DiskANNIndexOptions {
        parameters: DiskANNIndexParams::for_dimensions(dimensions)?,
        read: DiskANNReadLimits {
            resident_bytes: 65_536,
            cache_bytes: PAGE_BYTES,
            max_in_flight_page_bytes: 2 * PAGE_BYTES,
            max_record_bytes: 8192,
        },
        partitions: DiskANNPartitionOptions {
            max_partition_points: 8,
            coarse_training: PQTrainingOptions {
                max_centroids: 3,
                ..training
            },
            max_depth: 0,
        },
        merge: DiskANNMergeOptions {
            sort_buffer_records: 8,
        },
        generation: DiskANNGenerationOptions {
            training,
            code_batch_nodes: 4,
            side_batch_entries: 4,
            max_record_bytes: 8192,
        },
    })
}

fn canonical(store: &Arc<dyn KeyValueStore>) -> StorageBackendResult<KeyValueDiskANNCanonical> {
    KeyValueDiskANNCanonical::new(store.clone(), TABLE, FIELD, 2)
}

pub(super) fn runtime(
    store: &Arc<dyn KeyValueStore>,
    temporary: &DiskANNTemporaryBudget,
    control: &StorageReadControl,
) -> StorageBackendResult<impl VectorIndex> {
    let options = diskann_runtime_fixture_options(2)?;
    PersistentDiskANNIndex::new(
        canonical(store)?.bind(
            row([91; 16])?.relation,
            Arc::new(Resolver),
            options.read,
            control,
        )?,
        options,
        temporary,
    )
}

fn generation(
    store: &Arc<dyn KeyValueStore>,
    control: &StorageReadControl,
) -> StorageBackendResult<DiskANNGeneration> {
    canonical(store)?
        .retain_for_index(&row([91; 16])?.relation, control)?
        .selected_source(&Resolver, control)?
        .map(|source| source.generation())
        .ok_or_else(|| {
            crate::StorageBackendError::Other(
                "runtime fixture has no selected DiskANN generation".into(),
            )
        })
}

fn scores(index: &dyn VectorIndex, expected: &[(u64, f64)]) -> StorageBackendResult<()> {
    let postings = index.search_knn(&[1.0, 0.0], 10)?;
    expect_eq(
        &postings
            .iter()
            .map(|entry| (entry.doc_id, entry.payload.score.to_bits()))
            .collect::<Vec<_>>(),
        &expected
            .iter()
            .map(|&(id, score)| (id, score.to_bits()))
            .collect::<Vec<_>>(),
        "runtime literal complete tensor scores",
    )
}

fn population(
    index: &dyn VectorIndex,
    current: u64,
    changed: u64,
    control: &StorageReadControl,
) -> StorageBackendResult<()> {
    expect_eq(
        &index
            .diskann_query_metadata(control)?
            .expect("DiskANN metadata")
            .canonical_counts,
        &Some(crate::diskann_index::DiskANNCanonicalCounts::new(
            current, changed,
        )?),
        "runtime metadata preserves complete canonical populations",
    )
}

fn capture_metadata(
    store: &Arc<dyn KeyValueStore>,
    index: &dyn VectorIndex,
    held: &dyn VectorIndex,
    temporary: &DiskANNTemporaryBudget,
    control: &StorageReadControl,
) -> StorageBackendResult<DiskANNQueryMetadata> {
    let expected = index
        .diskann_query_metadata(control)?
        .expect("DiskANN metadata");
    require_metadata(held, &expected, control)?;
    expect_eq(
        &expected.manifest.input().coverage.vector_count(),
        &4,
        "metadata complete base tensor count",
    )?;
    let mut options = diskann_runtime_fixture_options(2)?;
    options.read.resident_bytes = 0;
    let handle = canonical(store)?.bind(
        row([91; 16])?.relation,
        Arc::new(Resolver),
        options.read,
        control,
    )?;
    let before = control.memory().used();
    let metadata = handle.query_metadata()?;
    expect_eq(
        &metadata.manifest,
        &expected.manifest,
        "bounded metadata selects the actual manifest",
    )?;
    expect_eq(
        &metadata.corpus_fingerprint,
        &expected.corpus_fingerprint,
        "bounded metadata selects the actual private field",
    )?;
    expect_eq(
        &metadata.read_limits,
        &options.read,
        "metadata retains requested read limits",
    )?;
    expect_eq(
        &control.memory().used(),
        &before,
        "metadata releases temporary provider readers",
    )?;
    expect(
        PersistentDiskANNIndex::new(handle, options, temporary).is_err(),
        "executable restore still requires resident PQ admission",
    )?;
    Ok(expected)
}

fn require_metadata(
    index: &dyn VectorIndex,
    expected: &DiskANNQueryMetadata,
    control: &StorageReadControl,
) -> StorageBackendResult<()> {
    expect_eq(
        &index.diskann_query_metadata(control)?.as_ref(),
        &Some(expected),
        "selected metadata matches the original captured view",
    )
}

fn changed_metadata(
    index: &dyn VectorIndex,
    held: &dyn VectorIndex,
    expected: &DiskANNQueryMetadata,
    control: &StorageReadControl,
) -> StorageBackendResult<()> {
    let changed = index
        .diskann_query_metadata(control)?
        .expect("changed metadata");
    expect_eq(
        &changed.manifest,
        &expected.manifest,
        "row changes preserve physical metadata",
    )?;
    expect(
        changed.corpus_fingerprint != expected.corpus_fingerprint,
        "private changes update live canonical identity",
    )?;
    require_metadata(held, expected, control)
}

fn create(
    store: &Arc<dyn KeyValueStore>,
    temporary: &DiskANNTemporaryBudget,
    control: &StorageReadControl,
) -> StorageBackendResult<()> {
    setup(store)?;
    let catalog = KeyValueCatalog::new(store.clone());
    let options = diskann_runtime_fixture_options(2)?;
    let definition = row([91; 16])?;
    catalog.save_catalog_index_row(&definition)?;
    let mut raw = KeyValueVectorIndex::new(store.clone(), TABLE, FIELD, 2);
    raw.add(1, vec![1.0, -0.0])?;
    raw.add_many(2, vec![vec![0.0, 1.0], vec![1.0, 0.0]])?;
    raw.add(3, vec![-1.0, 0.0])?;
    expect(
        runtime(store, temporary, control).is_err(),
        "restore never creates a missing head",
    )?;
    expect(
        canonical(store)?
            .create_index(&definition.relation, &Resolver, options, temporary, control)
            .is_err(),
        "creation requires the caller transaction",
    )?;
    store.begin_transaction()?;
    store.put(b"runtime-outer-write", b"kept")?;
    catalog.save_catalog_index_row(&definition)?;
    let mut bad = options;
    bad.generation.max_record_bytes = 1;
    expect(
        canonical(store)?
            .create_index(&definition.relation, &Resolver, bad, temporary, control)
            .is_err(),
        "failed construction rolls back adoption",
    )?;
    let mut unreadable = options;
    unreadable.read.resident_bytes = 1;
    expect(
        canonical(store)?
            .create_index(
                &definition.relation,
                &Resolver,
                unreadable,
                temporary,
                control,
            )
            .is_err(),
        "creation rejects a sealed generation that cannot be admitted for queries",
    )?;
    expect_eq(
        &raw.count()?,
        &4,
        "failed creation preserves original raw ordinals",
    )?;
    expect(
        canonical(store)?
            .retain(control)?
            .origin(1, control)
            .is_err(),
        "failed creation leaves the original unstamped raw field",
    )?;
    expect_eq(
        &store.get(b"runtime-outer-write")?,
        &Some(b"kept".to_vec()),
        "structural undo keeps earlier caller effects",
    )?;
    canonical(store)?.create_index(&definition.relation, &Resolver, options, temporary, control)?;
    expect(
        store.in_transaction(),
        "construction never commits the outer transaction",
    )?;
    Ok(())
}

/// Exercise real first construction, bounded legacy adoption, rebuild and empty-generation publication through the public `VectorIndex` interface, including private undo and failure atomicity.
pub fn verify_diskann_runtime_lifecycle(
    store: &Arc<dyn KeyValueStore>,
) -> StorageBackendResult<DiskANNGeneration> {
    let control = StorageReadControl::with_limit(1 << 21);
    let temporary = DiskANNTemporaryBudget::new(1 << 20);
    let options = diskann_runtime_fixture_options(2)?;
    let definition = row([91; 16])?;
    create(store, &temporary, &control)?;
    let mut index = runtime(store, &temporary, &control)?;
    scores(&index, &[(1, 1.0), (2, 1.0), (3, -1.0)])?;
    expect_eq(&index.count()?, &4, "adoption preserves complete tensors")?;
    let first = generation(store, &control)?;
    population(&index, 4, 0, &control)?;
    verify_diskann_vector_statistics(&index, first, (3, 4), (0, 0))?;
    let held = index.snapshot()?;
    let metadata = capture_metadata(store, &index, &*held, &temporary, &control)?;
    store.savepoint("runtime-user")?;
    index.add(1, vec![-1.0, 0.0])?;
    index.delete(2)?;
    index.add(4, vec![0.0, 1.0])?;
    changed_metadata(&index, &*held, &metadata, &control)?;
    population(&index, 3, 2, &control)?;
    verify_diskann_vector_statistics(&index, first, (1, 1), (2, 2))?;
    verify_diskann_vector_statistics(&*held, first, (3, 4), (0, 0))?;
    population(&*held, 4, 0, &control)?;
    expect_eq(
        &generation(store, &control)?,
        &first,
        "ordinary writes preserve the base",
    )?;
    index.initialize()?;
    expect(
        generation(store, &control)? != first,
        "initialize publishes a replacement base",
    )?;
    let rebuilt = index.snapshot()?;
    population(&index, 3, 0, &control)?;
    scores(&*rebuilt, &[(1, -1.0), (3, -1.0), (4, 0.0)])?;
    index.clear()?;
    let empty = index.snapshot()?;
    population(&index, 0, 0, &control)?;
    verify_diskann_vector_statistics(&*empty, generation(store, &control)?, (0, 0), (0, 0))?;
    expect_eq(
        &empty.count()?,
        &0,
        "clear publishes empty canonical membership",
    )?;
    let empty_source = canonical(store)?
        .retain_for_index(&definition.relation, &control)?
        .into_vector_index(&Resolver, options.read, &control)?
        .ok_or_else(|| crate::StorageBackendError::Other("missing empty runtime head".into()))?;
    expect_eq(
        &empty_source.manifest().input().nodes,
        &0,
        "clear publishes an actually empty graph",
    )?;
    store.rollback_to_savepoint("runtime-user")?;
    require_metadata(&index, &metadata, &control)?;
    population(&*rebuilt, 3, 0, &control)?;
    population(&*empty, 0, 0, &control)?;
    scores(&index, &[(1, 1.0), (2, 1.0), (3, -1.0)])?;
    scores(&*rebuilt, &[(1, -1.0), (3, -1.0), (4, 0.0)])?;
    expect_eq(&empty.count()?, &0, "undone empty snapshot remains valid")?;
    expect_eq(
        &generation(store, &control)?,
        &first,
        "savepoint restores the original head",
    )?;
    store.commit_transaction()?;
    expect(
        index.clear().is_err(),
        "clear requires a caller transaction",
    )?;
    expect(
        index.initialize().is_err(),
        "rebuild requires a caller transaction",
    )?;
    store.begin_transaction()?;
    reject_failed_clear(store, &temporary, &control)?;
    scores(&index, &[(1, 1.0), (2, 1.0), (3, -1.0)])?;
    expect_eq(
        &generation(store, &control)?,
        &first,
        "failed clear preserves the original head",
    )?;
    index.clear()?;
    index.add(5, vec![1.0, 0.0])?;
    index.initialize()?;
    store.commit_transaction()?;
    scores(&index, &[(5, 1.0)])?;
    population(&index, 1, 0, &control)?;
    scores(&*held, &[(1, 1.0), (2, 1.0), (3, -1.0)])?;
    let final_generation = generation(store, &control)?;
    verify_diskann_vector_statistics(&index, final_generation, (1, 1), (0, 0))?;
    verify_diskann_vector_statistics(&*held, first, (3, 4), (0, 0))?;
    expect_eq(
        &temporary.used(),
        &0,
        "finished builds release encrypted temporary input",
    )?;
    Ok(final_generation)
}

fn reject_failed_clear(
    store: &Arc<dyn KeyValueStore>,
    temporary: &DiskANNTemporaryBudget,
    control: &StorageReadControl,
) -> StorageBackendResult<()> {
    let mut bad = diskann_runtime_fixture_options(2)?;
    bad.generation.max_record_bytes = 1;
    let mut constrained = PersistentDiskANNIndex::new(
        canonical(store)?.bind(
            row([91; 16])?.relation,
            Arc::new(Resolver),
            bad.read,
            control,
        )?,
        bad,
        temporary,
    )?;
    expect(
        constrained.clear().is_err(),
        "failed empty build restores its own clear",
    )
}

/// Reopen the actual provider after all original owners have been dropped.
pub fn verify_diskann_runtime_reopen(
    store: &Arc<dyn KeyValueStore>,
    expected: DiskANNGeneration,
) -> StorageBackendResult<()> {
    let control = StorageReadControl::with_limit(1 << 20);
    let temporary = DiskANNTemporaryBudget::new(1 << 20);
    let mut index = runtime(store, &temporary, &control)?;
    expect_eq(
        &generation(store, &control)?,
        &expected,
        "cold restore keeps its actual generation",
    )?;
    scores(&index, &[(5, 1.0)])?;
    index.add(6, vec![0.0, 1.0])?;
    population(&index, 2, 1, &control)?;
    scores(&index, &[(5, 1.0), (6, 0.0)])?;
    expect_eq(
        &generation(store, &control)?,
        &expected,
        "cold writes do not rebuild",
    )
}

/// A constructor adopts one complete raw namespace. A racing legacy insertion must conflict even when its key did not exist in the captured source.
pub fn verify_diskann_runtime_adoption_conflicts(
    store: &Arc<dyn KeyValueStore>,
) -> StorageBackendResult<()> {
    setup(store)?;
    let catalog = KeyValueCatalog::new(store.clone());
    let definition = row([91; 16])?;
    catalog.save_catalog_index_row(&definition)?;
    let peer = store.open_session()?;
    let options = diskann_runtime_fixture_options(2)?;
    let temporary = DiskANNTemporaryBudget::new(1 << 20);
    let control = StorageReadControl::with_limit(1 << 21);
    verify_adoption_gaps(store, &control)?;
    let mut raw = KeyValueVectorIndex::new(store.clone(), TABLE, FIELD, 2);
    let mut late = KeyValueVectorIndex::new(peer.clone(), TABLE, FIELD, 2);
    for creator_first in [false, true] {
        raw.clear()?;
        raw.add(1, vec![1.0, 0.0])?;
        peer.begin_transaction()?;
        late.add(99, vec![0.0, 1.0])?;
        store.begin_transaction()?;
        canonical(store)?.create_index(
            &definition.relation,
            &Resolver,
            options,
            &temporary,
            &control,
        )?;
        if creator_first {
            store.commit_transaction()?;
            expect(
                peer.commit_transaction().is_err(),
                "older unstamped insertion conflicts with completed adoption",
            )?;
            peer.rollback_transaction()?;
            scores(&runtime(store, &temporary, &control)?, &[(1, 1.0)])?;
        } else {
            peer.commit_transaction()?;
            expect(
                store.commit_transaction().is_err(),
                "adoption conflicts with a newly committed raw ordinal",
            )?;
            store.rollback_transaction()?;
            expect_eq(
                &raw.count()?,
                &2,
                "failed adoption preserves the winning raw insertion",
            )?;
            expect(
                canonical(store)?
                    .retain_for_index(&definition.relation, &control)?
                    .selected_source(&Resolver, &control)?
                    .is_none(),
                "failed adoption publishes no head",
            )?;
        }
    }
    Ok(())
}

fn verify_adoption_gaps(
    store: &Arc<dyn KeyValueStore>,
    control: &StorageReadControl,
) -> StorageBackendResult<()> {
    let mut raw = KeyValueVectorIndex::new(store.clone(), TABLE, FIELD, 2);
    raw.add(1, vec![1.0, 0.0])?;
    raw.add_many(2, vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![-1.0, 0.0]])?;
    let peer = store.open_session()?;
    for (missing, expected) in [
        (
            0,
            "invalid persisted vector ordinal sequence for document 2: expected 0, found 1",
        ),
        (
            1,
            "invalid persisted vector ordinal sequence for document 2: expected 1, found 2",
        ),
    ] {
        let key = crate::key_value::codec::vector_key(TABLE, FIELD, 2, missing)?;
        let value = store.get(&key)?.expect("corrupt fixture starts complete");
        store.delete(&key)?;
        store.begin_transaction()?;
        store.put(b"before-rejected-adoption", b"kept")?;
        let before = store.scan_prefix(b"")?;
        let result = canonical(store)?.create_index(
            &row([91; 16])?.relation,
            &Resolver,
            diskann_runtime_fixture_options(2)?,
            &DiskANNTemporaryBudget::new(1 << 20),
            control,
        );
        expect_eq(
            &result.as_ref().err().map(std::string::ToString::to_string),
            &Some(expected.to_owned()),
            "the shared canonical decoder rejects ordinal gaps during adoption",
        )?;
        expect_eq(
            &store.scan_prefix(b"")?,
            &before,
            "failed adoption preserves all outer private records",
        )?;
        expect(
            peer.scan_prefix_keys_after(b"\0uqa-diskann-v1\0", None, 1)?
                .is_empty(),
            "rejected adoption allocates no physical generation or data identity",
        )?;
        store.rollback_transaction()?;
        store.put(&key, &value)?;
    }
    Ok(())
}