omgbase_surface/
graph.rs

1//! The `graph` neighborhood macro (`spec/surface/README.md` §4): compiled to
2//! an OQX `follow doc.out` / `doc.in` walk run through the shared runner and
3//! shaped into `{ documents, edges, frontier }`. Port of
4//! `packages/core/src/mcp/graph.ts`; nothing here walks the graph itself.
5
6use std::cmp::Ordering;
7use std::collections::HashMap;
8
9use omgbase_search::EmbeddingProvider;
10use omgbase_store::Store;
11use serde_json::{Map, Value as Json, json};
12
13use crate::error::{Result, SurfaceError};
14use crate::query::{QueryOptions, query};
15use crate::read::find_doc_by_ref;
16
17const DEFAULT_DEGREES: i64 = 1;
18const DEFAULT_MAX_DOCUMENTS: i64 = 200;
19const MAX_DEPTH: i64 = 8;
20
21/// `graph`'s arguments.
22#[derive(Clone, Debug, Default)]
23pub struct GraphArgs {
24    pub roots: Vec<String>,
25    pub degrees: Option<i64>,
26    /// `in` | `out` | `both` (default).
27    pub direction: Option<String>,
28    pub predicate: Option<String>,
29    pub select: Vec<String>,
30    pub max_documents: Option<i64>,
31}
32
33const EDGE_COLLECT: &str = "{ id: $id, src: $src, dst: $dst, dst_path: $dst_path, dst_uri: $dst_uri, dst_kind: dst_kind, predicate: predicate, provenance: provenance, anchor: anchor, src_field: src_field }";
34
35fn build_user_select(select: &[String]) -> (String, Vec<Option<String>>) {
36    let mut items = Vec::new();
37    let mut out_names: Vec<Option<String>> = vec![None; select.len()];
38    let mut used: Vec<String> = Vec::new();
39    for (i, expr) in select.iter().enumerate() {
40        let trimmed = expr.trim();
41        if trimmed.is_empty() {
42            continue;
43        }
44        let ident = trimmed.strip_prefix('$').unwrap_or(trimmed);
45        let is_ident = ident
46            .chars()
47            .next()
48            .is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
49            && ident.chars().all(|c| c.is_ascii_alphanumeric() || c == '_');
50        let mut name = if is_ident {
51            ident.to_owned()
52        } else {
53            format!("sel_{i}")
54        };
55        while used.contains(&name) {
56            name = format!("{name}_{i}");
57        }
58        used.push(name.clone());
59        out_names[i] = Some(name);
60        items.push(format!("_u{i}: {trimmed}"));
61    }
62    let clause = if items.is_empty() {
63        String::new()
64    } else {
65        format!(", {}", items.join(", "))
66    };
67    (clause, out_names)
68}
69
70fn build_query(seed: &str, dir: &str, depth: i64, user_select: &str) -> String {
71    format!(
72        "select _depth: $depth, _stop: $stop, _edges: doc.{dir}_edges collect {EDGE_COLLECT}{user_select} from docs where {seed} follow distinct doc.{dir} {{ depth {depth} }}"
73    )
74}
75
76/// Bytewise order (§9: the reference's `localeCompare` was replaced).
77fn locale_compare(a: &str, b: &str) -> Ordering {
78    a.cmp(b)
79}
80
81struct Doc {
82    id: String,
83    path: String,
84    degree: i64,
85    extra: Vec<(String, Json)>,
86}
87
88/// Run the macro.
89pub fn graph_neighborhood(
90    store: &Store,
91    repo_id: &str,
92    args: &GraphArgs,
93    provider: Option<&dyn EmbeddingProvider>,
94) -> Result<Json> {
95    if args.roots.is_empty() {
96        return Err(SurfaceError::new(
97            "target_missing",
98            "graph requires at least one root (path or id)",
99        ));
100    }
101    let mut root_ids: Vec<String> = Vec::new();
102    for r in &args.roots {
103        let Some(info) = find_doc_by_ref(store.conn(), repo_id, r)? else {
104            return Err(SurfaceError::with_data(
105                "doc_missing",
106                format!("no document for {}", Json::String(r.clone())),
107                json!({ "root": r }),
108            ));
109        };
110        if !root_ids.contains(&info.doc_id) {
111            root_ids.push(info.doc_id);
112        }
113    }
114    let degrees = args.degrees.unwrap_or(DEFAULT_DEGREES).max(0);
115    let depth = MAX_DEPTH.min(degrees + 1);
116    let effective_degrees = depth - 1;
117    let direction = args.direction.clone().unwrap_or_else(|| "both".to_owned());
118    let max_documents =
119        usize::try_from(args.max_documents.unwrap_or(DEFAULT_MAX_DOCUMENTS).max(1)).unwrap_or(1);
120    let dirs: Vec<&str> = match direction.as_str() {
121        "both" => vec!["out", "in"],
122        "in" => vec!["in"],
123        _ => vec!["out"],
124    };
125    let seed = root_ids
126        .iter()
127        .map(|id| format!("$id == {}", Json::String(id.clone())))
128        .collect::<Vec<_>>()
129        .join(" || ");
130    let (user_select, out_names) = build_user_select(&args.select);
131
132    let mut queries = Vec::new();
133    let mut docs: Vec<Doc> = Vec::new();
134    let mut edges: Vec<Json> = Vec::new();
135    let mut query_truncated = false;
136    for dir in &dirs {
137        let q = build_query(&seed, dir, depth, &user_select);
138        queries.push(q.clone());
139        let res = query(
140            store,
141            repo_id,
142            &q,
143            QueryOptions {
144                limit: Some(max_documents + 1),
145                cursor: None,
146                provider,
147                in_memory: false,
148            },
149        )?;
150        if res.truncated {
151            query_truncated = true;
152        }
153        for hit in &res.hits {
154            let id = hit["id"].as_str().unwrap_or_default().to_owned();
155            let path = hit["path"].as_str().unwrap_or_default().to_owned();
156            let hop = hit["_depth"].as_f64().unwrap_or(f64::NAN);
157            let degree = (hop - 1.0) as i64;
158            let replace = docs
159                .iter()
160                .position(|d| d.id == id)
161                .map(|i| (i, degree < docs[i].degree));
162            match replace {
163                Some((_, false)) => {}
164                found => {
165                    let extra: Vec<(String, Json)> = out_names
166                        .iter()
167                        .enumerate()
168                        .filter_map(|(i, n)| {
169                            n.as_ref().map(|name| {
170                                (
171                                    name.clone(),
172                                    hit.get(format!("_u{i}")).cloned().unwrap_or(Json::Null),
173                                )
174                            })
175                        })
176                        .collect();
177                    let doc = Doc {
178                        id: id.clone(),
179                        path,
180                        degree,
181                        extra,
182                    };
183                    match found {
184                        Some((i, true)) => docs[i] = doc,
185                        _ => docs.push(doc),
186                    }
187                }
188            }
189            if let Some(raw) = hit["_edges"].as_array() {
190                for e in raw {
191                    let eid = e["id"].as_str().unwrap_or_default();
192                    if !edges.iter().any(|x| x["id"].as_str() == Some(eid)) {
193                        edges.push(e.clone());
194                    }
195                }
196            }
197        }
198    }
199
200    if let Some(p) = &args.predicate {
201        let mut adj: HashMap<String, Vec<String>> = HashMap::new();
202        for e in &edges {
203            if e["predicate"].as_str() != Some(p.as_str()) {
204                continue;
205            }
206            let src = e["src"].as_str().unwrap_or_default().to_owned();
207            let dst = e["dst"].as_str().unwrap_or_default().to_owned();
208            if dirs.contains(&"out") {
209                adj.entry(src.clone()).or_default().push(dst.clone());
210            }
211            if dirs.contains(&"in") {
212                adj.entry(dst).or_default().push(src);
213            }
214        }
215        let mut depth_of: Vec<(String, i64)> = root_ids.iter().map(|id| (id.clone(), 0)).collect();
216        let mut wave: Vec<String> = root_ids.clone();
217        let mut lvl = 1;
218        while lvl <= effective_degrees && !wave.is_empty() {
219            let mut next = Vec::new();
220            for from in &wave {
221                for to in adj.get(from).map_or(&[][..], Vec::as_slice) {
222                    if docs.iter().any(|d| &d.id == to) && !depth_of.iter().any(|(id, _)| id == to)
223                    {
224                        depth_of.push((to.clone(), lvl));
225                        next.push(to.clone());
226                    }
227                }
228            }
229            wave = next;
230            lvl += 1;
231        }
232        let mut restricted: Vec<Doc> = Vec::new();
233        for (id, d) in depth_of {
234            if let Some(orig) = docs.iter().find(|x| x.id == id) {
235                restricted.push(Doc {
236                    id: orig.id.clone(),
237                    path: orig.path.clone(),
238                    degree: d,
239                    extra: orig.extra.clone(),
240                });
241            }
242        }
243        docs = restricted;
244    }
245
246    docs.sort_by(|a, b| {
247        a.degree
248            .cmp(&b.degree)
249            .then_with(|| locale_compare(&a.path, &b.path))
250    });
251    let capped = docs.len() > max_documents;
252    docs.truncate(max_documents);
253    let truncated = query_truncated || capped;
254
255    let mut frontier = Vec::new();
256    let documents: Vec<Json> = docs
257        .iter()
258        .map(|d| {
259            let is_frontier = d.degree == effective_degrees;
260            if is_frontier {
261                frontier.push(json!({ "id": d.id, "path": d.path, "degree": d.degree }));
262            }
263            let mut m = Map::new();
264            m.insert("id".to_owned(), json!(d.id));
265            m.insert("path".to_owned(), json!(d.path));
266            m.insert("degree".to_owned(), json!(d.degree));
267            m.insert("frontier".to_owned(), json!(is_frontier));
268            for (k, v) in &d.extra {
269                m.insert(k.clone(), v.clone());
270            }
271            Json::Object(m)
272        })
273        .collect();
274
275    let reached = |id: &str| docs.iter().any(|d| d.id == id);
276    let mut kept: Vec<Json> = edges
277        .into_iter()
278        .filter(|e| {
279            if let Some(p) = &args.predicate {
280                if e["predicate"].as_str() != Some(p.as_str()) {
281                    return false;
282                }
283            }
284            let src_in = reached(e["src"].as_str().unwrap_or_default());
285            let dst_kind = e["dst_kind"].as_str().unwrap_or_default();
286            let dst_dangling =
287                dst_kind == "external" || (dst_kind == "document" && e["dst_path"].is_null());
288            let dst_in = reached(e["dst"].as_str().unwrap_or_default()) || dst_dangling;
289            src_in && dst_in
290        })
291        .collect();
292    kept.sort_by(|a, b| {
293        let (sa, sb) = (
294            a["src"].as_str().unwrap_or_default(),
295            b["src"].as_str().unwrap_or_default(),
296        );
297        if sa == sb {
298            locale_compare(
299                a["id"].as_str().unwrap_or_default(),
300                b["id"].as_str().unwrap_or_default(),
301            )
302        } else {
303            locale_compare(sa, sb)
304        }
305    });
306
307    Ok(json!({
308        "roots": root_ids,
309        "degrees": effective_degrees,
310        "direction": direction,
311        "documents": documents,
312        "edges": kept,
313        "frontier": frontier,
314        "truncated": truncated,
315        "queries": queries,
316    }))
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[test]
324    fn user_select_aliases() {
325        let (clause, names) = build_user_select(&[
326            "layer".into(),
327            "$path".into(),
328            "a + 1".into(),
329            "layer".into(),
330        ]);
331        assert_eq!(clause, ", _u0: layer, _u1: $path, _u2: a + 1, _u3: layer");
332        assert_eq!(
333            names,
334            [
335                Some("layer".into()),
336                Some("path".into()),
337                Some("sel_2".into()),
338                Some("layer_3".into())
339            ]
340        );
341        assert_eq!(build_user_select(&[]).0, "");
342    }
343
344    #[test]
345    fn query_shape() {
346        let q = build_query("$id == \"d_0\"", "out", 2, "");
347        assert!(
348            q.starts_with("select _depth: $depth, _stop: $stop, _edges: doc.out_edges collect {")
349        );
350        assert!(q.ends_with("from docs where $id == \"d_0\" follow distinct doc.out { depth 2 }"));
351        assert!(oqx::parse_string(&q).is_ok());
352    }
353
354    #[test]
355    fn order_is_bytewise() {
356        assert_eq!(locale_compare("B", "a"), Ordering::Less);
357        assert_eq!(locale_compare("a", "b"), Ordering::Less);
358    }
359}
omgbase_surface/graph.rs

omgbase_surface/
graph.rs
