steeldb/

tokenql.rs

//! Token-only IKL — the single retrieval language, an s-expression over the postings store.
//! Ported 1:1 from the TS `tokenql.ts` evaluator:
//!   `<atom>`        a tag or glob pattern → UNION of matching tokens' postings, ∩ scope
//!   (and A B …)   → ∩   (also the default for a bare list)
//!   (or  A B …)   → ∪
//!   (not A)       → scope − A   (closed-world inversion)
//! (The `num` typed range predicate is stubbed here — it belongs to the columnar numeric layer,
//! which is a separate benchmark; set-algebra is what we're measuring.)

use crate::bitmap::Postings;

pub trait TokenStore<B: Postings> {
    fn atom(&self, pattern: &str) -> B;
    fn universe(&self) -> B;
    /// Situations whose numeric `field` satisfies `field op value` (op ∈ ge|gt|le|lt|eq|ne). Stores
    /// without a numeric layer return empty.
    fn numeric(&self, _field: &str, _op: &str, _value: f64) -> B {
        B::empty()
    }
    /// Situations where `token` holds with per-situation belief ≥ `min_bel` (paper §6, the `evidence` atom).
    /// Stores without a polarity layer fall back to plain membership.
    fn evidence(&self, token: &str, _min_bel: f64) -> B {
        self.atom(token)
    }
    /// Situations lying on an s-path between two tokens: a chain where consecutive steps share at least `s`
    /// situations (paper §3.3). Stores without a topological layer return nothing.
    fn s_path(&self, _a: &str, _b: &str, _s: usize) -> Option<B> {
        None
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Node {
    Atom(String),
    List(Vec<Node>),
}

/// Parse an s-expression. Atoms stay strings; `"quoted atoms"` keep spaces.
pub fn parse(expr: &str) -> Node {
    let mut stack: Vec<Vec<Node>> = Vec::new();
    let mut out: Vec<Node> = Vec::new();
    let mut buf = String::new();
    let mut inq = false;

    macro_rules! flush {
        () => {
            if !buf.is_empty() {
                out.push(Node::Atom(std::mem::take(&mut buf)));
            }
        };
    }

    for ch in expr.chars() {
        if inq {
            if ch == '"' {
                out.push(Node::Atom(std::mem::take(&mut buf)));
                inq = false;
            } else {
                buf.push(ch);
            }
        } else if ch == '"' {
            flush!();
            inq = true;
        } else if ch == '(' {
            flush!();
            stack.push(std::mem::take(&mut out));
        } else if ch == ')' {
            flush!();
            // A stray `)` is ignored rather than panicking. This is a `pub fn` reachable with any string, and
            // it used to abort on `query(")")` — in the wasm build a panic takes the whole module down. Callers
            // that care about well-formedness get that from the linter, which reports the imbalance; parsing
            // recovers so the failure is a refusal rather than a crash.
            if let Some(prev) = stack.pop() {
                let node = Node::List(std::mem::take(&mut out));
                out = prev;
                out.push(node);
            }
        } else if ch.is_whitespace() {
            flush!();
        } else {
            buf.push(ch);
        }
    }
    flush!();
    if out.len() == 1 {
        out.pop().unwrap()
    } else {
        Node::List(out)
    }
}

/// Why a query could not be evaluated. Separate from "matched nothing", which is an answer.
#[derive(Debug, Clone, PartialEq)]
pub enum QueryError {
    /// the sources in a `combine-ds` block disagree beyond the block's own threshold
    Conflict { conflict: f64, threshold: f64 },
    /// a `combine-ds` or `stream` block was not shaped as §6.1 requires
    Malformed(String),
}

impl std::fmt::Display for QueryError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            QueryError::Conflict { conflict, threshold } => write!(
                f,
                "evidential conflict K={conflict:.4} exceeds :max-conflict {threshold:.4}; \
                 the evidence streams disagree too much to fuse"
            ),
            QueryError::Malformed(what) => write!(f, "malformed query: {what}"),
        }
    }
}

