Skip to main content

omgbase_surface/
query.rs

1//! The runner (`spec/surface/README.md` §1.4): parse the source, rewrite the
2//! row functions to `$self` methods, run it through the engine over the
3//! store context, and shape the engine's result into the surface's
4//! `OqxResult` (lean `{ id, path, … }` hits, keyset paging, the consumer
5//! scalars). Port of `packages/core/src/oqx-js/run.ts`.
6//!
7//! By default the tier-3 planner ([`crate::planner`]) pre-filters the scan in
8//! SQL and the in-memory engine finishes the residual over the produced rows;
9//! [`QueryOptions::in_memory`] forces the pure in-memory engine (the
10//! differential gate runs every case both ways — a planner must be invisible).
11
12use std::collections::HashMap;
13
14use omgbase_search::{EmbeddingProvider, f32_to_blob};
15use omgbase_store::Store;
16use oqx::ast::{Expr, Follow, OpNode, OrderSpec, Query, SelectItem, Subquery, Where};
17use oqx::{Consumer, Engine, InMemoryEngine, Value};
18use serde_json::{Map, Value as Json};
19
20use crate::context::{SemanticVec, StoreContext, render_row_values};
21use crate::cursor::{decode_cursor, encode_cursor};
22use crate::error::{Result, SurfaceError};
23use crate::planner::SqlitePlanner;
24
25/// The default page size.
26pub const DEFAULT_LIMIT: usize = 50;
27
28/// Row-scoped domain functions: authored as free calls that implicitly
29/// reference the current row; rewritten to `$self.fn(…)`.
30const ROW_FNS: [&str; 12] = [
31    "text",
32    "semantic",
33    "under",
34    "under_heading",
35    "within",
36    "under_kind",
37    "yaml_path",
38    "json_pointer",
39    "has_edge",
40    "has_anchor",
41    "child_count",
42    "parent_type",
43];
44
45const ID_KEY: &str = "__oqx_id";
46const PATH_KEY: &str = "__oqx_path";
47/// The reserved key a top-level `values` projection's single item is renamed
48/// to, so it rides through id/path injection, paging and distinct as an
49/// ordinary field and is peeled off at the end.
50const VALUE_KEY: &str = "__oqx_value";
51
52/// `query`'s options.
53#[derive(Clone, Copy, Default)]
54pub struct QueryOptions<'a> {
55    /// The page cap (default 50).
56    pub limit: Option<usize>,
57    /// Resume after a truncated page's cursor.
58    pub cursor: Option<&'a str>,
59    /// The provider behind `semantic(...)`; `None` → `semantic_unavailable`
60    /// when the query names a phrase.
61    pub provider: Option<&'a dyn EmbeddingProvider>,
62    /// Force the pure in-memory engine (skip the tier-3 pushdown planner).
63    /// The default plans; the differential gate runs both and compares.
64    pub in_memory: bool,
65}
66
67/// §1.4 `OqxResult`.
68#[derive(Clone, Debug, PartialEq)]
69pub struct OqxResult {
70    pub hits: Vec<Json>,
71    pub truncated: bool,
72    pub cursor: Option<String>,
73    pub consumer: Consumer,
74    pub count: Option<f64>,
75    pub exists: Option<bool>,
76    pub none: Option<bool>,
77    /// A top-level `values` projection's bare values, in place of `hits`.
78    pub values: Option<Vec<Json>>,
79}
80
81impl OqxResult {
82    fn scalar(consumer: Consumer) -> Self {
83        Self {
84            hits: Vec::new(),
85            truncated: false,
86            cursor: None,
87            consumer,
88            count: None,
89            exists: None,
90            none: None,
91            values: None,
92        }
93    }
94
95    /// The wire shape: `{ hits, truncated, cursor, consumer, count?, exists?,
96    /// none?, values? }`.
97    #[must_use]
98    pub fn to_json(&self) -> Json {
99        let mut m = Map::new();
100        m.insert("hits".to_owned(), Json::Array(self.hits.clone()));
101        m.insert("truncated".to_owned(), Json::Bool(self.truncated));
102        m.insert(
103            "cursor".to_owned(),
104            self.cursor.clone().map_or(Json::Null, Json::String),
105        );
106        m.insert(
107            "consumer".to_owned(),
108            Json::String(self.consumer.as_str().to_owned()),
109        );
110        if let Some(n) = self.count {
111            m.insert("count".to_owned(), Value::Number(n).to_canonical_json());
112        }
113        if let Some(b) = self.exists {
114            m.insert("exists".to_owned(), Json::Bool(b));
115        }
116        if let Some(b) = self.none {
117            m.insert("none".to_owned(), Json::Bool(b));
118        }
119        if let Some(v) = &self.values {
120            m.insert("values".to_owned(), Json::Array(v.clone()));
121        }
122        Json::Object(m)
123    }
124}
125
126// ---- the `$self` rewrite -----------------------------------------------------------
127
128fn self_ref() -> Expr {
129    Expr::Ident {
130        name: "$self".to_owned(),
131    }
132}
133
134fn rewrite_expr(e: &Expr) -> Expr {
135    match e {
136        Expr::Member { recv, name } => Expr::Member {
137            recv: Box::new(rewrite_expr(recv)),
138            name: name.clone(),
139        },
140        Expr::Index { recv, index } => Expr::Index {
141            recv: Box::new(rewrite_expr(recv)),
142            index: Box::new(rewrite_expr(index)),
143        },
144        Expr::Unary { op, expr } => Expr::Unary {
145            op: *op,
146            expr: Box::new(rewrite_expr(expr)),
147        },
148        Expr::Binary { op, left, right } => Expr::Binary {
149            op: *op,
150            left: Box::new(rewrite_expr(left)),
151            right: Box::new(rewrite_expr(right)),
152        },
153        Expr::Logical { op, left, right } => Expr::Logical {
154            op: *op,
155            left: Box::new(rewrite_expr(left)),
156            right: Box::new(rewrite_expr(right)),
157        },
158        Expr::In { left, right } => Expr::In {
159            left: Box::new(rewrite_expr(left)),
160            right: Box::new(rewrite_expr(right)),
161        },
162        Expr::Range {
163            lo,
164            hi,
165            exclusive_end,
166        } => Expr::Range {
167            lo: lo.as_ref().map(|x| Box::new(rewrite_expr(x))),
168            hi: hi.as_ref().map(|x| Box::new(rewrite_expr(x))),
169            exclusive_end: *exclusive_end,
170        },
171        Expr::Call { recv, name, args } => {
172            let args = args.iter().map(rewrite_expr).collect();
173            match recv {
174                None if ROW_FNS.contains(&name.as_str()) => Expr::Call {
175                    recv: Some(Box::new(self_ref())),
176                    name: name.clone(),
177                    args,
178                },
179                None => Expr::Call {
180                    recv: None,
181                    name: name.clone(),
182                    args,
183                },
184                Some(r) => Expr::Call {
185                    recv: Some(Box::new(rewrite_expr(r))),
186                    name: name.clone(),
187                    args,
188                },
189            }
190        }
191        Expr::Lit(_) | Expr::Ident { .. } | Expr::Outer { .. } | Expr::Binding { .. } => e.clone(),
192    }
193}
194
195fn rewrite_where(w: &Where) -> Where {
196    match w {
197        Where::And { parts } => Where::And {
198            parts: parts.iter().map(rewrite_where).collect(),
199        },
200        Where::Or { parts } => Where::Or {
201            parts: parts.iter().map(rewrite_where).collect(),
202        },
203        Where::Not { expr } => Where::Not {
204            expr: Box::new(rewrite_where(expr)),
205        },
206        Where::Scalar { expr } => Where::Scalar {
207            expr: rewrite_expr(expr),
208        },
209        Where::Op(op) => Where::Op(Box::new(rewrite_op(op))),
210    }
211}
212
213fn rewrite_op(op: &OpNode) -> OpNode {
214    OpNode {
215        receiver: rewrite_expr(&op.receiver),
216        op: op.op,
217        sub: rewrite_sub(&op.sub),
218        count_cmp: op.count_cmp.clone(),
219        distinct: op.distinct,
220    }
221}
222
223fn rewrite_follow(f: &Follow) -> Follow {
224    Follow {
225        receiver: rewrite_expr(&f.receiver),
226        distinct: f.distinct,
227        r#where: f.r#where.as_ref().map(rewrite_expr),
228        frontier: f.frontier.as_ref().map(rewrite_expr),
229        depth: f.depth,
230        by: f.by.as_ref().map(rewrite_expr),
231    }
232}
233
234fn rewrite_select(items: &[SelectItem]) -> Vec<SelectItem> {
235    items
236        .iter()
237        .map(|it| match it {
238            SelectItem::Field { name, expr, lift } => SelectItem::Field {
239                name: name.clone(),
240                expr: rewrite_expr(expr),
241                lift: *lift,
242            },
243            SelectItem::Collect { name, op } => SelectItem::Collect {
244                name: name.clone(),
245                op: Box::new(rewrite_op(op)),
246            },
247        })
248        .collect()
249}
250
251fn rewrite_order(o: Option<&Vec<OrderSpec>>) -> Option<Vec<OrderSpec>> {
252    o.map(|specs| {
253        specs
254            .iter()
255            .map(|s| OrderSpec {
256                expr: rewrite_expr(&s.expr),
257                desc: s.desc,
258            })
259            .collect()
260    })
261}
262
263fn rewrite_sub(s: &Subquery) -> Subquery {
264    Subquery {
265        from: s.from.iter().map(rewrite_expr).collect(),
266        r#where: s.r#where.as_ref().map(rewrite_where),
267        select: rewrite_select(&s.select),
268        order_by: rewrite_order(s.order_by.as_ref()),
269        follow: s.follow.as_ref().map(rewrite_follow),
270        values: s.values,
271        limit: s.limit.as_ref().map(rewrite_expr),
272        offset: s.offset.as_ref().map(rewrite_expr),
273    }
274}
275
276/// Rewrite every row function in a parsed query to a `$self` method call.
277#[must_use]
278pub fn rewrite_query(q: &Query) -> Query {
279    Query {
280        source: rewrite_expr(&q.source),
281        from: q.from.iter().map(rewrite_expr).collect(),
282        r#where: q.r#where.as_ref().map(rewrite_where),
283        select: rewrite_select(&q.select),
284        order_by: rewrite_order(q.order_by.as_ref()),
285        consumer: q.consumer,
286        follow: q.follow.as_ref().map(rewrite_follow),
287        distinct: q.distinct,
288        values: q.values,
289        limit: q.limit.as_ref().map(rewrite_expr),
290        offset: q.offset.as_ref().map(rewrite_expr),
291    }
292}
293
294// ---- semantic phrases -----------------------------------------------------------------
295
296fn visit_expr(e: &Expr, out: &mut Vec<String>) {
297    match e {
298        Expr::Call { recv, name, args } => {
299            if recv.is_none() && name == "semantic" {
300                if let Some(Expr::Lit(Value::Str(s))) = args.first() {
301                    if !out.contains(s) {
302                        out.push(s.clone());
303                    }
304                }
305            }
306            if let Some(r) = recv {
307                visit_expr(r, out);
308            }
309            for a in args {
310                visit_expr(a, out);
311            }
312        }
313        Expr::Member { recv, .. } => visit_expr(recv, out),
314        Expr::Index { recv, index } => {
315            visit_expr(recv, out);
316            visit_expr(index, out);
317        }
318        Expr::Unary { expr, .. } => visit_expr(expr, out),
319        Expr::Binary { left, right, .. }
320        | Expr::Logical { left, right, .. }
321        | Expr::In { left, right } => {
322            visit_expr(left, out);
323            visit_expr(right, out);
324        }
325        Expr::Range { lo, hi, .. } => {
326            if let Some(l) = lo {
327                visit_expr(l, out);
328            }
329            if let Some(h) = hi {
330                visit_expr(h, out);
331            }
332        }
333        Expr::Lit(_) | Expr::Ident { .. } | Expr::Outer { .. } | Expr::Binding { .. } => {}
334    }
335}
336
337fn visit_where(w: &Where, out: &mut Vec<String>) {
338    match w {
339        Where::And { parts } | Where::Or { parts } => {
340            parts.iter().for_each(|p| visit_where(p, out))
341        }
342        Where::Not { expr } => visit_where(expr, out),
343        Where::Scalar { expr } => visit_expr(expr, out),
344        Where::Op(op) => visit_op(op, out),
345    }
346}
347
348fn visit_op(op: &OpNode, out: &mut Vec<String>) {
349    visit_expr(&op.receiver, out);
350    visit_sub(&op.sub, out);
351}
352
353fn visit_select(items: &[SelectItem], out: &mut Vec<String>) {
354    for it in items {
355        match it {
356            SelectItem::Field { expr, .. } => visit_expr(expr, out),
357            SelectItem::Collect { op, .. } => visit_op(op, out),
358        }
359    }
360}
361
362fn visit_follow(f: &Follow, out: &mut Vec<String>) {
363    visit_expr(&f.receiver, out);
364    for x in [&f.r#where, &f.frontier, &f.by].into_iter().flatten() {
365        visit_expr(x, out);
366    }
367}
368
369fn visit_sub(s: &Subquery, out: &mut Vec<String>) {
370    s.from.iter().for_each(|e| visit_expr(e, out));
371    if let Some(w) = &s.r#where {
372        visit_where(w, out);
373    }
374    visit_select(&s.select, out);
375    if let Some(o) = &s.order_by {
376        o.iter().for_each(|spec| visit_expr(&spec.expr, out));
377    }
378    if let Some(f) = &s.follow {
379        visit_follow(f, out);
380    }
381}
382
383/// The distinct phrases `semantic("…")` names (free calls in the raw parse);
384/// empty when the source does not parse.
385#[must_use]
386pub fn collect_semantic_phrases(source: &str) -> Vec<String> {
387    let Ok(q) = oqx::parse_string(source) else {
388        return Vec::new();
389    };
390    let mut out = Vec::new();
391    visit_expr(&q.source, &mut out);
392    q.from.iter().for_each(|e| visit_expr(e, &mut out));
393    if let Some(w) = &q.r#where {
394        visit_where(w, &mut out);
395    }
396    visit_select(&q.select, &mut out);
397    if let Some(o) = &q.order_by {
398        o.iter().for_each(|spec| visit_expr(&spec.expr, &mut out));
399    }
400    if let Some(f) = &q.follow {
401        visit_follow(f, &mut out);
402    }
403    out
404}
405
406// ---- hits ----------------------------------------------------------------------------
407
408/// A projected row as a hit: `{ id, path, ...rest }` with the injected
409/// columns peeled off (JavaScript's `String()` on the id, `""` for an absent
410/// path). The projection's values are rendered per §1.4 "rows as values": a
411/// store row nested in the result (an empty-projection `collect { }` and
412/// friends) becomes `{ id, path }`.
413fn to_hit(row: Value) -> Value {
414    let Value::Object(o) = render_row_values(row) else {
415        return Value::Object(oqx::Object::new());
416    };
417    let mut id = Value::Undefined;
418    let mut path = Value::Undefined;
419    let mut rest = Vec::new();
420    for (k, v) in o {
421        match k.as_str() {
422            ID_KEY => id = v,
423            PATH_KEY => path = v,
424            _ => rest.push((k, v)),
425        }
426    }
427    let mut hit = oqx::Object::with_capacity(rest.len() + 2);
428    hit.insert("id", Value::Str(id.to_string()));
429    hit.insert(
430        "path",
431        Value::Str(if path.is_absent() {
432            String::new()
433        } else {
434            path.to_string()
435        }),
436    );
437    for (k, v) in rest {
438        hit.insert(k, v);
439    }
440    Value::Object(hit)
441}
442
443fn hit_str(hit: &Value, key: &str) -> String {
444    hit.as_object()
445        .and_then(|o| o.get(key))
446        .map(|v| v.to_string())
447        .unwrap_or_default()
448}
449
450/// Top-level `select distinct`: dedup hits by their USER projection (every
451/// field but `id`/`path`), keeping the first. The key is the canonical JSON
452/// of the sorted `[key, value]` pairs (`JSON.stringify` in the reference).
453fn dedup_hits_by_projection(hits: Vec<Value>) -> Vec<Value> {
454    let mut seen: Vec<String> = Vec::new();
455    let mut out = Vec::new();
456    for h in hits {
457        let mut pairs: Vec<(String, Value)> = h
458            .as_object()
459            .map(|o| {
460                o.iter()
461                    .filter(|(k, _)| *k != "id" && *k != "path")
462                    .map(|(k, v)| (k.to_owned(), v.clone()))
463                    .collect()
464            })
465            .unwrap_or_default();
466        pairs.sort_by(|a, b| a.0.cmp(&b.0));
467        let key = Value::Array(
468            pairs
469                .into_iter()
470                .map(|(k, v)| Value::Array(vec![Value::Str(k), v]))
471                .collect(),
472        )
473        .to_canonical_json()
474        .to_string();
475        if seen.contains(&key) {
476            continue;
477        }
478        seen.push(key);
479        out.push(h);
480    }
481    out
482}
483
484/// A top-level `limit`/`offset` on the collect path is applied by the runner,
485/// so it must be a plain non-negative integer literal.
486fn const_bound(e: Option<&Expr>, word: &str) -> Result<Option<usize>> {
487    match e {
488        None => Ok(None),
489        Some(Expr::Lit(Value::Number(n))) if n.fract() == 0.0 && *n >= 0.0 && n.is_finite() => {
490            Ok(Some(*n as usize))
491        }
492        Some(_) => Err(SurfaceError::filter_invalid(
493            format!("top-level {word} must be a non-negative integer literal"),
494            "OQX",
495        )),
496    }
497}
498
499fn value_of(hit: &Value) -> Value {
500    hit.as_object()
501        .and_then(|o| o.get(VALUE_KEY))
502        .cloned()
503        .unwrap_or(Value::Undefined)
504}
505
506fn without_value_key(hit: Value) -> Json {
507    hit.to_canonical_json()
508}
509
510// ---- the run --------------------------------------------------------------------------
511
512/// Run an OQX query against `repo_id` (§1.4).
513pub fn query(
514    store: &Store,
515    repo_id: &str,
516    source: &str,
517    opts: QueryOptions<'_>,
518) -> Result<OqxResult> {
519    // Without a provider the phrases stay unembedded and the context reports
520    // `filter_invalid` ("needs an embedding provider") when one is reached —
521    // after its target check, so `semantic()` on nodes names the targets
522    // (§9; the `query` tool's pre-check is what reports `semantic_unavailable`).
523    let phrases = collect_semantic_phrases(source);
524    let mut semantic: HashMap<String, SemanticVec> = HashMap::new();
525    if let Some(provider) = opts.provider.filter(|_| !phrases.is_empty()) {
526        for phrase in phrases {
527            let vec = provider
528                .embed_query(&phrase)
529                .map_err(|e| SurfaceError::new(e.code(), e.to_string()))?;
530            semantic.insert(
531                phrase,
532                SemanticVec {
533                    model: provider.model().to_owned(),
534                    vec: f32_to_blob(&vec),
535                },
536            );
537        }
538    }
539    let runner = Runner {
540        store,
541        repo_id,
542        semantic,
543        planned: !opts.in_memory,
544    };
545    run_inner(&runner, source, opts)
546}
547
548/// One query's engine: a fresh store context per run (the planned path gives
549/// the residual a context serving the produced rows as its root). A failure
550/// inside a property read or row function is the engine's own error (the
551/// context's `get` / `call_method` return `Err`); only a failed root scan,
552/// which the `root` seam cannot raise, is read back after the run.
553struct Runner<'a> {
554    store: &'a Store,
555    repo_id: &'a str,
556    semantic: HashMap<String, SemanticVec>,
557    planned: bool,
558}
559
560impl Runner<'_> {
561    /// Tier-3 pushdown reduces the scan in SQL and the in-memory engine
562    /// finishes the residual over the produced rows (a declined plan, or
563    /// `in_memory`, runs the whole query in memory over a full scan), so
564    /// results match a pure scan. This is `oqx::PlannedEngine::run` inlined:
565    /// the store context borrows the connection, so it cannot be the
566    /// `'static` context a `Plan` carries.
567    fn run(&self, q: &Query) -> Result<oqx::OqxResult> {
568        let conn = self.store.conn();
569        let ctx = StoreContext::new(conn, self.repo_id, self.semantic.clone());
570        let plan = if self.planned {
571            SqlitePlanner::new(conn, self.repo_id)
572                .try_plan(q, &[])
573                .map_err(|e| SurfaceError::other(format!("sqlite: {e}")))?
574        } else {
575            None
576        };
577        let (ctx, residual) = match plan {
578            Some(plan) => (ctx.with_rows_root(plan.rows), Some(plan.residual)),
579            None => (ctx, None),
580        };
581        let engine = InMemoryEngine::new(ctx);
582        let out = engine.run(residual.as_ref().unwrap_or(q), &[]);
583        // `root` has no error channel: a store failure during a root scan was
584        // served as an empty scan and wins over whatever the run made of it.
585        if let Some(failed) = engine.context().take_root_failure() {
586            return Err(failed.into());
587        }
588        Ok(out?)
589    }
590}
591
592fn run_inner(engine: &Runner<'_>, source: &str, opts: QueryOptions<'_>) -> Result<OqxResult> {
593    let parsed = rewrite_query(&oqx::parse_string(source)?);
594    let consumer = parsed.consumer;
595
596    match consumer {
597        Consumer::Exists => {
598            let res = engine.run(&parsed)?;
599            let mut r = OqxResult::scalar(consumer);
600            r.exists = Some(matches!(res, oqx::OqxResult::Exists(true)));
601            return Ok(r);
602        }
603        Consumer::Count => {
604            let res = engine.run(&parsed)?;
605            let mut r = OqxResult::scalar(consumer);
606            r.count = Some(match res {
607                oqx::OqxResult::Count(n) => n,
608                _ => 0.0,
609            });
610            return Ok(r);
611        }
612        Consumer::None => {
613            let res = engine.run(&parsed)?;
614            let mut r = OqxResult::scalar(consumer);
615            r.none = Some(match res {
616                oqx::OqxResult::None(b) => b,
617                _ => true,
618            });
619            return Ok(r);
620        }
621        Consumer::Collect | Consumer::First | Consumer::Single => {}
622    }
623
624    // collect / first / single: inject id + path so every hit carries them. A
625    // top-level `select distinct` is applied HERE, not in the engine (the
626    // injected id/path are unique per row and would defeat the engine's
627    // projection dedup). A top-level `values` projection runs as a RECORD
628    // projection whose single item is renamed to VALUE_KEY.
629    let top_distinct = parsed.distinct;
630    let top_values = parsed.values;
631    let user_select: Vec<SelectItem> = if top_values {
632        parsed
633            .select
634            .first()
635            .map(|it| match it {
636                SelectItem::Field { expr, lift, .. } => SelectItem::Field {
637                    name: VALUE_KEY.to_owned(),
638                    expr: expr.clone(),
639                    lift: *lift,
640                },
641                SelectItem::Collect { op, .. } => SelectItem::Collect {
642                    name: VALUE_KEY.to_owned(),
643                    op: op.clone(),
644                },
645            })
646            .into_iter()
647            .collect()
648    } else {
649        parsed.select.clone()
650    };
651    let id_item = SelectItem::Field {
652        name: ID_KEY.to_owned(),
653        expr: Expr::Ident {
654            name: "$id".to_owned(),
655        },
656        lift: 0,
657    };
658    let path_item = SelectItem::Field {
659        name: PATH_KEY.to_owned(),
660        expr: Expr::Ident {
661            name: "$path".to_owned(),
662        },
663        lift: 0,
664    };
665    let mut select = vec![id_item, path_item];
666    select.extend(user_select);
667    // On the collect path the query's own limit/offset is taken out of the
668    // engine query and applied after the runner's distinct; first/single keep
669    // theirs (the engine's offset-aware cap is exactly right for them).
670    let (top_limit, top_offset) = (parsed.limit.clone(), parsed.offset.clone());
671    let q = Query {
672        distinct: false,
673        values: false,
674        select,
675        limit: if consumer == Consumer::Collect {
676            None
677        } else {
678            parsed.limit.clone()
679        },
680        offset: if consumer == Consumer::Collect {
681            None
682        } else {
683            parsed.offset.clone()
684        },
685        ..parsed.clone()
686    };
687    let res = engine.run(&q)?;
688
689    if matches!(consumer, Consumer::First | Consumer::Single) {
690        let row = match res {
691            oqx::OqxResult::First(r) | oqx::OqxResult::Single(r) => r,
692            _ => None,
693        };
694        let mut out = OqxResult::scalar(consumer);
695        match row {
696            None => {
697                if top_values {
698                    out.values = Some(Vec::new());
699                }
700            }
701            Some(r) => {
702                let hit = to_hit(r);
703                if top_values {
704                    out.values = Some(vec![value_of(&hit).to_canonical_json()]);
705                } else {
706                    out.hits = vec![without_value_key(hit)];
707                }
708            }
709        }
710        return Ok(out);
711    }
712
713    // collect: keyset pagination on (path, id) when the order is the default.
714    let mut rows: Vec<Value> = match res {
715        oqx::OqxResult::Collect(rows) => rows.into_iter().map(to_hit).collect(),
716        _ => Vec::new(),
717    };
718    if top_distinct {
719        rows = dedup_hits_by_projection(rows);
720    }
721    let offset = const_bound(top_offset.as_ref(), "offset")?.unwrap_or(0);
722    let limit = const_bound(top_limit.as_ref(), "limit")?;
723    if offset > 0 || limit.is_some() {
724        let end = limit.map_or(rows.len(), |l| (offset + l).min(rows.len()));
725        rows = if offset >= rows.len() {
726            Vec::new()
727        } else {
728            rows[offset..end].to_vec()
729        };
730    }
731    let custom = parsed.order_by.as_ref().is_some_and(|o| !o.is_empty());
732    let cap = opts.limit.unwrap_or(DEFAULT_LIMIT);
733    let mut page = rows;
734    if !custom {
735        if let Some(cursor) = opts.cursor.filter(|c| !c.is_empty()) {
736            let parts = decode_cursor(cursor, "query", 2)?;
737            let (path, id) = (&parts[0], &parts[1]);
738            page.retain(|h| {
739                let hp = hit_str(h, "path");
740                let hi = hit_str(h, "id");
741                hp > *path || (hp == *path && hi > *id)
742            });
743        }
744    }
745    let truncated = page.len() > cap;
746    page.truncate(cap);
747    let cursor = if truncated && !custom {
748        page.last()
749            .map(|last| encode_cursor(&[&hit_str(last, "path"), &hit_str(last, "id")]))
750    } else {
751        None
752    };
753    let mut out = OqxResult::scalar(Consumer::Collect);
754    out.truncated = truncated;
755    out.cursor = cursor;
756    if top_values {
757        out.values = Some(
758            page.iter()
759                .map(|h| value_of(h).to_canonical_json())
760                .collect(),
761        );
762    } else {
763        out.hits = page.into_iter().map(without_value_key).collect();
764    }
765    Ok(out)
766}
767
768#[cfg(test)]
769mod tests {
770    use super::*;
771
772    #[test]
773    fn row_functions_become_self_methods() {
774        let q = oqx::parse_string(
775            "from blocks where text(\"x\") && under_heading(\"h\") && size(attrs) > 0 && doc.$path.startsWith(\"a\")",
776        )
777        .unwrap();
778        let r = rewrite_query(&q);
779        let Some(Where::And { parts }) = &r.r#where else {
780            panic!("and")
781        };
782        let Where::Scalar { expr } = &parts[0] else {
783            panic!("scalar")
784        };
785        assert!(
786            matches!(expr, Expr::Call { recv: Some(r), name, .. } if name == "text" && **r == self_ref())
787        );
788        let Where::Scalar { expr } = &parts[2] else {
789            panic!("scalar")
790        };
791        assert!(
792            matches!(expr, Expr::Binary { left, .. } if matches!(&**left, Expr::Call { recv: None, name, .. } if name == "size"))
793        );
794    }
795
796    #[test]
797    fn semantic_phrases_are_collected_distinct() {
798        let phrases = collect_semantic_phrases(
799            "select s: semantic(\"alpha\") from docs where semantic(\"alpha\") > 0.5 || nodes exists { where semantic(\"beta\") > 0 } order by semantic(\"gamma\") desc",
800        );
801        assert_eq!(phrases, ["alpha", "beta", "gamma"]);
802        assert!(collect_semantic_phrases("not a query {{").is_empty());
803        assert!(collect_semantic_phrases("from docs").is_empty());
804    }
805
806    #[test]
807    fn hits_peel_the_injected_columns() {
808        let mut o = oqx::Object::new();
809        o.insert(ID_KEY, Value::Str("d_1".into()));
810        o.insert(PATH_KEY, Value::Null);
811        o.insert("layer", Value::Str("canon".into()));
812        let hit = to_hit(Value::Object(o));
813        let ho = hit.as_object().unwrap();
814        assert_eq!(ho.keys().collect::<Vec<_>>(), ["id", "path", "layer"]);
815        assert_eq!(ho.get("path"), Some(&Value::Str(String::new())));
816        // A user field named `id` overrides the injected one in place.
817        let mut o = oqx::Object::new();
818        o.insert(ID_KEY, Value::Str("d_1".into()));
819        o.insert(PATH_KEY, Value::Str("a.md".into()));
820        o.insert("id", Value::Number(7.0));
821        let hit = to_hit(Value::Object(o));
822        let ho = hit.as_object().unwrap();
823        assert_eq!(ho.keys().collect::<Vec<_>>(), ["id", "path"]);
824        assert_eq!(ho.get("id"), Some(&Value::Number(7.0)));
825    }
826
827    #[test]
828    fn distinct_dedups_by_user_projection_first_wins() {
829        let mk = |id: &str, t: &str| {
830            let mut o = oqx::Object::new();
831            o.insert("id", Value::Str(id.into()));
832            o.insert("path", Value::Str("p".into()));
833            o.insert("type", Value::Str(t.into()));
834            Value::Object(o)
835        };
836        let out = dedup_hits_by_projection(vec![mk("1", "a"), mk("2", "b"), mk("3", "a")]);
837        assert_eq!(out.len(), 2);
838        assert_eq!(hit_str(&out[0], "id"), "1");
839        assert_eq!(hit_str(&out[1], "id"), "2");
840    }
841
842    #[test]
843    fn top_level_bounds_must_be_literals() {
844        assert_eq!(const_bound(None, "limit").unwrap(), None);
845        assert_eq!(
846            const_bound(Some(&Expr::Lit(Value::Number(3.0))), "limit").unwrap(),
847            Some(3)
848        );
849        let e = const_bound(Some(&Expr::Lit(Value::Number(-1.0))), "offset").unwrap_err();
850        assert_eq!(e.code, "filter_invalid");
851        assert!(e.message.contains("top-level offset"));
852    }
853}