infino 0.5.6

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Infino Authors

//! Convenience builders for test fixtures.
//!
//! Three test contexts share these helpers:
//!
//! - **Unit tests** (`#[cfg(test)] mod tests` inside `src/`)
//!   reach this module via `crate::test_helpers::...` —
//!   `cfg(test)` always enables it.
//! - **Integration tests** (`tests/...`) reach it via
//!   `infino::test_helpers::...` — the `test-helpers` Cargo
//!   feature is auto-enabled by the `dev-dependencies` self-
//!   reference in `Cargo.toml`.
//! - **Benches** (`benches/...`) reach it the same way.
//!
//! Scope: small atomic idioms that repeat across dozens of
//! test / bench fixtures (Decimal128 id construction, default
//! tokenizer, default vector config). Higher-level "build a
//! test corpus" / "build a full test superfile" stays in the
//! test files themselves — those vary too much per scenario
//! to share usefully.
//!
//! [`brute_force_bm25`] is the textbook BM25 reference impl
//! used as the FTS correctness oracle.

pub mod brute_force_bm25;
pub mod cas_conformance;
pub mod fault_storage;

/// Observability probe for the vector query path's EFFECTIVE served
/// shortlist budget — the regression guard for the serve-the-law scope
/// bug (#520 review): recall-floor tests are insensitive to a re-shadowed
/// `options` (the `rm=256` constant yields MORE survivors, so recall
/// stays equal-or-better and only latency regresses), but the served
/// budget is not. The pooled warm arm records the limit and cell floor
/// it passes to the global shortlist; tests read them after a query.
/// Observability probes for the hidden vector path's two routing stages
/// (#515 diagnosis): the 1-bit admit window (which cells get exact fine
/// rescoring) and the exact-fine cell ranking selection serves from.
/// Append logs — drain between queries; race-tolerant under test
/// parallelism (a concurrent test adds entries, never removes yours).
pub mod admit_trace {
    use std::sync::Mutex;

    static ADMITS: Mutex<Vec<Vec<u32>>> = Mutex::new(Vec::new());
    static FINES: Mutex<Vec<Vec<(u32, f32)>>> = Mutex::new(Vec::new());

    /// The admit shortlist (global cell ids) chosen by the 1-bit window.
    pub fn record_admit(cells: Vec<u32>) {
        ADMITS.lock().expect("admit probe lock").push(cells);
    }

    /// The exact-fine cell ranking (cell id, best exact fine score).
    pub fn record_fine(ranked: Vec<(u32, f32)>) {
        FINES.lock().expect("fine probe lock").push(ranked);
    }

    /// Drain both logs recorded since the last drain.
    #[allow(clippy::type_complexity)]
    pub fn drain() -> (Vec<Vec<u32>>, Vec<Vec<(u32, f32)>>) {
        (
            std::mem::take(&mut *ADMITS.lock().expect("admit probe lock")),
            std::mem::take(&mut *FINES.lock().expect("fine probe lock")),
        )
    }
}

pub mod served_shortlist_probe {
    use std::sync::Mutex;

    static RECORDS: Mutex<Vec<(usize, usize)>> = Mutex::new(Vec::new());

    /// Called by the query path (test-helpers builds only).
    pub fn record(limit: usize, cell_floor: usize) {
        RECORDS
            .lock()
            .expect("shortlist probe lock")
            .push((limit, cell_floor));
    }

    /// Drain every `(limit, cell_floor)` recorded since the last drain.
    /// Append-log semantics keep the probe race-tolerant under the test
    /// binary's parallelism: a concurrent test adds tuples but can never
    /// remove this test's — assert with `contains`, not equality.
    pub fn drain() -> Vec<(usize, usize)> {
        std::mem::take(&mut *RECORDS.lock().expect("shortlist probe lock"))
    }
}

use std::{collections::HashSet, path::Path, sync::Arc};