impl std::error::Error for QueryError {}

/// Evaluate a query, returning empty on refusal. Kept for callers that cannot act on an error.
pub fn evaluate<B: Postings, S: TokenStore<B>>(store: &S, expr: &str) -> B {
    try_evaluate(store, expr).unwrap_or_else(|_| B::empty())
}

/// Evaluate a query, surfacing refusals.
///
/// A conflict refusal is not the same as an empty result and must not be collapsed into one: "the sources
/// contradict each other" is information the caller can act on, and "no rows matched" is an answer.
pub fn try_evaluate<B: Postings, S: TokenStore<B>>(store: &S, expr: &str) -> Result<B, QueryError> {
    let scope = store.universe();
    try_eval_node(store, &parse(expr), &scope)
}

/// Read a `:keyword value` pair out of an argument list.
fn keyword_arg(args: &[Node], name: &str) -> Option<String> {
    args.iter().position(|n| matches!(n, Node::Atom(a) if a == name)).and_then(|i| match args.get(i + 1) {
        Some(Node::Atom(v)) => Some(v.clone()),
        _ => None,
    })
}

/// `(mass (<atom>+) <float>)` — a focal set and its mass. The atoms are a DISJUNCTION: mass on
/// "one of these, cannot say which", so the focal set is their union.
fn focal_of<B: Postings, S: TokenStore<B>>(store: &S, node: &Node, scope: &B) -> Option<(B, f64)> {
    let Node::List(items) = node else { return None };
    if items.len() != 3 || !matches!(&items[0], Node::Atom(a) if a == "mass") {
        return None;
    }
    let mass = match &items[2] {
        Node::Atom(v) => v.parse::<f64>().ok()?,
        _ => return None,
    };
    let set = match &items[1] {
        Node::List(atoms) => {
            let mut u = B::empty();
            for a in atoms {
                if let Node::Atom(t) = a {
                    u.or_inplace(&store.atom(t));
                }
            }
            u.and(scope)
        }
        Node::Atom(t) => store.atom(t).and(scope),
    };
    Some((set, mass))
}

