hayson-engine 0.22.0

use std::collections::{HashMap, HashSet};

use crate::errors::HaysonNotFoundError;
use crate::query::planner::{OverlayCandidates, QueryOverlay};
use crate::registry::exact_index::value_keys;
use crate::registry::fulltext_index::CorpusDelta;
use crate::schema::{primary_key_of, validate_record, HaysonRecord, ValidateOptions};
use crate::table::{Table, TableError, TxnOp};

use super::unit_of_work::StagedView;

/// Every op staged for one table, resolved once into the state the unit will leave behind, with the
/// same two indexes over it that the committed registry keeps: an exact bucket per single-field
/// `exact`/`sorted` index, and the full-text tokens of each staged record. Built per query, because
/// the ops it reads grow between two queries on the same unit.
///
/// The indexes exist so a query through a unit does not have to check every record the unit staged:
/// an equality on an indexed field is answered from the buckets, and a `search` is scored over the
/// committed corpus with [`StagedIndex::corpus_delta`] spliced in, which is what lets a staged
/// record join the ranking and not merely the answer.
pub struct StagedIndex {
    records: HashMap<String, HaysonRecord>,
    deleted: HashSet<String>,
    buckets: HashMap<String, HashMap<String, Vec<String>>>,
    tokens: HashMap<String, Vec<String>>,
}

impl StagedIndex {
    /// Resolves `ops` against `table`'s committed state — the same replay as
    /// [`StagedOverlay::resolve`], run once for every key the ops touch instead of once per lookup.
    pub fn build(ops: &[TxnOp], table: &Table) -> Result<Self, TableError> {
        let overlay = StagedOverlay::new(ops);
        let mut index = StagedIndex {
            records: HashMap::new(),
            deleted: HashSet::new(),
            buckets: table
                .indexed_fields()
                .into_iter()
                .map(|field| (field, HashMap::new()))
                .collect(),
            tokens: HashMap::new(),
        };
        let indexes_fulltext = table.has_fulltext();
        for key in overlay.touched_keys(table)? {
            match overlay.resolve(&key, table)? {
                StagedView::Put(record) => {
                    for (field, buckets) in &mut index.buckets {
                        for encoded in value_keys(record.get(field)) {
                            buckets.entry(encoded).or_default().push(key.clone());
                        }
                    }
                    if indexes_fulltext {
                        index.tokens.insert(key.clone(), table.tokens_of(&record));
                    }
                    index.records.insert(key, record);
                }
                StagedView::Deleted => {
                    index.deleted.insert(key);
                }
                StagedView::Untouched => {}
            }
        }
        Ok(index)
    }

    /// The staged writes as a change to the committed corpus: every staged record added under its
    /// own tokens, every staged delete removed. A re-staged key is both — `FulltextIndex::search`
    /// hides an added key's committed posting before splicing the staged one in.
    pub fn corpus_delta(&self) -> CorpusDelta {
        CorpusDelta {
            added: self.tokens.clone(),
            removed: self.deleted.iter().cloned().collect(),
        }
    }
}

impl QueryOverlay for StagedIndex {
    fn records(&self) -> &HashMap<String, HaysonRecord> {
        &self.records
    }

    fn deleted(&self) -> &HashSet<String> {
        &self.deleted
    }

    fn candidates_for(&self, field: &str, encoded: &str) -> OverlayCandidates {
        match self.buckets.get(field) {
            None => OverlayCandidates::NotIndexed,
            Some(buckets) => OverlayCandidates::Narrowed(buckets.get(encoded).cloned().unwrap_or_default()),
        }
    }
}

/// The ops staged for one table within a `UnitOfWork`, replayed on demand against that table's
/// committed state — the predicate-only slice of ADR-3h3r's overlay (see `UnitOfWork::query`).
/// Read access by single key already existed as `UnitOfWork::preview`; this type carries the same
/// replay so `preview` and query overlay share it instead of diverging.
pub struct StagedOverlay<'a> {
    ops: &'a [TxnOp],
}

impl<'a> StagedOverlay<'a> {
    pub fn new(ops: &'a [TxnOp]) -> Self {
        StagedOverlay { ops }
    }

    /// What `key` would read as once every op here lands, replayed in staging order over `table`'s
    /// committed state — identical rule to `UnitOfWork::preview`.
    pub fn resolve(&self, key: &str, table: &Table) -> Result<StagedView, TableError> {
        let mut view = StagedView::Untouched;
        for op in self.ops {
            match op {
                TxnOp::Delete(op_key) => {
                    if op_key == key {
                        view = StagedView::Deleted;
                    }
                }
                TxnOp::Put(document) => {
                    let accepted = table.accepted_document(document);
                    let record = validate_record(table.schema(), accepted.as_ref(), ValidateOptions::default())?;
                    if primary_key_of(table.schema(), &record)? == key {
                        view = StagedView::Put(record);
                    }
                }
                TxnOp::Patch(op_key, changes) => {
                    if op_key != key {
                        continue;
                    }
                    let base = match &view {
                        StagedView::Put(record) => Some(record.clone()),
                        StagedView::Deleted => None,
                        StagedView::Untouched => table.get(key)?,
                    };
                    let Some(base) = base else {
                        return Err(TableError::NotFound(HaysonNotFoundError {
                            table: table.name().to_string(),
                            key: key.to_string(),
                        }));
                    };
                    let accepted = table.accepted_document(changes);
                    let options = ValidateOptions {
                        partial: true,
                        now: None,
                    };
                    let partial = validate_record(table.schema(), accepted.as_ref(), options)?;
                    let mut merged = base;
                    for (field, value) in &partial {
                        merged.insert(field.clone(), value.clone());
                    }
                    view = StagedView::Put(merged);
                }
            }
        }
        Ok(view)
    }

    /// Every key at least one staged op names — the whole of what an overlay has to resolve, since
    /// a key the unit does not touch reads the same through the unit as off the committed log.
    fn touched_keys(&self, table: &Table) -> Result<Vec<String>, TableError> {
        let mut seen = HashSet::new();
        let mut keys = Vec::new();
        for op in self.ops {
            let key = match op {
                TxnOp::Delete(key) => key.clone(),
                TxnOp::Patch(key, _) => key.clone(),
                TxnOp::Put(document) => {
                    let accepted = table.accepted_document(document);
                    let record = validate_record(table.schema(), accepted.as_ref(), ValidateOptions::default())?;
                    primary_key_of(table.schema(), &record)?
                }
            };
            if seen.insert(key.clone()) {
                keys.push(key);
            }
        }
        Ok(keys)
    }
}