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

omgbase_surface/
graph.rs
