kevy-embedded 7.0.0

//! Write-path index maintenance — the `commit_write` hook body plus
//! the catalog↔shard segment-list reconciliation (split out of
//! `ops_index.rs` to keep it under the 500-LOC project ceiling;
//! behaviour unchanged).

use kevy_index::{AggRow, IndexKind, IndexSpec, Segment};

use crate::ops_index::{IndexReg, ShardSegs};

impl ShardSegs {
    /// Σ approximate heap bytes of this shard's index segments, every
    /// kind — the tier's `reserved_bytes` floor feed. Served from the
    /// generation cache: an idle store recomputes nothing (—
    /// the walk behind this sum was a consumer's measured 300-500× idle
    /// CPU term).
    #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
    pub(crate) fn reserved_bytes(&mut self) -> u64 {
        if !self.stats_dirty {
            return self.reserved_cache;
        }
        let mut sum: u64 = self.segs.iter().map(|(_, s)| s.stats().approx_bytes).sum();
        sum += self.agg.iter().map(|(_, a)| a.stats().approx_bytes).sum::<u64>();
        #[cfg(feature = "text")]
        {
            sum += self.text.iter().map(|(_, t)| t.stats().approx_bytes).sum::<u64>();
        }
        #[cfg(feature = "vector")]
        {
            sum += self.ann.iter().map(|(_, g)| g.stats().approx_bytes).sum::<u64>();
        }
        self.reserved_cache = sum;
        self.stats_dirty = false;
        sum
    }
}

/// Tiering floor refusal (mirrors the
/// server's IDX.CREATE precheck, same wire message): indexes are the
/// fixed layer demotion can never reclaim; when the existing floor
/// already exhausts the tier's demotable headroom, a new index is
/// refused by name. The floor is refreshed from the live segments
/// first so the check never trails the reaper tick.
#[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
pub(crate) fn tier_floor_check(shards: &crate::store::Shards) -> crate::KevyResult<()> {
    for shard in shards.iter() {
        let mut g = crate::store::lock_write(shard);
        if !g.store.tier_enabled() {
            break;
        }
        let reserved = g.idx_segs.reserved_bytes() + g.view_segs.reserved_bytes();
        g.store.set_tier_reserved(reserved);
        if g.store.tier_index_floor_blocked(0) {
            return Err(crate::KevyError::InvalidInput(
                "index memory floor exceeds the tiering budget".into(),
            ));
        }
    }
    Ok(())
}

/// A fresh text segment shaped by the spec: as many separately scored
/// fields as it declares (the breakdown `IN <field…>` scopes to), and a
/// positional side-channel when it asked for `WITH POSITIONS`.
#[cfg(feature = "text")]
pub(crate) fn new_text(spec: &IndexSpec) -> kevy_text::TextSegment {
    kevy_text::TextSegment::with_shape(
        kevy_text::SegmentShape::default()
            .with_fields(spec.fields().len())
            .with_positions(spec.has_positions())
            .with_values(spec.values().len()),
    )
}

#[cfg(feature = "vector")]
pub(crate) fn new_graph(spec: &IndexSpec) -> kevy_vector::Hnsw {
    let a = spec.ann().expect("an ann index carries its parameters");
    kevy_vector::Hnsw::new(
        a.dim as usize,
        kevy_vector::HnswParams::default()
            .with_m(a.m as usize)
            .with_ef_construction(a.ef as usize)
            .with_distance(match a.distance {
                1 => kevy_vector::Distance::L2,
                2 => kevy_vector::Distance::Ip,
                _ => kevy_vector::Distance::Cosine,
            }),
    )
}

/// Reconcile one shard's segment list with the catalog; new indexes
/// backfill from this shard's live keys (we hold the shard's write
/// lock — no concurrent writes can race the scan).
pub(crate) fn sync_segs(reg: &IndexReg, shard_segs: &mut ShardSegs, store: &mut kevy_store::Store) {
    {
        let g = reg.catalog.read().unwrap_or_else(std::sync::PoisonError::into_inner);
        let (ver, cat) = &*g;
        if shard_segs.version != *ver {
            rebuild_seg_lists(cat, *ver, shard_segs, store);
        }
    }
    // every row the store saw written since the last drain, however it
    // was written
    if crate::ops_index_changes::drain(shard_segs, store) {
        shard_segs.mark_stats_dirty();
    }
}

