omgbase_surface/
graph.rs

1//! The `graph` neighborhood macro (`spec/surface/README.md` §4): compiled to
2//! an OQX `follow _out_links collect { from _dst }` / `follow _in_links
3//! collect { from _doc }` walk over the link rows (3.0 follow-up: there are no
4//! `out`/`in` relations — `_dst` is a link's destination document, `_doc` its
5//! source; an absent `_dst`, an external or phantom target, contributes no
6//! successor) run through the shared runner and
7//! shaped into `{ documents, links, frontier }`. Port of
8//! `packages/core/src/mcp/graph.ts`; nothing here walks the graph itself.
9
10use std::cmp::Ordering;
11use std::collections::HashMap;
12
13use omgbase_search::EmbeddingProvider;
14use omgbase_store::Store;
15use oqx::ast::{BinaryOp, Consumer, FollowDestination, SelectItem, Where};
16use oqx::{build, print_query};
17use serde_json::{Map, Value as Json, json};
18
19use crate::error::{Result, SurfaceError};
20use crate::query::{QueryOptions, query};
21use crate::read::find_doc_by_ref;
22
23const DEFAULT_DEGREES: i64 = 1;
24const DEFAULT_MAX_DOCUMENTS: i64 = 200;
25const MAX_DEPTH: i64 = 8;
26
27/// `graph`'s arguments.
28#[derive(Clone, Debug, Default)]
29pub struct GraphArgs {
30    pub roots: Vec<String>,
31    pub degrees: Option<i64>,
32    /// `in` | `out` | `both` (default).
33    pub direction: Option<String>,
34    pub predicate: Option<String>,
35    pub select: Vec<String>,
36    pub max_documents: Option<i64>,
37}
38
39/// The edge fields projected from `_out_links`/`_in_links` — the `links`
40/// target's own names, projected under themselves (`spec/surface` §1.2, 5.0:
41/// no rename between the language and the result) — identical for both
42/// directions so a single filter handles either scan (an edge is a directional
43/// fact regardless of which endpoint the walk reached it from).
44const EDGE_FIELDS: [&str; 12] = [
45    "_id",
46    "_doc_id",
47    "_block_id",
48    "_via_id",
49    "_dst_id",
50    "_dst_path",
51    "_uri",
52    "_dst_kind",
53    "_predicate",
54    "_provenance",
55    "_anchor",
56    "_field",
57];
58
59// The macro's own aliases: `__oqx_`/`__omg_` (two underscores) are the
60// runner's and the selector compiler's reserved prefixes; a single underscore
61// is the engine prefix of 5.0, which an ALIAS may use freely (an alias is never
62// a row read), and keeps these off any authored key.
63const DEPTH_ALIAS: &str = "_depth";
64const STOP_ALIAS: &str = "_stop";
65const LINKS_ALIAS: &str = "_links";
66
67/// The extra `select` items for caller projections. Each expression gets a
68/// stable internal alias (`_u0`, `_u1`, …) so it can never collide with a
69/// reserved word or the macro's own aliases; the clean output name is derived
70/// separately. An expression is parsed as the item of a projection, so it may
71/// be anything a `select` item may be and nothing more; a malformed one is
72/// `filter_invalid`, as it would be in the `query` tool.
73fn build_user_select(select: &[String]) -> Result<(Vec<SelectItem>, Vec<Option<String>>)> {
74    let mut items = Vec::new();
75    let mut out_names: Vec<Option<String>> = vec![None; select.len()];
76    let mut used: Vec<String> = Vec::new();
77    for (i, expr) in select.iter().enumerate() {
78        let trimmed = expr.trim();
79        if trimmed.is_empty() {
80            continue;
81        }
82        // An identifier keeps its spelling as the output key (5.0: a `layer`
83        // or `_title` selection is its own key in the document, as a `query`
84        // hit's would be; an OQX `$` name too); anything else is `sel_<i>`,
85        // so a projection can never overwrite the envelope here.
86        let ident = trimmed.strip_prefix('$').unwrap_or(trimmed);
87        let is_ident = ident
88            .chars()
89            .next()
90            .is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
91            && ident.chars().all(|c| c.is_ascii_alphanumeric() || c == '_');
92        let mut name = if is_ident {
93            trimmed.to_owned()
94        } else {
95            format!("sel_{i}")
96        };
97        // The envelope keys are the row's own identity: an `_id`/`_path`
98        // selection lands on them rather than beside them (the engine reads
99        // the same value).
100        if name != "_id" && name != "_path" {
101            while used.contains(&name) {
102                name = format!("{name}_{i}");
103            }
104            used.push(name.clone());
105        }
106        out_names[i] = Some(name);
107        let parsed = oqx::parse_string(&format!("select _u{i}: {trimmed} from docs"))?;
108        items.extend(parsed.select);
109    }
110    Ok((items, out_names))
111}
112
113/// The walk, as an AST (`spec/oqx/AST.md` §7) printed canonically — never
114/// spliced strings: the seed ids are literals, the caller's projections parsed
115/// items.
116///
117/// ```text
118/// select _depth: $depth, _stop: $stop, _links: _<dir>_links collect { <edge fields> } <user items>
119/// from docs where _id == <root> or … follow distinct _<dir>_links collect { from <_dst | _doc> depth <depth> }
120/// ```
121fn build_query(root_ids: &[String], dir: &str, depth: u32, user_select: &[SelectItem]) -> String {
122    let mut edges_body = build::subquery();
123    edges_body.select = EDGE_FIELDS
124        .iter()
125        .map(|name| build::field(name, build::ident(name)))
126        .collect();
127    let edges = build::op(
128        build::ident(&format!("_{dir}_links")),
129        Consumer::Collect,
130        edges_body,
131    );
132    let mut seed: Vec<Where> = root_ids
133        .iter()
134        .map(|id| {
135            build::scalar(build::binary(
136                BinaryOp::Eq,
137                build::ident("_id"),
138                build::lit(id.as_str()),
139            ))
140        })
141        .collect();
142    let mut q = build::query(build::ident("docs"));
143    q.select = vec![
144        build::field(DEPTH_ALIAS, build::ident("$depth")),
145        build::field(STOP_ALIAS, build::ident("$stop")),
146        build::collect(LINKS_ALIAS, edges),
147    ];
148    q.select.extend(user_select.iter().cloned());
149    q.r#where = Some(if seed.len() == 1 {
150        seed.pop().expect("one seed")
151    } else {
152        build::or(seed)
153    });
154    // `depth` is the destination block's own clause (oqx 0.20).
155    let mut step = build::subquery();
156    step.from = vec![build::ident(if dir == "out" { "_dst" } else { "_doc" })];
157    step.depth = Some(depth);
158    let mut walk = build::follow(vec![FollowDestination::Block(Box::new(build::op(
159        build::ident(&format!("_{dir}_links")),
160        Consumer::Collect,
161        step,
162    )))]);
163    walk.distinct = true;
164    q.follow = Some(walk);
165    print_query(&q).expect("a built query has no bindings")
166}
167
168/// Bytewise order (§9: the reference's `localeCompare` was replaced).
169fn locale_compare(a: &str, b: &str) -> Ordering {
170    a.cmp(b)
171}
172
173struct Doc {
174    id: String,
175    path: String,
176    degree: i64,
177    extra: Vec<(String, Json)>,
178}
179
180/// Run the macro.
181pub fn graph_neighborhood(
182    store: &Store,
183    repo_id: &str,
184    args: &GraphArgs,
185    provider: Option<&dyn EmbeddingProvider>,
186    now_ms: Option<i64>,
187) -> Result<Json> {
188    if args.roots.is_empty() {
189        return Err(SurfaceError::new(
190            "target_missing",
191            "graph requires at least one root (path or id)",
192        ));
193    }
194    let mut root_ids: Vec<String> = Vec::new();
195    for r in &args.roots {
196        let Some(info) = find_doc_by_ref(store.conn(), repo_id, r)? else {
197            return Err(SurfaceError::with_data(
198                "doc_missing",
199                format!("no document for {}", Json::String(r.clone())),
200                json!({ "root": r }),
201            ));
202        };
203        if !root_ids.contains(&info.doc_id) {
204            root_ids.push(info.doc_id);
205        }
206    }
207    let degrees = args.degrees.unwrap_or(DEFAULT_DEGREES).max(0);
208    let depth = MAX_DEPTH.min(degrees + 1);
209    let walk_depth = u32::try_from(depth).expect("1..=8");
210    let effective_degrees = depth - 1;
211    let direction = args.direction.clone().unwrap_or_else(|| "both".to_owned());
212    let max_documents =
213        usize::try_from(args.max_documents.unwrap_or(DEFAULT_MAX_DOCUMENTS).max(1)).unwrap_or(1);
214    let dirs: Vec<&str> = match direction.as_str() {
215        "both" => vec!["out", "in"],
216        "in" => vec!["in"],
217        _ => vec!["out"],
218    };
219    let (user_select, out_names) = build_user_select(&args.select)?;
220
221    let mut queries = Vec::new();
222    let mut docs: Vec<Doc> = Vec::new();
223    let mut edges: Vec<Json> = Vec::new();
224    let mut query_truncated = false;
225    for dir in &dirs {
226        let q = build_query(&root_ids, dir, walk_depth, &user_select);
227        queries.push(q.clone());
228        let res = query(
229            store,
230            repo_id,
231            &q,
232            QueryOptions {
233                limit: Some(max_documents + 1),
234                cursor: None,
235                provider,
236                in_memory: false,
237                bare: false,
238                unknown_names: None,
239                vocabulary: None,
240                now_ms,
241                request: None,
242            },
243        )?;
244        if res.truncated {
245            query_truncated = true;
246        }
247        for hit in &res.hits {
248            // The walk's hits carry `_id`/`_path` (§1.4, 5.0), and so does
249            // `graph`'s own envelope (§4): one spelling for a document's identity.
250            let id = hit["_id"].as_str().unwrap_or_default().to_owned();
251            let path = hit["_path"].as_str().unwrap_or_default().to_owned();
252            let hop = hit[DEPTH_ALIAS].as_f64().unwrap_or(f64::NAN);
253            let degree = (hop - 1.0) as i64;
254            let replace = docs
255                .iter()
256                .position(|d| d.id == id)
257                .map(|i| (i, degree < docs[i].degree));
258            match replace {
259                Some((_, false)) => {}
260                found => {
261                    let extra: Vec<(String, Json)> = out_names
262                        .iter()
263                        .enumerate()
264                        .filter_map(|(i, n)| {
265                            n.as_ref().map(|name| {
266                                (
267                                    name.clone(),
268                                    hit.get(format!("_u{i}")).cloned().unwrap_or(Json::Null),
269                                )
270                            })
271                        })
272                        .collect();
273                    let doc = Doc {
274                        id: id.clone(),
275                        path,
276                        degree,
277                        extra,
278                    };
279                    match found {
280                        Some((i, true)) => docs[i] = doc,
281                        _ => docs.push(doc),
282                    }
283                }
284            }
285            if let Some(raw) = hit[LINKS_ALIAS].as_array() {
286                for e in raw {
287                    let eid = e["_id"].as_str().unwrap_or_default();
288                    if !edges.iter().any(|x| x["_id"].as_str() == Some(eid)) {
289                        edges.push(e.clone());
290                    }
291                }
292            }
293        }
294    }
295
296    if let Some(p) = &args.predicate {
297        let mut adj: HashMap<String, Vec<String>> = HashMap::new();
298        for e in &edges {
299            if e["_predicate"].as_str() != Some(p.as_str()) {
300                continue;
301            }
302            let src = e["_doc_id"].as_str().unwrap_or_default().to_owned();
303            let dst = e["_dst_id"].as_str().unwrap_or_default().to_owned();
304            if dirs.contains(&"out") {
305                adj.entry(src.clone()).or_default().push(dst.clone());
306            }
307            if dirs.contains(&"in") {
308                adj.entry(dst).or_default().push(src);
309            }
310        }
311        let mut depth_of: Vec<(String, i64)> = root_ids.iter().map(|id| (id.clone(), 0)).collect();
312        let mut wave: Vec<String> = root_ids.clone();
313        let mut lvl = 1;
314        while lvl <= effective_degrees && !wave.is_empty() {
315            let mut next = Vec::new();
316            for from in &wave {
317                for to in adj.get(from).map_or(&[][..], Vec::as_slice) {
318                    if docs.iter().any(|d| &d.id == to) && !depth_of.iter().any(|(id, _)| id == to)
319                    {
320                        depth_of.push((to.clone(), lvl));
321                        next.push(to.clone());
322                    }
323                }
324            }
325            wave = next;
326            lvl += 1;
327        }
328        let mut restricted: Vec<Doc> = Vec::new();
329        for (id, d) in depth_of {
330            if let Some(orig) = docs.iter().find(|x| x.id == id) {
331                restricted.push(Doc {
332                    id: orig.id.clone(),
333                    path: orig.path.clone(),
334                    degree: d,
335                    extra: orig.extra.clone(),
336                });
337            }
338        }
339        docs = restricted;
340    }
341
342    docs.sort_by(|a, b| {
343        a.degree
344            .cmp(&b.degree)
345            .then_with(|| locale_compare(&a.path, &b.path))
346    });
347    let capped = docs.len() > max_documents;
348    docs.truncate(max_documents);
349    let truncated = query_truncated || capped;
350
351    let mut frontier = Vec::new();
352    let documents: Vec<Json> = docs
353        .iter()
354        .map(|d| {
355            let is_frontier = d.degree == effective_degrees;
356            if is_frontier {
357                frontier.push(json!({ "_id": d.id, "_path": d.path, "degree": d.degree }));
358            }
359            let mut m = Map::new();
360            m.insert("_id".to_owned(), json!(d.id));
361            m.insert("_path".to_owned(), json!(d.path));
362            m.insert("degree".to_owned(), json!(d.degree));
363            m.insert("frontier".to_owned(), json!(is_frontier));
364            for (k, v) in &d.extra {
365                m.insert(k.clone(), v.clone());
366            }
367            Json::Object(m)
368        })
369        .collect();
370
371    let reached = |id: &str| docs.iter().any(|d| d.id == id);
372    let mut kept: Vec<Json> = edges
373        .into_iter()
374        .filter(|e| {
375            if let Some(p) = &args.predicate {
376                if e["_predicate"].as_str() != Some(p.as_str()) {
377                    return false;
378                }
379            }
380            let src_in = reached(e["_doc_id"].as_str().unwrap_or_default());
381            let dst_kind = e["_dst_kind"].as_str().unwrap_or_default();
382            let dst_dangling =
383                dst_kind == "external" || (dst_kind == "document" && e["_dst_path"].is_null());
384            let dst_in = reached(e["_dst_id"].as_str().unwrap_or_default()) || dst_dangling;
385            src_in && dst_in
386        })
387        .collect();
388    kept.sort_by(|a, b| {
389        let (sa, sb) = (
390            a["_doc_id"].as_str().unwrap_or_default(),
391            b["_doc_id"].as_str().unwrap_or_default(),
392        );
393        if sa == sb {
394            locale_compare(
395                a["_id"].as_str().unwrap_or_default(),
396                b["_id"].as_str().unwrap_or_default(),
397            )
398        } else {
399            locale_compare(sa, sb)
400        }
401    });
402
403    Ok(json!({
404        "roots": root_ids,
405        "degrees": effective_degrees,
406        "direction": direction,
407        "documents": documents,
408        "links": kept,
409        "frontier": frontier,
410        "truncated": truncated,
411        "queries": queries,
412    }))
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    #[test]
420    fn user_select_aliases() {
421        let (items, names) = build_user_select(&[
422            "layer".into(),
423            "_path".into(),
424            "a + 1".into(),
425            "layer".into(),
426        ])
427        .unwrap();
428        let printed: Vec<String> = items
429            .iter()
430            .map(|it| oqx::print(oqx::Node::Select(it)).unwrap())
431            .collect();
432        assert_eq!(
433            printed,
434            ["_u0: layer", "_u1: _path", "_u2: a + 1", "_u3: layer"]
435        );
436        assert_eq!(
437            names,
438            [
439                Some("layer".into()),
440                Some("_path".into()),
441                Some("sel_2".into()),
442                Some("layer_3".into())
443            ]
444        );
445        assert!(build_user_select(&[]).unwrap().0.is_empty());
446        let e = build_user_select(&["a +".into()]).unwrap_err();
447        assert_eq!(e.code, "filter_invalid");
448    }
449
450    #[test]
451    fn query_shape() {
452        let q = build_query(&["d_0".to_owned()], "out", 2, &[]);
453        assert!(q.starts_with("select _depth: $depth, _stop: $stop, _links: _out_links collect {"));
454        assert!(q.ends_with(
455            "from docs where _id == \"d_0\" follow distinct _out_links collect { from _dst depth 2 }"
456        ));
457        assert!(oqx::parse_string(&q).is_ok());
458        // several roots: an `or` seed (the canonical print of the `or` node since
459        // oqx 0.18); the generated text is exactly the macro's pinned shape
460        // (spec/surface `reads.json`)
461        let q = build_query(&["d_14".to_owned(), "d_15".to_owned()], "in", 8, &[]);
462        assert_eq!(
463            q,
464            "select _depth: $depth, _stop: $stop, _links: _in_links collect { _id: _id, _doc_id: _doc_id, _block_id: _block_id, _via_id: _via_id, _dst_id: _dst_id, _dst_path: _dst_path, _uri: _uri, _dst_kind: _dst_kind, _predicate: _predicate, _provenance: _provenance, _anchor: _anchor, _field: _field } from docs where _id == \"d_14\" or _id == \"d_15\" follow distinct _in_links collect { from _doc depth 8 }"
465        );
466    }
467
468    #[test]
469    fn order_is_bytewise() {
470        assert_eq!(locale_compare("B", "a"), Ordering::Less);
471        assert_eq!(locale_compare("a", "b"), Ordering::Less);
472    }
473}
omgbase_surface/graph.rs

omgbase_surface/
graph.rs
