Skip to main content

partiql_ast/pretty/
graph.rs

1use crate::ast::*;
2use crate::pretty::pretty_parenthesized_expr;
3use partiql_common::pretty::{
4    pretty_doc_list, pretty_list, pretty_parenthesized_doc, pretty_surrounded,
5    pretty_surrounded_doc, PrettyDoc, PRETTY_INDENT_MINOR_NEST,
6};
7use pretty::{DocAllocator, DocBuilder};
8use std::borrow::Cow;
9
10impl PrettyDoc for GraphMatch {
11    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
12    where
13        D: DocAllocator<'b, A>,
14        D::Doc: Clone,
15        A: Clone,
16    {
17        let GraphMatch {
18            expr,
19            pattern,
20            shape,
21        } = self;
22        let head = arena.intersperse([expr.pretty_doc(arena), arena.text("MATCH")], arena.space());
23        let patterns = pattern.pretty_doc(arena).group();
24        let mut match_expr = arena.intersperse([head, patterns], arena.space());
25        let shapes = [
26            shape.rows.as_ref().map(|d| d.pretty_doc(arena)),
27            shape.cols.as_ref().map(|d| d.pretty_doc(arena)),
28            shape.export.as_ref().map(|d| d.pretty_doc(arena)),
29        ]
30        .into_iter()
31        .flatten()
32        .collect::<Vec<_>>();
33        if !shapes.is_empty() {
34            let docs = std::iter::once(match_expr).chain(shapes);
35            match_expr = arena.intersperse(docs, arena.hardline());
36        }
37
38        pretty_parenthesized_doc(match_expr, arena)
39    }
40}
41
42impl PrettyDoc for GraphTableRows {
43    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
44    where
45        D: DocAllocator<'b, A>,
46        D::Doc: Clone,
47        A: Clone,
48    {
49        match self {
50            GraphTableRows::OneRowPerMatch => arena.text("ONE ROW PER MATCH"),
51            GraphTableRows::OneRowPerVertex { v, in_paths } => {
52                let prefix = arena.text("ONE ROW PER NODE");
53                let spec = pretty_parenthesized_doc(arena.text(&v.value), arena);
54                let in_paths = in_paths.as_ref().map(|paths| {
55                    let paths = pretty_parenthesized_doc(
56                        pretty_doc_list(paths.iter().map(|p| arena.text(&p.value)), 0, arena),
57                        arena,
58                    );
59                    [arena.text("IN"), paths]
60                });
61                arena.intersperse(
62                    [Some([prefix, spec]), in_paths]
63                        .into_iter()
64                        .flatten()
65                        .flatten(),
66                    arena.softline(),
67                )
68            }
69            GraphTableRows::OneRowPerStep {
70                v1,
71                e,
72                v2,
73                in_paths,
74            } => {
75                let prefix = arena.text("ONE ROW PER STEP");
76                let step = pretty_doc_list(
77                    [v1, e, v2].into_iter().map(|n| arena.text(&n.value)),
78                    0,
79                    arena,
80                );
81                let spec = pretty_parenthesized_doc(step, arena);
82                let in_paths = in_paths.as_ref().map(|paths| {
83                    let paths =
84                        pretty_doc_list(paths.iter().map(|p| arena.text(&p.value)), 0, arena);
85                    [arena.text("IN"), pretty_parenthesized_doc(paths, arena)]
86                });
87                arena.intersperse(
88                    [Some([prefix, spec]), in_paths]
89                        .into_iter()
90                        .flatten()
91                        .flatten(),
92                    arena.softline(),
93                )
94            }
95        }
96    }
97}
98
99impl PrettyDoc for GraphTableColumns {
100    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
101    where
102        D: DocAllocator<'b, A>,
103        D::Doc: Clone,
104        A: Clone,
105    {
106        let col_defs = pretty_list(&self.columns, 0, arena);
107        arena.intersperse(
108            [
109                arena.text("COLUMNS"),
110                pretty_parenthesized_doc(col_defs, arena),
111            ],
112            arena.space(),
113        )
114    }
115}
116
117impl PrettyDoc for GraphTableColumnDef {
118    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
119    where
120        D: DocAllocator<'b, A>,
121        D::Doc: Clone,
122        A: Clone,
123    {
124        match self {
125            GraphTableColumnDef::Expr(expr, as_ident) => {
126                let parts = if let Some(as_ident) = as_ident {
127                    vec![
128                        expr.pretty_doc(arena),
129                        arena.text("AS"),
130                        arena.text(&as_ident.value),
131                    ]
132                } else {
133                    vec![expr.pretty_doc(arena)]
134                };
135                arena.intersperse(parts, arena.space())
136            }
137            GraphTableColumnDef::AllProperties(_) => {
138                unreachable!()
139            }
140        }
141    }
142}
143
144impl PrettyDoc for GraphTableExport {
145    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
146    where
147        D: DocAllocator<'b, A>,
148        D::Doc: Clone,
149        A: Clone,
150    {
151        match self {
152            GraphTableExport::AllSingletons { except } => {
153                let prefix = arena.text("EXPORT ALL SINGLETONS");
154                if let Some(except) = except {
155                    let except =
156                        pretty_doc_list(except.iter().map(|s| arena.text(&s.value)), 0, arena);
157                    let parts = [
158                        prefix,
159                        arena.text("EXCEPT"),
160                        pretty_parenthesized_doc(except, arena),
161                    ];
162                    arena.intersperse(parts, arena.space())
163                } else {
164                    prefix
165                }
166            }
167            GraphTableExport::Singletons { exports } => {
168                let prefix = arena.text("EXPORT SINGLETONS");
169                let exports =
170                    pretty_doc_list(exports.iter().map(|e| arena.text(&e.value)), 0, arena);
171                let parts = [prefix, pretty_parenthesized_doc(exports, arena)];
172                arena.intersperse(parts, arena.space())
173            }
174            GraphTableExport::NoSingletons => arena.text("EXPORT NO SINGLETONS"),
175        }
176    }
177}
178
179impl PrettyDoc for GraphPattern {
180    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
181    where
182        D: DocAllocator<'b, A>,
183        D::Doc: Clone,
184        A: Clone,
185    {
186        let GraphPattern {
187            match_mode,
188            patterns,
189            keep,
190            where_clause,
191        } = &self;
192        let patterns = pretty_list(patterns, PRETTY_INDENT_MINOR_NEST, arena).group();
193        let parts = [
194            match_mode.as_ref().map(|inner| inner.pretty_doc(arena)),
195            Some(patterns),
196            keep.as_ref().map(|keep| {
197                arena.intersperse([arena.text("KEEP"), keep.pretty_doc(arena)], arena.space())
198            }),
199            where_clause.as_ref().map(|clause| {
200                arena.intersperse(
201                    [arena.text("WHERE"), clause.pretty_doc(arena)],
202                    arena.space(),
203                )
204            }),
205        ]
206        .into_iter()
207        .flatten();
208
209        arena.intersperse(parts, arena.space())
210    }
211}
212
213impl PrettyDoc for GraphMatchMode {
214    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
215    where
216        D: DocAllocator<'b, A>,
217        D::Doc: Clone,
218        A: Clone,
219    {
220        match self {
221            GraphMatchMode::DifferentEdges => arena.text("DIFFERENT EDGES"),
222            GraphMatchMode::RepeatableElements => arena.text("REPEATABLE ELEMENTS"),
223        }
224    }
225}
226
227impl PrettyDoc for GraphPathPrefix {
228    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
229    where
230        D: DocAllocator<'b, A>,
231        D::Doc: Clone,
232        A: Clone,
233    {
234        match self {
235            GraphPathPrefix::Mode(mode) => mode.pretty_doc(arena),
236            GraphPathPrefix::Search(search, mode) => {
237                let mode = mode.as_ref().map(|mode| mode.pretty_doc(arena));
238                let parts = match search {
239                    GraphPathSearchPrefix::All => vec![Some(arena.text("ALL")), mode],
240                    GraphPathSearchPrefix::Any => vec![Some(arena.text("ANY")), mode],
241                    GraphPathSearchPrefix::AnyK(k) => vec![
242                        Some(arena.text("ANY")),
243                        Some(arena.text(k.to_string())),
244                        mode,
245                    ],
246                    GraphPathSearchPrefix::AllShortest => {
247                        vec![Some(arena.text("ALL")), Some(arena.text("SHORTEST")), mode]
248                    }
249                    GraphPathSearchPrefix::AnyShortest => {
250                        vec![Some(arena.text("ANY")), Some(arena.text("SHORTEST")), mode]
251                    }
252                    GraphPathSearchPrefix::ShortestK(k) => vec![
253                        Some(arena.text("SHORTEST")),
254                        Some(arena.text(k.to_string())),
255                        mode,
256                    ],
257                    GraphPathSearchPrefix::ShortestKGroup(k) => vec![
258                        Some(arena.text("SHORTEST")),
259                        k.as_ref().map(|k| arena.text(k.to_string())),
260                        mode,
261                        Some(arena.text("GROUPS")),
262                    ],
263                }
264                .into_iter()
265                .flatten()
266                .collect::<Vec<_>>();
267                arena.intersperse(parts, arena.space())
268            }
269        }
270    }
271}
272
273impl PrettyDoc for GraphPathMode {
274    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
275    where
276        D: DocAllocator<'b, A>,
277        D::Doc: Clone,
278        A: Clone,
279    {
280        let txt = match self {
281            GraphPathMode::Walk => "WALK",
282            GraphPathMode::Trail => "TRAIL",
283            GraphPathMode::Acyclic => "ACYCLIC",
284            GraphPathMode::Simple => "SIMPLE",
285        };
286        arena.text(txt)
287    }
288}
289
290impl PrettyDoc for GraphPathPattern {
291    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
292    where
293        D: DocAllocator<'b, A>,
294        D::Doc: Clone,
295        A: Clone,
296    {
297        let var = self
298            .variable
299            .as_ref()
300            .map(|var| arena.intersperse([arena.text(&var.value), arena.text("=")], arena.space()));
301
302        let prefix = self.prefix.as_ref().map(|prefix| prefix.pretty_doc(arena));
303
304        let parts = [var, prefix, Some(self.path.pretty_doc(arena))]
305            .into_iter()
306            .flatten();
307        arena.intersperse(parts, arena.space())
308    }
309}
310
311impl PrettyDoc for GraphPathSubPattern {
312    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
313    where
314        D: DocAllocator<'b, A>,
315        D::Doc: Clone,
316        A: Clone,
317    {
318        let var = self
319            .variable
320            .as_ref()
321            .map(|var| arena.intersperse([arena.text(&var.value), arena.text("=")], arena.space()));
322
323        let mode = self.mode.as_ref().map(|prefix| prefix.pretty_doc(arena));
324        let where_clause = self.where_clause.as_ref().map(|clause| {
325            arena.intersperse(
326                [arena.text("WHERE"), clause.pretty_doc(arena)],
327                arena.space(),
328            )
329        });
330
331        let parts = [var, mode, Some(self.path.pretty_doc(arena)), where_clause]
332            .into_iter()
333            .flatten();
334        arena.intersperse(parts, arena.space())
335    }
336}
337
338impl PrettyDoc for GraphMatchPathPattern {
339    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
340    where
341        D: DocAllocator<'b, A>,
342        D::Doc: Clone,
343        A: Clone,
344    {
345        match self {
346            GraphMatchPathPattern::Path(path) => {
347                arena.intersperse(path.iter().map(|e| e.pretty_doc(arena)), arena.space())
348            }
349            GraphMatchPathPattern::Quantified(GraphMatchPathPatternQuantified { path, quant }) => {
350                arena.concat([path.pretty_doc(arena), quant.pretty_doc(arena)])
351            }
352            GraphMatchPathPattern::Questioned(p) => {
353                arena.concat([p.pretty_doc(arena), arena.text("?")])
354            }
355            GraphMatchPathPattern::Sub(path) => pretty_surrounded(path, "(", ")", arena),
356            GraphMatchPathPattern::Node(node) => node.pretty_doc(arena),
357            GraphMatchPathPattern::Edge(edge) => edge.pretty_doc(arena),
358            GraphMatchPathPattern::Union(u) => {
359                arena.intersperse(u.iter().map(|l| l.pretty_doc(arena)), arena.text(" | "))
360            }
361            GraphMatchPathPattern::Multiset(s) => {
362                arena.intersperse(s.iter().map(|l| l.pretty_doc(arena)), arena.text(" |+| "))
363            }
364            GraphMatchPathPattern::Simplified(simplified) => simplified.pretty_doc(arena),
365        }
366    }
367}
368
369impl PrettyDoc for GraphMatchSimplified {
370    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
371    where
372        D: DocAllocator<'b, A>,
373        D::Doc: Clone,
374        A: Clone,
375    {
376        let ends = match self.dir {
377            GraphMatchDirection::Left => ["<-/", "/-"],
378            GraphMatchDirection::Undirected => ["~/", "/~"],
379            GraphMatchDirection::Right => ["-/", "/->"],
380            GraphMatchDirection::LeftOrUndirected => ["<~/", "/~"],
381            GraphMatchDirection::UndirectedOrRight => ["~/", "/~>"],
382            GraphMatchDirection::LeftOrRight => ["<-/", "/->"],
383            GraphMatchDirection::LeftOrUndirectedOrRight => ["-/", "/-"],
384        };
385
386        let parts = [
387            arena.text(ends[0]),
388            self.pattern.pretty_doc(arena),
389            arena.text(ends[1]),
390        ];
391        arena.intersperse(parts, arena.space()).group()
392    }
393}
394
395impl PrettyDoc for GraphMatchSimplifiedPattern {
396    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
397    where
398        D: DocAllocator<'b, A>,
399        D::Doc: Clone,
400        A: Clone,
401    {
402        match self {
403            GraphMatchSimplifiedPattern::Union(u) => {
404                arena.intersperse(u.iter().map(|l| l.pretty_doc(arena)), arena.text(" | "))
405            }
406            GraphMatchSimplifiedPattern::Multiset(s) => {
407                arena.intersperse(s.iter().map(|l| l.pretty_doc(arena)), arena.text(" |+| "))
408            }
409            GraphMatchSimplifiedPattern::Path(path) => {
410                arena.intersperse(path.iter().map(|e| e.pretty_doc(arena)), arena.space())
411            }
412            GraphMatchSimplifiedPattern::Sub(path) => pretty_surrounded(path, "(", ")", arena),
413            GraphMatchSimplifiedPattern::Conjunction(c) => {
414                arena.intersperse(c.iter().map(|l| l.pretty_doc(arena)), arena.text("&"))
415            }
416            GraphMatchSimplifiedPattern::Questioned(p) => {
417                arena.concat([p.pretty_doc(arena), arena.text("?")])
418            }
419            GraphMatchSimplifiedPattern::Quantified(GraphMatchSimplifiedPatternQuantified {
420                path,
421                quant,
422            }) => arena.concat([path.pretty_doc(arena), quant.pretty_doc(arena)]),
423            GraphMatchSimplifiedPattern::Direction(GraphMatchSimplifiedPatternDirected {
424                dir,
425                path,
426            }) => {
427                let path = path.pretty_doc(arena);
428                let parts = match dir {
429                    GraphMatchDirection::Left => vec![arena.text("<"), path],
430                    GraphMatchDirection::Undirected => vec![arena.text("~"), path],
431                    GraphMatchDirection::Right => vec![path, arena.text(">")],
432                    GraphMatchDirection::LeftOrUndirected => vec![arena.text("<~"), path],
433                    GraphMatchDirection::UndirectedOrRight => {
434                        vec![arena.text("~"), path, arena.text(">")]
435                    }
436                    GraphMatchDirection::LeftOrRight => {
437                        vec![arena.text("<"), path, arena.text(">")]
438                    }
439                    GraphMatchDirection::LeftOrUndirectedOrRight => vec![arena.text("-"), path],
440                };
441                arena.concat(parts).group()
442            }
443            GraphMatchSimplifiedPattern::Negated(l) => {
444                arena.concat([arena.text("!"), l.pretty_doc(arena)])
445            }
446            GraphMatchSimplifiedPattern::Label(l) => arena.text(&l.value),
447        }
448    }
449}
450
451impl PrettyDoc for GraphMatchNode {
452    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
453    where
454        D: DocAllocator<'b, A>,
455        D::Doc: Clone,
456        A: Clone,
457    {
458        let name: Vec<_> = [
459            self.variable.as_ref().map(|var| arena.text(&var.value)),
460            self.label
461                .as_ref()
462                .map(|label| arena.concat([arena.text(":"), label.pretty_doc(arena)])),
463        ]
464        .into_iter()
465        .flatten()
466        .collect();
467        let name = if name.is_empty() {
468            None
469        } else {
470            Some(arena.concat(name))
471        };
472        let where_clause = self.where_clause.as_ref().map(|clause| {
473            arena.intersperse(
474                [arena.text("WHERE"), clause.pretty_doc(arena)],
475                arena.space(),
476            )
477        });
478        let parts = [name, where_clause].into_iter().flatten();
479
480        let spec = arena.intersperse(parts, arena.space());
481        pretty_surrounded_doc(spec, "(", ")", arena).group()
482    }
483}
484
485impl PrettyDoc for GraphMatchEdge {
486    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
487    where
488        D: DocAllocator<'b, A>,
489        D::Doc: Clone,
490        A: Clone,
491    {
492        let name: Vec<_> = [
493            self.variable.as_ref().map(|var| arena.text(&var.value)),
494            self.label
495                .as_ref()
496                .map(|label| arena.concat([arena.text(":"), label.pretty_doc(arena)])),
497        ]
498        .into_iter()
499        .flatten()
500        .collect();
501        let name = if name.is_empty() {
502            None
503        } else {
504            Some(arena.concat(name))
505        };
506        let where_clause = self.where_clause.as_ref().map(|clause| {
507            arena.intersperse(
508                [arena.text("WHERE"), clause.pretty_doc(arena)],
509                arena.space(),
510            )
511        });
512        let parts = [name, where_clause]
513            .into_iter()
514            .flatten()
515            .collect::<Vec<_>>();
516
517        let edge = if !parts.is_empty() {
518            let (prefix, suffix) = match self.direction {
519                GraphMatchDirection::Right => ("-[", "]->"),
520                GraphMatchDirection::Left => ("<-[", "]-"),
521                GraphMatchDirection::Undirected => ("~[", "]~"),
522                GraphMatchDirection::UndirectedOrRight => ("~[", "]~>"),
523                GraphMatchDirection::LeftOrUndirected => ("<~[", "]~"),
524                GraphMatchDirection::LeftOrRight => ("<-[", "]->"),
525                GraphMatchDirection::LeftOrUndirectedOrRight => ("-[", "]-"),
526            };
527            let spec = arena.intersperse(parts, arena.space());
528            pretty_surrounded_doc(spec, prefix, suffix, arena)
529        } else {
530            let edge = match self.direction {
531                GraphMatchDirection::Right => "->",
532                GraphMatchDirection::Left => "<-",
533                GraphMatchDirection::Undirected => "~",
534                GraphMatchDirection::UndirectedOrRight => "~>",
535                GraphMatchDirection::LeftOrUndirected => "<~",
536                GraphMatchDirection::LeftOrRight => "<->",
537                GraphMatchDirection::LeftOrUndirectedOrRight => "-",
538            };
539            arena.text(edge)
540        };
541        edge.group()
542    }
543}
544
545impl PrettyDoc for GraphMatchLabel {
546    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
547    where
548        D: DocAllocator<'b, A>,
549        D::Doc: Clone,
550        A: Clone,
551    {
552        match self {
553            GraphMatchLabel::Name(name) => arena.text(&name.value),
554            GraphMatchLabel::Wildcard => arena.text("%"),
555            GraphMatchLabel::Negated(l) => {
556                arena.concat([arena.text("!"), pretty_parenthesized_expr(l, 0, arena)])
557            }
558            GraphMatchLabel::Conjunction(c) => pretty_parenthesized_doc(
559                arena.intersperse(c.iter().map(|l| l.pretty_doc(arena)), arena.text("&")),
560                arena,
561            ),
562            GraphMatchLabel::Disjunction(d) => pretty_parenthesized_doc(
563                arena.intersperse(d.iter().map(|l| l.pretty_doc(arena)), arena.text("|")),
564                arena,
565            ),
566        }
567    }
568}
569
570impl PrettyDoc for GraphMatchQuantifier {
571    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
572    where
573        D: DocAllocator<'b, A>,
574        D::Doc: Clone,
575        A: Clone,
576    {
577        let GraphMatchQuantifier { lower, upper } = &self;
578        match (lower, upper) {
579            (0, None) => arena.text("*"),
580            (1, None) => arena.text("+"),
581            (l, u) => {
582                let l = Cow::Owned(l.to_string());
583                let u = u.map(|u| Cow::Owned(u.to_string())).unwrap_or("".into());
584                pretty_surrounded_doc(arena.concat([l, ",".into(), u]), "{", "}", arena)
585            }
586        }
587    }
588}