/// The out-of-date half of `sync_segs`: rebuild every per-kind segment
/// list against the catalog (existing segments move, new specs backfill).
// Every kind appears twice below, once for the build that has the engine
// and once for the build that answers NotFound instead, which is where the
// length comes from — the control flow is one `for`.
// LOC-WAIVER: a match over IndexKind, one push per arm, doubled by #[cfg].
fn rebuild_seg_lists(
    cat: &kevy_index::Catalog,
    ver: u64,
    shard_segs: &mut ShardSegs,
    store: &mut kevy_store::Store,
) {
    let mut next: Vec<(IndexSpec, Segment)> = Vec::new();
    #[cfg(feature = "text")]
    let mut next_text: Vec<(IndexSpec, kevy_text::TextSegment)> = Vec::new();
    #[cfg(feature = "vector")]
    let mut next_ann: Vec<(IndexSpec, kevy_vector::Hnsw)> = Vec::new();
    let mut next_agg: Vec<(IndexSpec, kevy_index::AggSegment)> = Vec::new();
    for (spec, _) in cat.iter() {
        let (segs, st) = (&mut *shard_segs, &mut *store);
        match spec.kind() {
            IndexKind::Agg => next_agg.push(take_or_backfill(
                &mut segs.agg,
                spec,
                st,
                kevy_index::AggSegment::new,
                apply_agg_key,
            )),
            #[cfg(feature = "vector")]
            IndexKind::Ann => next_ann.push(take_or_backfill(
                &mut segs.ann,
                spec,
                st,
                || new_graph(spec),
                apply_ann_key,
            )),
            // Engine compiled out (idx_create rejects the kind; a
            // sidecar-loaded spec gets no segment, so queries answer
            // NotFound instead of silently mis-indexing).
            #[cfg(not(feature = "vector"))]
            IndexKind::Ann => {}
            #[cfg(feature = "text")]
            IndexKind::Text => next_text.push(take_or_backfill(
                &mut segs.text,
                spec,
                st,
                || new_text(spec),
                apply_text_key,
            )),
            #[cfg(not(feature = "text"))]
            IndexKind::Text => {}
            _ => {
                let fresh = !segs.segs.iter().any(|(s, _)| s == spec);
                let mut built = take_or_backfill(
                    &mut segs.segs,
                    spec,
                    st,
                    || new_scalar(spec),
                    |st, sp, sg, k| apply_key(st, sp, sg, k, None),
                );
                if fresh {
                    // the walk filled leaves in hash order; pack them
                    built.1.repack();
                }
                next.push(built);
            }
        }
    }
    shard_segs.segs = next;
    #[cfg(feature = "text")]
    {
        shard_segs.text = next_text;
    }
    #[cfg(feature = "vector")]
    {
        shard_segs.ann = next_ann;
    }
    shard_segs.agg = next_agg;
    shard_segs.version = ver;
    shard_segs.mark_stats_dirty();
}

/// Buckets visited per page of a backfill walk.
const BACKFILL_PAGE_BUCKETS: usize = 4096;

/// Keep `spec`'s existing segment from `have` (position move), or
/// backfill a fresh one from this shard's live keys in the spec's
/// prefix domain.
fn take_or_backfill<S>(
    have: &mut Vec<(IndexSpec, S)>,
    spec: &IndexSpec,
    store: &mut kevy_store::Store,
    empty: impl FnOnce() -> S,
    apply: impl Fn(&mut kevy_store::Store, &IndexSpec, &mut S, &[u8]),
) -> (IndexSpec, S) {
    if let Some(i) = have.iter().position(|(s, _)| s == spec) {
        return have.swap_remove(i);
    }
    let mut seg = empty();
    let mut pat = spec.prefix().to_vec();
    pat.push(b'*');
    // a page of keys at a time rather than a copy of every key under the
    // prefix; the store is held for the whole walk, so no key is missed
    let mut cursor = 0;
    loop {
        let (next, keys) = store.walk_page(cursor, BACKFILL_PAGE_BUCKETS, Some(&pat));
        for key in &keys {
            apply(store, spec, &mut seg, key);
        }
        if next == 0 {
            return (spec.clone(), seg);
        }
        cursor = next;
    }
}

fn apply_agg_key(
    store: &mut kevy_store::Store,
    spec: &IndexSpec,
    a: &mut kevy_index::AggSegment,
    key: &[u8],
) {
    // Both fields in ONE row peek (server twin: `apply_row_agg`) —
    // one record read on a cold row, no promotion, no gate mark; the
    // `Ok(None)`/`Err` arms carry the old `exists()` distinction.
    let group_field = spec.group_by().unwrap_or_default();
    match store.peek_hash_fields(key, &[group_field, spec.field()]) {
        Ok(Some(mut vals)) => {
            let group = vals[0].take();
            let val =
                vals[1].take().and_then(|raw| kevy_index::IndexValue::coerce(spec.ty(), &raw));
            match (group, val) {
                (Some(g), Some(v)) => a.apply(key, AggRow::Member { group: g, value: v }),
                _ => a.apply(key, AggRow::Excluded),
            }
        }
        Ok(None) => a.apply(key, AggRow::Removed),
        Err(_) => a.apply(key, AggRow::Excluded),
    }
}

