uqa-storage 0.3.7

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

use super::*;
use uqa_analysis::analyzer::standard_analyzer;

fn fields<const N: usize>(pairs: [(&str, &str); N]) -> BTreeMap<FieldName, String> {
    pairs
        .into_iter()
        .map(|(k, v)| (k.to_string(), v.to_string()))
        .collect()
}

#[test]
fn add_document_indexes_tokens() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "The Rust Programming Language")]))
        .unwrap();
    idx.add_document(2, fields([("title", "Programming with Rust")]))
        .unwrap();

    let pl = idx.get_posting_list("title", "rust").unwrap();
    let docs: Vec<_> = pl.doc_ids().collect();
    assert_eq!(docs, vec![1, 2]);

    // standard analyzer stems "programming" -> "program"
    let pl2 = idx.get_posting_list("title", "program").unwrap();
    let docs2: Vec<_> = pl2.doc_ids().collect();
    assert_eq!(docs2, vec![1, 2]);
}

#[test]
fn snapshot_creation_does_not_scale_with_document_count() {
    let allocations = [1, 1024].map(|count| {
        let mut index = MemoryInvertedIndex::new(uqa_analysis::whitespace_analyzer());
        for id in 0..count {
            index
                .add_document(id, fields([("body", "snapshot payload repeated tokens")]))
                .unwrap();
        }
        let mut snapshots = None;
        let allocation = allocation_counter::measure(|| {
            snapshots = Some((
                index.snapshot().unwrap(),
                index.writable_snapshot().unwrap(),
            ));
        });
        let (read, writable) = snapshots.unwrap();
        index.clear().unwrap();
        assert_eq!(read.doc_count().unwrap(), count);
        assert_eq!(writable.doc_count().unwrap(), count);
        allocation
    });
    assert_eq!(allocations[0].bytes_total, allocations[1].bytes_total);
    assert_eq!(allocations[0].count_total, allocations[1].count_total);
}

#[test]
fn writable_snapshots_keep_postings_counters_and_revisions_independent() {
    let mut index = MemoryInvertedIndex::new(uqa_analysis::whitespace_analyzer());
    index
        .add_document(1, fields([("body", "old shared")]))
        .unwrap();
    index.add_document(2, fields([("body", "shared")])).unwrap();
    let read = index.snapshot().unwrap();
    let mut writable = index.writable_snapshot().unwrap();
    let original_revision = index.search_analyzer_revision("body").unwrap();
    let key = TokenTermKey::from_text("shared");
    let original_occurrences = index.get_occurrences(1, "body", &key).unwrap();

    writable.remove_document(2).unwrap();
    writable
        .try_add_documents(vec![
            (1, fields([("body", "next next")])),
            (3, fields([("body", "next")])),
        ])
        .unwrap();
    writable
        .set_field_analyzer("body", standard_analyzer("english"), AnalyzerPhase::Search)
        .unwrap();
    index.add_document(1, fields([("body", "source")])).unwrap();

    assert_eq!(read.doc_freq("body", "shared").unwrap(), 2);
    assert_eq!(read.total_field_length("body").unwrap(), 3);
    assert_eq!(
        read.get_occurrences(1, "body", &key).unwrap(),
        original_occurrences
    );
    assert!(Arc::ptr_eq(
        &read.search_analyzer_revision("body").unwrap(),
        &original_revision
    ));
    assert_eq!(writable.doc_freq("body", "next").unwrap(), 2);
    assert_eq!(writable.total_field_length("body").unwrap(), 3);
    assert_eq!(writable.field_doc_count("body").unwrap(), 2);
    assert!(writable
        .indexed_field_metadata(2, "body")
        .unwrap()
        .is_none());
    assert_eq!(index.doc_freq("body", "shared").unwrap(), 1);
    assert_eq!(index.doc_freq("body", "next").unwrap(), 0);
    assert!(Arc::ptr_eq(
        &index.search_analyzer_revision("body").unwrap(),
        &original_revision
    ));

    let second = writable.snapshot().unwrap();
    writable.clear().unwrap();
    writable
        .add_document(9, fields([("body", "reused")]))
        .unwrap();
    index.clear().unwrap();
    assert_eq!(second.doc_freq("body", "next").unwrap(), 2);
    assert_eq!(read.doc_freq("body", "shared").unwrap(), 2);
    assert_eq!(writable.doc_count().unwrap(), 1);
    assert_eq!(writable.doc_freq("body", "reused").unwrap(), 1);
}

#[test]
fn doc_freq_counts_documents() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "rust")])).unwrap();
    idx.add_document(2, fields([("title", "rust rust rust")]))
        .unwrap();
    idx.add_document(3, fields([("title", "go")])).unwrap();

    assert_eq!(idx.doc_freq("title", "rust").unwrap(), 2);
    assert_eq!(idx.doc_freq("title", "go").unwrap(), 1);
    assert_eq!(idx.doc_freq("title", "java").unwrap(), 0);
}