use arrow_array::{Decimal128Array, LargeStringArray, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use rayon::ThreadPoolBuilder;

use crate::{
    storage::StorageProvider,
    superfile::{
        builder::{FtsConfig, VectorConfig},
        fts::tokenize::{AsciiLowerTokenizer, Tokenizer},
        vector::{distance::Metric, rerank_codec::RerankCodec},
    },
    supertable::{
        SupertableOptions,
        reader_cache::{ColdFetchMode, DiskCacheConfig, DiskCacheStore, LruPolicy},
    },
};

/// 1 GiB disk-cache budget for tests.
const TEST_DISK_CACHE_BUDGET_BYTES: u64 = 1 << 30;
/// Parallel cold-fetch streams for the test disk cache.
const TEST_COLD_FETCH_STREAMS: usize = 4;
/// Cold-fetch range chunk (1 MiB) for the test disk cache.
const TEST_COLD_FETCH_CHUNK_BYTES: u64 = 1 << 20;

/// Build a `DiskCacheStore` with the standard test config: 1 GiB budget,
/// hybrid-with-prefetch cold fetch, mmap sweep timers disabled, LRU eviction,
/// CRC-on-open, and a no-op pin set (pinning is a perf optimization, not a
/// correctness requirement — an `Arc<SuperfileReader>` keeps the mmap alive
/// past eviction). Shared by the storage / query / disk-cache tests.
pub fn default_disk_cache(
    storage: Arc<dyn StorageProvider>,
    cache_root: &Path,
) -> Arc<DiskCacheStore> {
    let cfg = DiskCacheConfig {
        cache_root: cache_root.to_path_buf(),
        disk_budget_bytes: TEST_DISK_CACHE_BUDGET_BYTES,
        cold_fetch_mode: ColdFetchMode::HybridWithPrefetch,
        cold_fetch_streams: TEST_COLD_FETCH_STREAMS,
        cold_fetch_chunk_bytes: TEST_COLD_FETCH_CHUNK_BYTES,
        mmap_cold_threshold_secs: 0,
        mmap_sweep_interval_secs: 0,
        eviction: Box::new(LruPolicy::new()),
        verify_crc_on_open: true,
        ..Default::default()
    };
    let pinned: Arc<dyn Fn() -> HashSet<_> + Send + Sync> = Arc::new(HashSet::new);
    DiskCacheStore::new(storage, cfg, pinned).expect("test disk cache")
}

/// A `DiskCacheStore` in `LazyForegroundWithBackgroundFill`: the foreground
/// query reads through an `open_lazy` `StorageRangeSource`, so the superfile's
/// bytes stay non-resident and every cold read is an object-store GET. That is
/// the path the connection budget gates. `default_disk_cache`
/// (`HybridWithPrefetch`) instead collects the cold responses into a resident
/// in-memory reader, which reads warm and reserves nothing. Same budget,
/// timers, and eviction otherwise.
pub fn lazy_foreground_disk_cache(
    storage: Arc<dyn StorageProvider>,
    cache_root: &Path,
) -> Arc<DiskCacheStore> {
    let cfg = DiskCacheConfig {
        cache_root: cache_root.to_path_buf(),
        disk_budget_bytes: TEST_DISK_CACHE_BUDGET_BYTES,
        cold_fetch_mode: ColdFetchMode::LazyForegroundWithBackgroundFill,
        cold_fetch_streams: TEST_COLD_FETCH_STREAMS,
        cold_fetch_chunk_bytes: TEST_COLD_FETCH_CHUNK_BYTES,
        mmap_cold_threshold_secs: 0,
        mmap_sweep_interval_secs: 0,
        eviction: Box::new(LruPolicy::new()),
        verify_crc_on_open: true,
        ..Default::default()
    };
    let pinned: Arc<dyn Fn() -> HashSet<_> + Send + Sync> = Arc::new(HashSet::new);
    DiskCacheStore::new(storage, cfg, pinned).expect("test lazy-foreground disk cache")
}

/// Build a `Decimal128Array(38, 0)` from `u64` ids.
///
/// Centralizes the verbose three-step construction that
/// every test fixture reinvents:
///
/// ```ignore
/// Decimal128Array::from(ids.into_iter().map(|v| v as i128).collect::<Vec<_>>())
///     .with_precision_and_scale(38, 0)
///     .expect("decimal128")
/// ```
pub fn decimal128_ids<I: IntoIterator<Item = u64>>(ids: I) -> Decimal128Array {
    Decimal128Array::from(ids.into_iter().map(|v| v as i128).collect::<Vec<_>>())
        .with_precision_and_scale(38, 0)
        .expect("Decimal128(38, 0) is a valid precision/scale pair")
}

/// `Field` for the primary-key id column — `Decimal128(38, 0)`,
/// non-nullable. Caller supplies the column name (typically
/// `"_id"` at the supertable layer or `"doc_id"` in lower-level
/// superfile fixtures).
pub fn decimal128_id_field(name: &str) -> Field {
    Field::new(name, DataType::Decimal128(38, 0), false)
}

/// The default tokenizer used in tests + benches:
/// `AsciiLowerTokenizer` wrapped in `Arc<dyn Tokenizer>`.
///
/// Callers passing this into `BuilderOptions::new` wrap in
/// `Some(...)` at the call site:
///
/// ```ignore
/// BuilderOptions::new(schema, "doc_id", fts_cols, vec_cols,
///                     Some(default_tokenizer()));
/// ```
pub fn default_tokenizer() -> Arc<dyn Tokenizer> {
    Arc::new(AsciiLowerTokenizer)
}

/// Default `VectorConfig` for test fixtures: `dim=16`,
/// `n_cent=4`, `metric=Cosine`. Caller supplies the column
/// name and `rot_seed` — the only fields tests vary.
///
/// For realistic-scale vectors (e.g. `dim=384` in benches),
/// callers construct `VectorConfig` directly with their own
/// values.
pub fn default_vector_config(column: &str, rot_seed: u64) -> VectorConfig {
    VectorConfig {
        column: column.into(),
        dim: 16,
        rot_seed,
        metric: Metric::Cosine,
        rerank_codec: RerankCodec::Fp32,
        provided_centroids: None,
    }
}

/// Single-column user schema with `title: LargeUtf8`.
///
/// Mirrors the supertable's auto-`_id` model: the supertable
/// layer prepends `_id: Decimal128(38, 0)` automatically at
/// append time, so the user-facing schema only declares the
/// payload columns. Dozens of supertable tests reconstruct
/// this exact schema; centralizing keeps the
/// supertable-auto-injects-id contract in one place.
pub fn schema_id_title() -> Arc<Schema> {
    Arc::new(Schema::new(vec![Field::new(
        "title",
        DataType::LargeUtf8,
        false,
    )]))
}

/// Build a single-column `RecordBatch` of titles matching
/// [`schema_id_title`]. Caller supplies the title strings;
/// the rest is fixed.
pub fn build_title_batch(titles: &[&str]) -> RecordBatch {
    let titles_arr = LargeStringArray::from(titles.to_vec());
    RecordBatch::try_new(schema_id_title(), vec![Arc::new(titles_arr)])
        .expect("RecordBatch shape matches schema_id_title")
}

/// `SupertableOptions` with the test-fixture defaults:
/// [`schema_id_title`] schema, a single FTS column `title`,
/// no vector columns, and a 1-thread rayon writer pool.
///
/// Caller chains `.with_storage(...)` / `.with_disk_cache(...)`
/// / `.with_*(...)` for whatever the specific test needs.
/// Returning the un-storage-d shape lets each test decide
/// explicitly whether to attach storage.
pub fn default_supertable_options() -> SupertableOptions {
    let pool = Arc::new(
        ThreadPoolBuilder::new()
            .num_threads(1)
            .build()
            .expect("rayon ThreadPoolBuilder with num_threads(1) builds"),
    );
    SupertableOptions::new(
        schema_id_title(),
        vec![FtsConfig {
            column: "title".into(),
            positions: false,
        }],
        vec![],
        Some(default_tokenizer()),
    )
    .expect("SupertableOptions::new with default test fixture args")
    .with_writer_pool(pool)
}