#[cfg(feature = "vector")]
fn apply_ann_key(
    store: &mut kevy_store::Store,
    spec: &IndexSpec,
    g: &mut kevy_vector::Hnsw,
    key: &[u8],
) {
    // The row peek — one record read on cold, no promotion, no
    // gate mark (server twin: `apply_row`'s ann arm).
    let v = match store.peek_hash_fields(key, &[spec.field()]) {
        Ok(Some(mut vals)) => {
            vals[0].take().and_then(|raw| kevy_vector::parse_vector(&raw, g.dim()))
        }
        _ => None,
    };
    g.apply(key, v);
}

#[cfg(feature = "text")]
fn apply_text_key(
    store: &mut kevy_store::Store,
    spec: &IndexSpec,
    ts: &mut kevy_text::TextSegment,
    key: &[u8],
) {
    // The spec owns what it reads out of a row -- declared fields with
    // their weights, declared stored values -- so this path and the
    // server's cannot index the same row differently. Every
    // declared field + value prefetched with ONE row peek (one record
    // read on a cold row, no promotion, no gate mark); `read_row`
    // resolves from the prefetch, not per-field hgets.
    let names: Vec<&[u8]> = spec
        .fields()
        .iter()
        .map(|f| f.name.as_slice())
        .chain(spec.values().iter().map(|v| v.name.as_slice()))
        .collect();
    let fetched = store.peek_hash_fields(key, &names).ok().flatten();
    let (fields, values) = spec.read_row(|f| {
        let vals = fetched.as_ref()?;
        names.iter().position(|n| *n == f).and_then(|i| vals[i].clone())
    });
    let vals: Vec<Option<&[u8]>> = values.iter().map(|v| v.as_deref()).collect();
    if fields.is_empty() {
        ts.apply_doc(key, None, &vals);
    } else {
        ts.apply_doc(key, Some(&fields), &vals);
    }
}

/// The write-path hook body — called from `commit_write` with the
/// logged argv, under the shard lock. Extracts the written key(s)
/// EXACTLY and re-derives their index entries.
pub(crate) fn on_commit(
    reg: &IndexReg,
    shard_segs: &mut ShardSegs,
    store: &mut kevy_store::Store,
    parts: &[&[u8]],
) {
    {
        // Cheap gate: empty catalog = one read-lock + is_empty.
        let g = reg.catalog.read().unwrap_or_else(std::sync::PoisonError::into_inner);
        if g.1.is_empty() {
            return;
        }
    }
    sync_segs(reg, shard_segs, store);
    let verb = parts.first().copied().unwrap_or(b"");
    if verb.eq_ignore_ascii_case(b"FLUSHALL") || verb.eq_ignore_ascii_case(b"FLUSHDB") {
        reset_all_segs(shard_segs);
        shard_segs.mark_stats_dirty();
    }
}

/// The text half of one written key: re-derive its inverted entries,
/// then — on a windowed table — shadow any frozen entries this write
/// stales (rewrite / delete / revival; the statistics withdraw
/// exactly — server twin: `apply_row`'s text arm).
#[cfg(feature = "text")]
fn apply_text_arm(shard_segs: &mut ShardSegs, store: &mut kevy_store::Store, key: &[u8]) -> bool {
    let mut touched = false;
    for (spec, ts) in &mut shard_segs.text {
        if key.starts_with(spec.prefix()) {
            apply_text_key(store, spec, ts, key);
            touched = true;
        }
    }
    #[cfg(not(target_arch = "wasm32"))]
    {
        let ShardSegs { text, cold_text, .. } = &mut *shard_segs;
        for (name, dir) in cold_text.iter_mut() {
            if text.iter().any(|(s, _)| s.name() == name && key.starts_with(s.prefix())) {
                dir.on_row_write(key);
            }
        }
    }
    touched
}