#[test]
fn term_freq_counts_positions() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "rust rust rust")]))
        .unwrap();
    // After standard analyzer: ["rust", "rust", "rust"]
    assert_eq!(idx.get_term_freq(1, "title", "rust").unwrap(), 3);
}

#[test]
fn doc_length_tracks_token_count() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    // standard analyzer drops "the" / "is" stop words
    idx.add_document(1, fields([("title", "the rust language is fast")]))
        .unwrap();
    // Remaining tokens: ["rust", "languag", "fast"] -> 3
    assert_eq!(idx.get_doc_length(1, "title").unwrap(), 3);
}

#[test]
fn remove_document_clears_postings_and_lengths() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "rust")])).unwrap();
    idx.add_document(2, fields([("title", "rust")])).unwrap();
    idx.remove_document(1).unwrap();

    assert_eq!(idx.doc_freq("title", "rust").unwrap(), 1);
    assert_eq!(idx.get_doc_length(1, "title").unwrap(), 0);
    assert_eq!(idx.doc_count().unwrap(), 1);
}

#[test]
fn replacing_doc_replaces_postings() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "rust")])).unwrap();
    idx.add_document(1, fields([("title", "go")])).unwrap();

    assert_eq!(idx.doc_freq("title", "rust").unwrap(), 0);
    assert_eq!(idx.doc_freq("title", "go").unwrap(), 1);
    assert_eq!(idx.doc_count().unwrap(), 1);
}

#[test]
fn empty_field_map_removes_existing_index_document() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "rust")])).unwrap();
    idx.add_document(1, BTreeMap::new()).unwrap();
    idx.add_document(2, BTreeMap::new()).unwrap();

    assert_eq!(idx.doc_count().unwrap(), 0);
    assert_eq!(idx.doc_freq("title", "rust").unwrap(), 0);
    assert!(!idx.state.doc_terms.contains_key(&1));
    assert!(!idx.state.doc_terms.contains_key(&2));
}

#[test]
fn stats_avg_doc_length_correct() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "rust language")]))
        .unwrap();
    idx.add_document(2, fields([("title", "rust")])).unwrap();
    let s = idx.stats().unwrap();
    assert_eq!(s.total_docs, 2);
    // 2 + 1 = 3 tokens / 2 docs = 1.5
    assert!((s.avg_doc_length - 1.5).abs() < 1e-9);
    assert_eq!(s.doc_freq("title", "rust"), 2);
}

#[test]
fn token_position_format_requires_a_representable_end() {
    let mut occurrence = uqa_core::TokenOccurrence {
        position: u32::MAX - 1,
        position_length: 1,
        offsets: None,
    };
    occurrence.validate().unwrap();
    occurrence.position = u32::MAX;
    assert!(occurrence.validate().is_err());
}

#[test]
fn add_overflow_does_not_partially_insert_document() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    Arc::make_mut(&mut idx.state).doc_count = u64::MAX;

    let error = idx
        .add_document(7, fields([("title", "rust")]))
        .unwrap_err();
    assert!(error.to_string().contains("document count"));
    assert_eq!(idx.state.doc_count, u64::MAX);
    assert!(!idx.state.doc_terms.contains_key(&7));
    assert!(idx.get_posting_list("title", "rust").unwrap().is_empty());
}

#[test]
fn field_length_overflow_preserves_existing_document() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "rust")])).unwrap();
    Arc::make_mut(&mut idx.state)
        .total_length
        .insert("title".into(), u64::MAX);

    let error = idx.add_document(2, fields([("title", "go")])).unwrap_err();
    assert!(error.to_string().contains("total field length"));
    assert_eq!(idx.doc_count().unwrap(), 1);
    assert_eq!(idx.doc_freq("title", "rust").unwrap(), 1);
    assert_eq!(idx.doc_freq("title", "go").unwrap(), 0);
    assert!(!idx.state.doc_terms.contains_key(&2));
}

#[test]
fn corrupt_counter_rejects_remove_without_mutating_postings() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    idx.add_document(1, fields([("title", "rust")])).unwrap();
    Arc::make_mut(&mut idx.state)
        .total_length
        .insert("title".into(), 0);

    let error = idx.remove_document(1).unwrap_err();
    assert!(error.to_string().contains("total field length"));
    assert_eq!(idx.doc_count().unwrap(), 1);
    assert_eq!(idx.doc_freq("title", "rust").unwrap(), 1);
    assert_eq!(idx.get_doc_length(1, "title").unwrap(), 1);
}

#[test]
fn stats_reports_cross_field_total_overflow() {
    let mut idx = MemoryInvertedIndex::new(standard_analyzer("english"));
    Arc::make_mut(&mut idx.state).doc_count = 1;
    Arc::make_mut(&mut idx.state)
        .total_length
        .insert("a".into(), u64::MAX);
    Arc::make_mut(&mut idx.state)
        .total_length
        .insert("b".into(), 1);

    let error = idx.stats().unwrap_err();
    assert!(error.to_string().contains("total document length"));
}