fn try_eval_node<B: Postings, S: TokenStore<B>>(
    store: &S,
    node: &Node,
    scope: &B,
) -> Result<B, QueryError> {
    Ok(match node {
        Node::Atom(a) => store.atom(a).and(scope),
        Node::List(items) => {
            if items.is_empty() {
                return Ok(scope.clone());
            }
            let op = if let Node::Atom(s) = &items[0] { s.as_str() } else { "" };
            let args = &items[1..];
            match op {
                "and" => {
                    let mut r = scope.clone();
                    for a in args {
                        r = r.and(&try_eval_node(store, a, scope)?);
                    }
                    r
                }
                "or" => {
                    let mut r = B::empty();
                    for a in args {
                        let e = try_eval_node(store, a, scope)?;
                        r.or_inplace(&e);
                    }
                    r.and(scope)
                }
                // `(not)` with no argument used to index an empty slice and PANIC — which in WebAssembly
                // aborts the whole module, taking the page with it. A model emitting a truncated expression is
                // an expected input, not a bug to crash on.
                "not" => match args.first() {
                    Some(inner) => scope.and_not(&try_eval_node(store, inner, scope)?),
                    None => return Err(QueryError::Malformed("not needs one argument".into())),
                },
                // (evidence <atom> :min-bel <f> :max-pl <f>) — §6.1. This was lint-accepted but never
                // evaluated: it fell through to implicit AND, where ":min-bel" and "0.8" are atoms matching
                // nothing, so a belief-constrained query silently returned empty instead of filtering.
                "evidence" if !args.is_empty() => {
                    let token = match &args[0] {
                        Node::Atom(t) => t.clone(),
                        other => return try_eval_node(store, other, scope),
                    };
                    let min_bel = keyword_arg(args, ":min-bel")
                        .and_then(|v| v.parse::<f64>().ok())
                        .unwrap_or(0.0);
                    store.evidence(&token, min_bel).and(scope)
                }
                // (s-path :s <n> (source <atom>) (target <atom>)) — §3.3/§6.1. Returns the situations the
                // path passes through, so it composes with the rest of the algebra like any other set.
                "s-path" => {
                    let s_thr = keyword_arg(args, ":s").and_then(|v| v.parse::<usize>().ok()).unwrap_or(1);
                    let endpoint = |name: &str| -> Option<String> {
                        args.iter().find_map(|n| match n {
                            Node::List(it)
                                if matches!(it.first(), Some(Node::Atom(h)) if h == name) =>
                            {
                                match it.get(1) {
                                    Some(Node::Atom(t)) => Some(t.clone()),
                                    _ => None,
                                }
                            }
                            _ => None,
                        })
                    };
                    let (Some(src), Some(dst)) = (endpoint("source"), endpoint("target")) else {
                        return Err(QueryError::Malformed(
                            "s-path needs (source <atom>) and (target <atom>)".into(),
                        ));
                    };
                    // an optional (constraint <expr>) narrows the scope BEFORE the walk — the predicate
                    // pushdown of §3.3, which is what bounds the traversal
                    let mut walk_scope = scope.clone();
                    for a in args {
                        if let Node::List(it) = a {
                            if matches!(it.first(), Some(Node::Atom(h)) if h == "constraint") {
                                if let Some(inner) = it.get(1) {
                                    walk_scope = walk_scope.and(&try_eval_node(store, inner, scope)?);
                                }
                            }
                        }
                    }
                    match store.s_path(&src, &dst, s_thr) {
                        Some(found) => found.and(&walk_scope),
                        None => {
                            return Err(QueryError::Malformed(
                                "this store has no topological layer, so s-path cannot be evaluated".into(),
                            ))
                        }
                    }
                }
                // (combine-ds :max-conflict <f> (stream …)+) — §6.1. Fuses independent streams by
                // Dempster's rule and REFUSES when they disagree beyond the stated threshold.
                "combine-ds" => {
                    let threshold = keyword_arg(args, ":max-conflict")
                        .and_then(|v| v.parse::<f64>().ok())
                        .unwrap_or(1.0);
                    let mut masses: Vec<crate::evidence::Mass<B>> = Vec::new();
                    for a in args {
                        let Node::List(items) = a else { continue };
                        if !matches!(items.first(), Some(Node::Atom(h)) if h == "stream") {
                            continue;
                        }
                        // :mass-assignments ((mass (…) f) …)
                        let pos = items
                            .iter()
                            .position(|n| matches!(n, Node::Atom(k) if k == ":mass-assignments"));
                        let Some(Node::List(focals)) = pos.and_then(|i| items.get(i + 1)) else {
                            return Err(QueryError::Malformed(
                                "stream needs :mass-assignments ((mass (<atom>+) <float>) …)".into(),
                            ));
                        };
                        let mut built: Vec<(B, f64)> = Vec::new();
                        for f in focals {
                            match focal_of(store, f, scope) {
                                // a focal set that matches nothing in this corpus carries no evidence;
                                // dropping it keeps the mass function valid instead of failing the query
                                Some((set, m)) if !set.is_empty() => built.push((set, m)),
                                Some(_) => {}
                                None => {
                                    return Err(QueryError::Malformed(
                                        "expected (mass (<atom>+) <float>)".into(),
                                    ))
                                }
                            }
                        }
                        // renormalise what survived, so dropped focal sets cannot silently unbalance it
                        let total: f64 = built.iter().map(|(_, m)| *m).sum();
                        if total <= 0.0 {
                            continue;
                        }
                        for (_, m) in built.iter_mut() {
                            *m /= total;
                        }
                        match crate::evidence::Mass::new(built) {
                            Ok(m) => masses.push(m),
                            Err(e) => return Err(QueryError::Malformed(e.to_string())),
                        }
                    }
                    if masses.is_empty() {
                        return Err(QueryError::Malformed("combine-ds needs at least one stream".into()));
                    }
                    match crate::evidence::combine_all(&masses, threshold) {
                        Ok(fused) => {
                            // everything still possible after fusion
                            let mut out = B::empty();
                            for (set, _) in fused.focals() {
                                out.or_inplace(set);
                            }
                            out.and(scope)
                        }
                        Err((_, crate::evidence::EvidenceError::ConflictExceeded { conflict, threshold })) => {
                            return Err(QueryError::Conflict { conflict, threshold })
                        }
                        Err((_, crate::evidence::EvidenceError::TotalConflict)) => {
                            return Err(QueryError::Conflict { conflict: 1.0, threshold })
                        }
                        Err((i, e)) => {
                            return Err(QueryError::Malformed(format!("stream {i}: {e}")))
                        }
                    }
                }
                // (num <field> <op> <value>) — numeric range predicate over the columnar numeric layer
                "num" if args.len() == 3 => {
                    if let (Node::Atom(field), Node::Atom(op), Node::Atom(val)) = (&args[0], &args[1], &args[2]) {
                        match val.parse::<f64>().ok().or_else(|| crate::units::parse_number(val)) {
                            Some(v) => store.numeric(field, op, v).and(scope),
                            None => B::empty(),
                        }
                    } else {
                        B::empty()
                    }
                }
                // bare list with no leading operator = implicit AND
                // These heads ARE implemented; reaching here means the shape was wrong, which is a different
                // failure from "not implemented" and must not be reported as one.
                "num" => {
                    return Err(QueryError::Malformed(
                        "num needs (num <field> <op> <value>), op one of ge gt le lt eq ne".into(),
                    ))
                }
                "evidence" => {
                    return Err(QueryError::Malformed(
                        "evidence needs (evidence <atom> :min-bel <float>)".into(),
                    ))
                }
                // A head the linter accepts but this evaluator does not implement must FAIL LOUDLY. Falling
                // through to implicit AND turns it into an empty result, which reads as "no rows matched" — a
                // wrong answer to a query that passed type-checking. This is how `evidence`, `combine-ds` and
                // `s-path` each shipped broken.
                other
                    if crate::linter::STRUCTURAL_HEADS.contains(&other)
                        && !crate::linter::SUB_FORMS.contains(&other) =>
                {
                    return Err(QueryError::Malformed(format!(
                        "'{other}' is accepted by the linter but not implemented by the evaluator"
                    )))
                }
                _ => {
                    let mut r = scope.clone();
                    for a in items {
                        r = r.and(&try_eval_node(store, a, scope)?);
                    }
                    r
                }
            }
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_nested() {
        let n = parse("(and a (or b c) (not d))");
        match n {
            Node::List(v) => assert_eq!(v.len(), 4),
            _ => panic!("expected list"),
        }
    }

    // ── a tiny store, so the query semantics are tested without an index ──
    use crate::bitmap::RoarPostings as P;
    use std::collections::HashMap;

    struct Toy {
        sets: HashMap<String, P>,
        universe: P,
        /// tokens whose situations are refuted, to exercise the evidence atom
        refuted: HashMap<String, P>,
    }
    impl Toy {
        fn new(pairs: &[(&str, &[u32])]) -> Toy {
            let mut sets = HashMap::new();
            let mut all: Vec<u32> = Vec::new();
            for (k, ids) in pairs {
                sets.insert((*k).to_string(), P::from_sorted(ids));
                all.extend_from_slice(ids);
            }
            all.sort_unstable();
            all.dedup();
            Toy { sets, universe: P::from_sorted(&all), refuted: HashMap::new() }
        }
        fn refute(mut self, token: &str, ids: &[u32]) -> Toy {
            self.refuted.insert(token.to_string(), P::from_sorted(ids));
            self
        }
    }
    impl TokenStore<P> for Toy {
        fn atom(&self, p: &str) -> P {
            self.sets.get(p).cloned().unwrap_or_else(P::empty)
        }
        fn universe(&self) -> P {
            self.universe.clone()
        }
        fn evidence(&self, token: &str, min_bel: f64) -> P {
            let base = self.atom(token);
            if min_bel <= 0.0 {
                return base;
            }
            match self.refuted.get(token) {
                Some(bad) => base.and_not(bad),
                None => base,
            }
        }
    }

    #[test]
    fn the_evidence_atom_filters_by_belief_instead_of_returning_empty() {
        // situations 1..4 mention it; 3 and 4 refute it
        let store = Toy::new(&[("artifact/cell", &[1, 2, 3, 4])]).refute("artifact/cell", &[3, 4]);

        // the bug this covers: `evidence` used to fall through to implicit AND, where ":min-bel" and "0.8"
        // are atoms matching nothing, so the whole query collapsed to empty
        let got = try_evaluate(&store, "(evidence artifact/cell :min-bel 0.8)").unwrap();
        assert_eq!(got.to_sorted(), vec![1, 2], "must drop the refuted situations, not everything");

        // with no threshold it is plain membership
        let all = try_evaluate(&store, "(evidence artifact/cell :min-bel 0.0)").unwrap();
        assert_eq!(all.to_sorted(), vec![1, 2, 3, 4]);
    }

    #[test]
    fn combine_ds_fuses_agreeing_streams() {
        let store = Toy::new(&[("a", &[1, 2, 3]), ("b", &[3, 4, 5])]);
        let q = "(combine-ds :max-conflict 0.5 \
                   (stream :id s1 :mass-assignments ((mass (a) 0.7) (mass (a b) 0.3))) \
                   (stream :id s2 :mass-assignments ((mass (a) 0.6) (mass (a b) 0.4))))";
        let got = try_evaluate(&store, q).expect("compatible streams should fuse");
        // both streams favour a; the fused possibilities stay inside a ∪ b
        assert!(!got.is_empty());
        assert!(got.to_sorted().iter().all(|s| (1..=5).contains(s)), "{:?}", got.to_sorted());
    }

    #[test]
    fn combine_ds_refuses_when_streams_contradict() {
        // disjoint evidence: one stream says {1,2}, the other {8,9}
        let store = Toy::new(&[("x", &[1, 2]), ("y", &[8, 9])]);
        let q = "(combine-ds :max-conflict 0.2 \
                   (stream :id s1 :mass-assignments ((mass (x) 1.0))) \
                   (stream :id s2 :mass-assignments ((mass (y) 1.0))))";
        match try_evaluate(&store, q) {
            Err(QueryError::Conflict { conflict, threshold }) => {
                assert!(conflict > threshold, "K={conflict} should exceed {threshold}");
            }
            other => panic!("expected a conflict refusal, got {other:?}"),
        }
        // and the infallible wrapper must not present that refusal as an ordinary empty answer
        assert!(evaluate(&store, q).is_empty());
    }

    #[test]
    fn a_permissive_threshold_still_fuses_the_same_streams() {
        let store = Toy::new(&[("x", &[1, 2]), ("y", &[8, 9]), ("either", &[1, 2, 8, 9])]);
        let q = "(combine-ds :max-conflict 1.0 \
                   (stream :id s1 :mass-assignments ((mass (x either) 1.0))) \
                   (stream :id s2 :mass-assignments ((mass (y either) 1.0))))";
        assert!(try_evaluate(&store, q).is_ok(), "overlapping focal sets are compatible");
    }

    #[test]
    fn a_malformed_stream_is_reported_not_silently_ignored() {
        let store = Toy::new(&[("a", &[1])]);
        let bad = "(combine-ds :max-conflict 0.5 (stream :id s1))";
        assert!(matches!(try_evaluate(&store, bad), Err(QueryError::Malformed(_))));
    }

    #[test]
    fn plain_queries_are_unaffected_by_the_new_error_channel() {
        let store = Toy::new(&[("a", &[1, 2, 3]), ("b", &[2, 3, 4])]);
        assert_eq!(try_evaluate(&store, "(and a b)").unwrap().to_sorted(), vec![2, 3]);
        assert_eq!(try_evaluate(&store, "(or a b)").unwrap().to_sorted(), vec![1, 2, 3, 4]);
        assert_eq!(try_evaluate(&store, "(and a (not b))").unwrap().to_sorted(), vec![1]);
    }

    #[test]
    fn every_head_the_linter_accepts_is_implemented() {
        // The contract that used to drift silently. A head whitelisted by the linter with no evaluator arm
        // passed type-checking and then returned an empty set — indistinguishable from "no rows matched".
        // s-path, source, target and constraint all sat in that state.
        let store = Toy::new(&[("a", &[1])]);
        for head in crate::linter::STRUCTURAL_HEADS {
            if crate::linter::SUB_FORMS.contains(head) {
                continue; // only valid inside a parent form
            }
            let q = format!("({head})");
            match try_evaluate(&store, &q) {
                // implemented: may succeed, or fail for a reason specific to its own shape
                Ok(_) => {}
                Err(QueryError::Malformed(m)) => assert!(
                    !m.contains("not implemented by the evaluator"),
                    "{head} is accepted by the linter but has no evaluator arm"
                ),
                Err(_) => {}
            }
        }
    }

    #[test]
    fn s_path_walks_the_token_graph_at_the_given_threshold() {
        use crate::index::InfonIndex;
        let mut raw: HashMap<String, Vec<u32>> = HashMap::new();
        raw.insert("a/x".into(), vec![0, 1]);
        raw.insert("b/y".into(), vec![0, 2]);
        raw.insert("c/z".into(), vec![1, 2]);
        let ix: InfonIndex<P> = InfonIndex::from_postings(raw, 3);

        // s=1: one shared situation is enough, so a chain exists
        let hit = try_evaluate(&ix, "(s-path :s 1 (source a/x) (target c/z))").unwrap();
        assert!(!hit.is_empty(), "a chain should exist at s=1");

        // s=2: no pair shares two situations, so nothing connects them. An ANSWER, not an error.
        let none = try_evaluate(&ix, "(s-path :s 2 (source a/x) (target c/z))").unwrap();
        assert!(none.is_empty(), "raising s must break the weak link");

        // (constraint …) narrows the scope before the walk — the predicate pushdown of §3.3
        let constrained =
            try_evaluate(&ix, "(s-path :s 1 (source a/x) (target c/z) (constraint b/y))").unwrap();
        assert_eq!(constrained.to_sorted(), vec![0, 2], "the constraint must bound the result");
        assert!(constrained.len() < hit.len(), "constrained must be narrower than unconstrained");
    }

    #[test]
    fn a_malformed_s_path_is_refused_rather_than_answered() {
        use crate::index::InfonIndex;
        let mut raw: HashMap<String, Vec<u32>> = HashMap::new();
        raw.insert("a/x".into(), vec![0]);
        let ix: InfonIndex<P> = InfonIndex::from_postings(raw, 1);
        assert!(matches!(
            try_evaluate(&ix, "(s-path :s 1 (source a/x))"),
            Err(QueryError::Malformed(_))
        ));
    }

    #[test]
    fn an_unknown_endpoint_yields_nothing_without_erroring() {
        use crate::index::InfonIndex;
        let mut raw: HashMap<String, Vec<u32>> = HashMap::new();
        raw.insert("a/x".into(), vec![0]);
        let ix: InfonIndex<P> = InfonIndex::from_postings(raw, 1);
        let r = try_evaluate(&ix, "(s-path :s 1 (source a/x) (target nope/tok))").unwrap();
        assert!(r.is_empty(), "an absent endpoint has no path; that is an answer");
    }
}