/// Apply one written key to every matching segment of every kind;
/// `olds` is, per scalar segment, the value the row was indexed under
/// before the write. Returns whether any segment was touched (the
/// cache-invalidation signal).
pub(crate) fn apply_one_key(
    shard_segs: &mut ShardSegs,
    store: &mut kevy_store::Store,
    key: &[u8],
    olds: &[Option<kevy_index::IndexValue>],
) -> bool {
    let mut touched = false;
    for ((spec, seg), old) in shard_segs.segs.iter_mut().zip(olds) {
        if key.starts_with(spec.prefix()) {
            apply_key(store, spec, seg, key, old.as_ref());
            touched = true;
        }
    }
    // Windowed indexes: this write may shadow a cold entry (rewrite /
    // delete / revival) — the bloom decides if it earns a tombstone.
    #[cfg(not(target_arch = "wasm32"))]
    {
        let ShardSegs { segs, windows, .. } = &mut *shard_segs;
        for (name, win) in windows.iter_mut() {
            if segs.iter().any(|(s, _)| s.name() == name && key.starts_with(s.prefix())) {
                win.on_row_write(key);
            }
        }
    }
    #[cfg(feature = "text")]
    {
        touched |= apply_text_arm(shard_segs, store, key);
    }
    #[cfg(feature = "vector")]
    for (spec, g) in &mut shard_segs.ann {
        if key.starts_with(spec.prefix()) {
            apply_ann_key(store, spec, g, key);
            touched = true;
        }
    }
    for (spec, a) in &mut shard_segs.agg {
        if key.starts_with(spec.prefix()) {
            apply_agg_key(store, spec, a, key);
            touched = true;
        }
    }
    touched
}

/// FLUSHALL / FLUSHDB: every segment resets to empty.
pub(crate) fn reset_all_segs(shard_segs: &mut ShardSegs) {
    for (spec, seg) in &mut shard_segs.segs {
        // a view reads the new segment by key as it read the old one
        let key_dir = seg.key_dir().is_some();
        *seg = new_scalar(spec);
        seg.set_key_dir(key_dir);
    }
    #[cfg(feature = "text")]
    for (spec, ts) in &mut shard_segs.text {
        *ts = new_text(spec);
    }
    #[cfg(feature = "vector")]
    for (spec, g) in &mut shard_segs.ann {
        *g = new_graph(spec);
    }
    for (_, a) in &mut shard_segs.agg {
        *a = kevy_index::AggSegment::new();
    }
}

/// EXACT written-key walk per verb shape (the logged effect argv is a
/// closed set — new effect verbs must be added here; the parity test
/// in `store_tests_index.rs` pins the list).
pub(crate) fn each_written_key_pub(verb: &[u8], parts: &[&[u8]], f: impl FnMut(&[u8])) {
    each_written_key(verb, parts, f);
}

fn each_written_key(verb: &[u8], parts: &[&[u8]], mut f: impl FnMut(&[u8])) {
    let up = |v: &[u8], t: &[u8]| v.eq_ignore_ascii_case(t);
    if up(verb, b"DEL") || up(verb, b"UNLINK") {
        for k in &parts[1..] {
            f(k);
        }
    } else if up(verb, b"MSET") {
        let mut i = 1;
        while i + 1 < parts.len() {
            f(parts[i]);
            i += 2;
        }
    } else if up(verb, b"COPY") || up(verb, b"RENAME") || up(verb, b"RENAMENX") {
        if let Some(k) = parts.get(1) {
            f(k);
        }
        if let Some(k) = parts.get(2) {
            f(k);
        }
    } else if let Some(k) = parts.get(1) {
        f(k);
    }
}

/// A fresh scalar segment shaped by the spec.
fn new_scalar(spec: &IndexSpec) -> Segment {
    Segment::for_spec(spec)
}

/// The primary field AND every declared VALUES column read with
/// ONE `peek_hash_fields` row peek — a cold row costs one record read
/// plus one decode (never one per field), promotes nothing and never
/// advances the 2nd-touch gate (the server twin is
/// `index_runtime::apply_scalar_row`). The peek's `Ok(None)`/`Err`
/// arms replace the old `exists()` disambiguation probe exactly.
fn apply_key(
    store: &mut kevy_store::Store,
    spec: &IndexSpec,
    seg: &mut Segment,
    key: &[u8],
    old: Option<&kevy_index::IndexValue>,
) {
    // The driving columns (composite or single FIELD) + VALUES in one
    // row peek; the derivation is the spec's own
    // ([`IndexSpec::derive_scalar`]) — the single implementation the
    // server's `apply_scalar_row` applies too, so the two engines
    // cannot index one row differently.
    let names = spec.scalar_read_names();
    let w = spec.primary_width();
    match store.peek_hash_fields(key, &names) {
        Ok(None) | Err(_) => {
            if let Some(o) = old {
                seg.remove(key, o);
            }
        }
        Ok(Some(vals)) => {
            let primary = spec.derive_scalar(&vals[..w]);
            match primary {
                None => seg.apply_with_values(key, old, None, &[]),
                Some(v) if spec.values().is_empty() => seg.apply(key, old, Some(v)),
                Some(v) => {
                    let refs: Vec<Option<&[u8]>> = vals[w..].iter().map(|o| o.as_deref()).collect();
                    seg.apply_with_values(key, old, Some(v), &refs);
                }
            }
        }
    }
}