partiql_ast/pretty/
mod.rs

1use crate::ast::*;
2use partiql_common::pretty::{
3    pretty_list, pretty_parenthesized_doc, pretty_prefixed_doc, pretty_seperated,
4    pretty_seperated_doc, pretty_seq, pretty_seq_doc, PrettyDoc, PRETTY_INDENT_MINOR_NEST,
5    PRETTY_INDENT_SUBORDINATE_CLAUSE_NEST,
6};
7use pretty::{DocAllocator, DocBuilder};
8mod graph;
9
10impl<T> PrettyDoc for AstNode<T>
11where
12    T: PrettyDoc,
13{
14    #[inline]
15    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
16    where
17        D: DocAllocator<'b, A>,
18        D::Doc: Clone,
19        A: Clone,
20    {
21        self.node.pretty_doc(arena)
22    }
23}
24
25impl PrettyDoc for TopLevelQuery {
26    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
27    where
28        D: DocAllocator<'b, A>,
29        D::Doc: Clone,
30        A: Clone,
31    {
32        if self.with.is_some() {
33            todo!("WITH Clause")
34        }
35        self.query.pretty_doc(arena)
36    }
37}
38
39impl PrettyDoc for Query {
40    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
41    where
42        D: DocAllocator<'b, A>,
43        D::Doc: Clone,
44        A: Clone,
45    {
46        let Query {
47            set,
48            order_by,
49            limit_offset,
50        } = self;
51
52        let clauses = [
53            Some(set.pretty_doc(arena)),
54            order_by.as_ref().map(|inner| inner.pretty_doc(arena)),
55            limit_offset.as_ref().map(|inner| inner.pretty_doc(arena)),
56        ]
57        .into_iter()
58        .flatten();
59
60        arena.intersperse(clauses, arena.softline()).group()
61    }
62}
63
64impl PrettyDoc for QuerySet {
65    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
66    where
67        D: DocAllocator<'b, A>,
68        D::Doc: Clone,
69        A: Clone,
70    {
71        match self {
72            QuerySet::BagOp(op) => op.pretty_doc(arena),
73            QuerySet::Select(sel) => sel.pretty_doc(arena),
74            QuerySet::Expr(e) => e.pretty_doc(arena),
75            QuerySet::Values(v) => pretty_prefixed_doc("VALUES", pretty_list(v, 0, arena), arena),
76            QuerySet::Table(t) => pretty_prefixed_expr("TABLE", t, 0, arena),
77        }
78    }
79}
80
81impl PrettyDoc for BagOpExpr {
82    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
83    where
84        D: DocAllocator<'b, A>,
85        D::Doc: Clone,
86        A: Clone,
87    {
88        let op = match self.bag_op {
89            BagOperator::Union => "UNION",
90            BagOperator::Except => "EXCEPT",
91            BagOperator::Intersect => "INTERSECT",
92            BagOperator::OuterUnion => "OUTER UNION",
93            BagOperator::OuterExcept => "OUTER EXCEPT",
94            BagOperator::OuterIntersect => "OUTER INTERSECT",
95        };
96        let op = arena.text(op);
97        let op = match self.setq {
98            None => op,
99            Some(SetQuantifier::All) => op.append(" ALL"),
100            Some(SetQuantifier::Distinct) => op.append(" DISTINCT"),
101        };
102
103        let lhs = pretty_parenthesized_expr(&self.lhs, PRETTY_INDENT_MINOR_NEST, arena);
104        let rhs = pretty_parenthesized_expr(&self.rhs, PRETTY_INDENT_MINOR_NEST, arena);
105
106        arena.intersperse([lhs, op, rhs], arena.hardline()).group()
107    }
108}
109
110impl PrettyDoc for QueryTable {
111    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
112    where
113        D: DocAllocator<'b, A>,
114        D::Doc: Clone,
115        A: Clone,
116    {
117        self.table_name.pretty_doc(arena)
118    }
119}
120
121impl PrettyDoc for Select {
122    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
123    where
124        D: DocAllocator<'b, A>,
125        D::Doc: Clone,
126        A: Clone,
127    {
128        fn format<'b, C, D, A>(child: &'b C, arena: &'b D) -> DocBuilder<'b, D, A>
129        where
130            D: DocAllocator<'b, A>,
131            D::Doc: Clone,
132            A: Clone,
133            C: PrettyDoc,
134        {
135            child.pretty_doc(arena).group()
136        }
137
138        fn delegate<'b, C, D, A>(child: &'b Option<C>, arena: &'b D) -> Option<DocBuilder<'b, D, A>>
139        where
140            D: DocAllocator<'b, A>,
141            D::Doc: Clone,
142            A: Clone,
143            C: PrettyDoc,
144        {
145            child.as_ref().map(|inner| format(inner, arena))
146        }
147
148        let Select {
149            project,
150            exclude,
151            from,
152            from_let,
153            where_clause,
154            group_by,
155            having,
156        } = self;
157        let mut clauses = [
158            Some(format(project, arena)),
159            delegate(exclude, arena),
160            delegate(from, arena),
161            delegate(from_let, arena),
162            delegate(where_clause, arena),
163            delegate(group_by, arena),
164            delegate(having, arena),
165        ]
166        .into_iter()
167        .flatten();
168
169        let mut result = arena.nil();
170        let separator = arena.line();
171        if let Some(first) = clauses.next() {
172            let mut curr = first;
173
174            for clause in clauses {
175                result = result.append(curr.append(separator.clone()).group());
176                curr = clause;
177            }
178
179            result = result.append(curr);
180        }
181
182        result
183    }
184}
185
186impl PrettyDoc for Projection {
187    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
188    where
189        D: DocAllocator<'b, A>,
190        D::Doc: Clone,
191        A: Clone,
192    {
193        if self.setq.is_some() {
194            todo!("project SetQuantifier")
195        }
196        self.kind.pretty_doc(arena)
197    }
198}
199
200impl PrettyDoc for ProjectionKind {
201    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
202    where
203        D: DocAllocator<'b, A>,
204        D::Doc: Clone,
205        A: Clone,
206    {
207        match self {
208            ProjectionKind::ProjectStar => arena.text("SELECT *"),
209            ProjectionKind::ProjectList(l) => pretty_prefixed_doc(
210                "SELECT",
211                pretty_list(l, PRETTY_INDENT_MINOR_NEST, arena),
212                arena,
213            ),
214            ProjectionKind::ProjectPivot(ProjectPivot { key, value }) => {
215                let parts = [
216                    value.pretty_doc(arena),
217                    arena.text("AT"),
218                    key.pretty_doc(arena),
219                ];
220                let decl = arena.intersperse(parts, arena.space()).group();
221                pretty_prefixed_doc("PIVOT", decl, arena)
222            }
223            ProjectionKind::ProjectValue(ctor) => {
224                pretty_prefixed_expr("SELECT VALUE", ctor, PRETTY_INDENT_MINOR_NEST, arena)
225            }
226        }
227        .group()
228    }
229}
230
231impl PrettyDoc for ProjectItem {
232    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
233    where
234        D: DocAllocator<'b, A>,
235        D::Doc: Clone,
236        A: Clone,
237    {
238        match self {
239            ProjectItem::ProjectAll(_) => {
240                todo!("ProjectItem::ProjectAll; remove this?")
241            }
242            ProjectItem::ProjectExpr(e) => e.pretty_doc(arena),
243        }
244    }
245}
246
247impl PrettyDoc for ProjectExpr {
248    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
249    where
250        D: DocAllocator<'b, A>,
251        D::Doc: Clone,
252        A: Clone,
253    {
254        pretty_source_as_alias(&self.expr, self.as_alias.as_ref(), arena)
255            .unwrap_or_else(|| self.expr.pretty_doc(arena))
256    }
257}
258
259impl PrettyDoc for Exclusion {
260    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
261    where
262        D: DocAllocator<'b, A>,
263        D::Doc: Clone,
264        A: Clone,
265    {
266        pretty_prefixed_doc(
267            "EXCLUDE",
268            pretty_list(&self.items, PRETTY_INDENT_MINOR_NEST, arena),
269            arena,
270        )
271    }
272}
273
274impl PrettyDoc for ExcludePath {
275    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
276    where
277        D: DocAllocator<'b, A>,
278        D::Doc: Clone,
279        A: Clone,
280    {
281        let ExcludePath { root, steps } = self;
282        let mut path = root.pretty_doc(arena);
283        for step in steps {
284            path = path.append(match step {
285                ExcludePathStep::PathProject(e) => arena.text(".").append(e.pretty_doc(arena)),
286                ExcludePathStep::PathIndex(e) => arena
287                    .text("[")
288                    .append(e.pretty_doc(arena))
289                    .append(arena.text("]")),
290                ExcludePathStep::PathForEach => arena.text("[*]"),
291                ExcludePathStep::PathUnpivot => arena.text(".*"),
292            });
293        }
294
295        path
296    }
297}
298
299impl PrettyDoc for Expr {
300    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
301    where
302        D: DocAllocator<'b, A>,
303        D::Doc: Clone,
304        A: Clone,
305    {
306        match self {
307            Expr::Lit(inner) => inner.pretty_doc(arena),
308            Expr::VarRef(inner) => inner.pretty_doc(arena),
309            Expr::BinOp(inner) => inner.pretty_doc(arena),
310            Expr::UniOp(inner) => inner.pretty_doc(arena),
311            Expr::Like(inner) => inner.pretty_doc(arena),
312            Expr::Between(inner) => inner.pretty_doc(arena),
313            Expr::In(inner) => inner.pretty_doc(arena),
314            Expr::Case(inner) => inner.pretty_doc(arena),
315            Expr::Struct(inner) => inner.pretty_doc(arena),
316            Expr::Bag(inner) => inner.pretty_doc(arena),
317            Expr::List(inner) => inner.pretty_doc(arena),
318            Expr::Path(inner) => inner.pretty_doc(arena),
319            Expr::Call(inner) => inner.pretty_doc(arena),
320            Expr::CallAgg(inner) => inner.pretty_doc(arena),
321            Expr::GraphMatch(inner) => inner.pretty_doc(arena),
322            Expr::Query(inner) => {
323                let inner = inner.pretty_doc(arena).group();
324                arena
325                    .text("(")
326                    .append(inner.nest(PRETTY_INDENT_SUBORDINATE_CLAUSE_NEST))
327                    .append(arena.text(")"))
328            }
329            Expr::Error => {
330                unreachable!();
331            }
332        }
333        .group()
334    }
335}
336
337impl PrettyDoc for Path {
338    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
339    where
340        D: DocAllocator<'b, A>,
341        D::Doc: Clone,
342        A: Clone,
343    {
344        let Path { root, steps } = self;
345        let mut path = root.pretty_doc(arena);
346        for step in steps {
347            path = path.append(match step {
348                PathStep::PathProject(e) => arena.text(".").append(e.index.pretty_doc(arena)),
349                PathStep::PathIndex(e) => arena
350                    .text("[")
351                    .append(e.index.pretty_doc(arena))
352                    .append(arena.text("]")),
353                PathStep::PathForEach => arena.text("[*]"),
354                PathStep::PathUnpivot => arena.text(".*"),
355            });
356        }
357
358        path
359    }
360}
361
362impl PrettyDoc for VarRef {
363    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
364    where
365        D: DocAllocator<'b, A>,
366        D::Doc: Clone,
367        A: Clone,
368    {
369        let name = self.name.pretty_doc(arena);
370        match self.qualifier {
371            ScopeQualifier::Unqualified => name,
372            ScopeQualifier::Qualified => arena.text("@").append(name).group(),
373        }
374    }
375}
376
377impl PrettyDoc for Lit {
378    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
379    where
380        D: DocAllocator<'b, A>,
381        D::Doc: Clone,
382        A: Clone,
383    {
384        match self {
385            Lit::Null => arena.text("NULL"),
386            Lit::Missing => arena.text("MISSING"),
387            Lit::Int8Lit(inner) => arena.text(inner.to_string()),
388            Lit::Int16Lit(inner) => arena.text(inner.to_string()),
389            Lit::Int32Lit(inner) => arena.text(inner.to_string()),
390            Lit::Int64Lit(inner) => arena.text(inner.to_string()),
391            Lit::DecimalLit(inner) => inner.pretty_doc(arena),
392            Lit::NumericLit(inner) => inner.pretty_doc(arena),
393            Lit::RealLit(inner) => arena.text(inner.to_string()),
394            Lit::FloatLit(inner) => arena.text(inner.to_string()),
395            Lit::DoubleLit(inner) => arena.text(inner.to_string()),
396            Lit::BoolLit(inner) => arena.text(inner.to_string()),
397            Lit::EmbeddedDocLit(inner, _typ) => inner.pretty_doc(arena), // TODO better pretty for embedded doc: https://github.com/partiql/partiql-lang-rust/issues/508
398            Lit::CharStringLit(inner) => inner.pretty_doc(arena),
399            Lit::NationalCharStringLit(inner) => inner.pretty_doc(arena),
400            Lit::BitStringLit(inner) => inner.pretty_doc(arena),
401            Lit::HexStringLit(inner) => inner.pretty_doc(arena),
402            Lit::StructLit(inner) => inner.pretty_doc(arena),
403            Lit::BagLit(inner) => inner.pretty_doc(arena),
404            Lit::ListLit(inner) => inner.pretty_doc(arena),
405            Lit::TypedLit(s, ty) => {
406                let ty = ty.pretty_doc(arena);
407                let s = s.pretty_doc(arena);
408                pretty_seperated_doc(arena.space(), [ty, s], 0, arena)
409            }
410        }
411    }
412}
413
414impl PrettyDoc for Type {
415    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
416    where
417        D: DocAllocator<'b, A>,
418        D::Doc: Clone,
419        A: Clone,
420    {
421        match self {
422            Type::CustomType(cty) => cty.pretty_doc(arena),
423            Type::NullType => arena.text("NULL"),
424            Type::BooleanType => arena.text("BOOL"),
425            Type::Integer2Type => arena.text("INT2"),
426            Type::Integer4Type => arena.text("INT4"),
427            Type::Integer8Type => arena.text("INT8"),
428            Type::DecimalType => arena.text("DECIMAL"),
429            Type::NumericType => arena.text("NUMERIC"),
430            Type::RealType => arena.text("REAL"),
431            Type::DoublePrecisionType => arena.text("DOUBLE PRECISION"),
432            Type::TimestampType => arena.text("TIMESTAMP"),
433            Type::CharacterType => arena.text("CHAR"),
434            Type::CharacterVaryingType => arena.text("VARCHAR"),
435            Type::MissingType => arena.text("MISSING"),
436            Type::StringType => arena.text("STRING"),
437            Type::SymbolType => arena.text("SYMBOL"),
438            Type::BlobType => arena.text("BLOB"),
439            Type::ClobType => arena.text("CLOB"),
440            Type::DateType => arena.text("DATE"),
441            Type::TimeType => arena.text("TIME"),
442            Type::ZonedTimestampType => arena.text("TIMESTAMPTZ"),
443            Type::StructType => arena.text("STRUCT"),
444            Type::TupleType => arena.text("TUPLE"),
445            Type::ListType => arena.text("LIST"),
446            Type::BagType => arena.text("BAG"),
447            Type::AnyType => arena.text("ANY"),
448        }
449    }
450}
451
452impl PrettyDoc for CustomType {
453    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
454    where
455        D: DocAllocator<'b, A>,
456        D::Doc: Clone,
457        A: Clone,
458    {
459        pretty_seperated(arena.space(), &self.parts, 0, arena)
460    }
461}
462
463impl PrettyDoc for CustomTypePart {
464    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
465    where
466        D: DocAllocator<'b, A>,
467        D::Doc: Clone,
468        A: Clone,
469    {
470        match self {
471            CustomTypePart::Name(sym) => sym.pretty_doc(arena),
472            CustomTypePart::Parameterized(sym, param) => {
473                let sym = sym.pretty_doc(arena);
474                let list = pretty_list(param, 0, arena);
475                let list = pretty_parenthesized_doc(list, arena);
476                sym.append(list)
477            }
478        }
479    }
480}
481
482impl PrettyDoc for CustomTypeParam {
483    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
484    where
485        D: DocAllocator<'b, A>,
486        D::Doc: Clone,
487        A: Clone,
488    {
489        match self {
490            CustomTypeParam::Lit(l) => l.pretty_doc(arena),
491            CustomTypeParam::Type(ty) => ty.pretty_doc(arena),
492        }
493    }
494}
495
496impl PrettyDoc for BinOp {
497    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
498    where
499        D: DocAllocator<'b, A>,
500        D::Doc: Clone,
501        A: Clone,
502    {
503        let BinOp { kind, lhs, rhs } = self;
504        let (nest, sym) = match kind {
505            BinOpKind::Add => (0, "+"),
506            BinOpKind::Div => (0, "/"),
507            BinOpKind::Exp => (0, "^"),
508            BinOpKind::Mod => (0, "%"),
509            BinOpKind::Mul => (0, "*"),
510            BinOpKind::Sub => (0, "-"),
511            BinOpKind::And => (PRETTY_INDENT_MINOR_NEST, "AND"),
512            BinOpKind::Or => (PRETTY_INDENT_MINOR_NEST, "OR"),
513            BinOpKind::Concat => (0, "||"),
514            BinOpKind::Eq => (0, "="),
515            BinOpKind::Gt => (0, ">"),
516            BinOpKind::Gte => (0, ">="),
517            BinOpKind::Lt => (0, "<"),
518            BinOpKind::Lte => (0, "<="),
519            BinOpKind::Ne => (0, "<>"),
520            BinOpKind::Is => (0, "IS"),
521        };
522        let op = arena.text(sym);
523        let lhs = lhs.pretty_doc(arena).nest(nest);
524        let rhs = rhs.pretty_doc(arena).nest(nest);
525        let sep = if nest == 0 {
526            arena.space()
527        } else {
528            arena.softline()
529        };
530        let expr = arena.intersperse([lhs, op, rhs], sep).group();
531        pretty_parenthesized_doc(expr, arena).group()
532    }
533}
534
535impl PrettyDoc for UniOp {
536    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
537    where
538        D: DocAllocator<'b, A>,
539        D::Doc: Clone,
540        A: Clone,
541    {
542        // TODO NOT LIKE, NOT IN, NOT BETWEEN?
543        let UniOp { kind, expr } = self;
544        let (sym, paren) = match kind {
545            UniOpKind::Pos => ("+", false),
546            UniOpKind::Neg => ("-", false),
547            UniOpKind::Not => ("NOT ", true),
548        };
549        let op = arena.text(sym);
550        let expr = expr.pretty_doc(arena);
551        if paren {
552            let open = arena.text("(");
553            let close = arena.text(")");
554            arena.concat([op, open, expr, close]).group()
555        } else {
556            arena.concat([op, expr]).group()
557        }
558    }
559}
560
561impl PrettyDoc for Like {
562    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
563    where
564        D: DocAllocator<'b, A>,
565        D::Doc: Clone,
566        A: Clone,
567    {
568        let Like {
569            value,
570            pattern,
571            escape,
572        } = self;
573
574        let sep = arena.space();
575        let value = value.pretty_doc(arena);
576        let kw_like = arena.text("LIKE");
577        let pattern = pattern.pretty_doc(arena);
578        if let Some(escape) = escape {
579            let kw_esc = arena.text("ESCAPE");
580            let escape = escape.pretty_doc(arena);
581            arena.intersperse([value, kw_like, pattern, kw_esc, escape], sep)
582        } else {
583            arena.intersperse([value, kw_like, pattern], sep)
584        }
585        .group()
586    }
587}
588
589impl PrettyDoc for Between {
590    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
591    where
592        D: DocAllocator<'b, A>,
593        D::Doc: Clone,
594        A: Clone,
595    {
596        let Between { value, from, to } = self;
597
598        let value = value.pretty_doc(arena);
599        let kw_b = arena.text("BETWEEN");
600        let kw_a = arena.text("AND");
601        let from = from.pretty_doc(arena);
602        let to = to.pretty_doc(arena);
603        let sep = arena.space();
604        let expr = arena
605            .intersperse([value, kw_b, from, kw_a, to], sep)
606            .group();
607        expr.group()
608    }
609}
610
611impl PrettyDoc for In {
612    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
613    where
614        D: DocAllocator<'b, A>,
615        D::Doc: Clone,
616        A: Clone,
617    {
618        let In { lhs, rhs } = self;
619
620        let kw_in = arena.text("IN");
621        let lhs = lhs.pretty_doc(arena);
622        let rhs = rhs.pretty_doc(arena);
623        let sep = arena.space();
624        let expr = arena.intersperse([lhs, kw_in, rhs], sep).group();
625        expr.group()
626    }
627}
628
629impl PrettyDoc for Case {
630    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
631    where
632        D: DocAllocator<'b, A>,
633        D::Doc: Clone,
634        A: Clone,
635    {
636        match self {
637            Case::SimpleCase(inner) => inner.pretty_doc(arena),
638            Case::SearchedCase(inner) => inner.pretty_doc(arena),
639        }
640    }
641}
642
643impl PrettyDoc for SimpleCase {
644    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
645    where
646        D: DocAllocator<'b, A>,
647        D::Doc: Clone,
648        A: Clone,
649    {
650        let SimpleCase {
651            expr,
652            cases,
653            default,
654        } = self;
655
656        let search = expr.pretty_doc(arena);
657        let branches = case_branches(arena, cases, default);
658        pretty_seq_doc(
659            branches,
660            "CASE",
661            Some(search),
662            "END",
663            " ",
664            PRETTY_INDENT_MINOR_NEST,
665            arena,
666        )
667    }
668}
669
670impl PrettyDoc for SearchedCase {
671    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
672    where
673        D: DocAllocator<'b, A>,
674        D::Doc: Clone,
675        A: Clone,
676    {
677        let SearchedCase { cases, default } = self;
678
679        let branches = case_branches(arena, cases, default);
680        pretty_seq_doc(
681            branches,
682            "CASE",
683            None,
684            "END",
685            " ",
686            PRETTY_INDENT_MINOR_NEST,
687            arena,
688        )
689    }
690}
691
692impl PrettyDoc for Struct {
693    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
694    where
695        D: DocAllocator<'b, A>,
696        D::Doc: Clone,
697        A: Clone,
698    {
699        let fields = self.fields.iter().map(|expr_pair| {
700            let k = expr_pair.first.pretty_doc(arena);
701            let v = expr_pair.second.pretty_doc(arena);
702            let sep = arena.text(": ");
703
704            k.append(sep).group().append(v).group()
705        });
706        pretty_seq_doc(fields, "{", None, "}", ",", PRETTY_INDENT_MINOR_NEST, arena)
707    }
708}
709
710impl PrettyDoc for StructLit {
711    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
712    where
713        D: DocAllocator<'b, A>,
714        D::Doc: Clone,
715        A: Clone,
716    {
717        let fields = self.fields.iter().map(|expr_pair| {
718            let k = expr_pair.first.pretty_doc(arena);
719            let v = expr_pair.second.pretty_doc(arena);
720            let sep = arena.text(": ");
721
722            k.append(sep).group().append(v).group()
723        });
724        pretty_seq_doc(fields, "{", None, "}", ",", PRETTY_INDENT_MINOR_NEST, arena)
725    }
726}
727
728impl PrettyDoc for Bag {
729    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
730    where
731        D: DocAllocator<'b, A>,
732        D::Doc: Clone,
733        A: Clone,
734    {
735        pretty_seq(
736            &self.values,
737            "<<",
738            ">>",
739            ",",
740            PRETTY_INDENT_MINOR_NEST,
741            arena,
742        )
743    }
744}
745
746impl PrettyDoc for List {
747    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
748    where
749        D: DocAllocator<'b, A>,
750        D::Doc: Clone,
751        A: Clone,
752    {
753        pretty_seq(&self.values, "[", "]", ",", PRETTY_INDENT_MINOR_NEST, arena)
754    }
755}
756
757impl PrettyDoc for BagLit {
758    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
759    where
760        D: DocAllocator<'b, A>,
761        D::Doc: Clone,
762        A: Clone,
763    {
764        pretty_seq(
765            &self.values,
766            "<<",
767            ">>",
768            ",",
769            PRETTY_INDENT_MINOR_NEST,
770            arena,
771        )
772    }
773}
774
775impl PrettyDoc for ListLit {
776    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
777    where
778        D: DocAllocator<'b, A>,
779        D::Doc: Clone,
780        A: Clone,
781    {
782        pretty_seq(&self.values, "[", "]", ",", PRETTY_INDENT_MINOR_NEST, arena)
783    }
784}
785
786impl PrettyDoc for Call {
787    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
788    where
789        D: DocAllocator<'b, A>,
790        D::Doc: Clone,
791        A: Clone,
792    {
793        let name = self.func_name.pretty_doc(arena);
794        let list = pretty_list(&self.args, 0, arena);
795        name.append(arena.text("("))
796            .append(list.nest(PRETTY_INDENT_MINOR_NEST))
797            .append(arena.text(")"))
798    }
799}
800
801impl PrettyDoc for CallAgg {
802    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
803    where
804        D: DocAllocator<'b, A>,
805        D::Doc: Clone,
806        A: Clone,
807    {
808        let name = self.func_name.pretty_doc(arena);
809        let list = pretty_list(&self.args, 0, arena);
810        name.append(arena.text("("))
811            .append(list.nest(PRETTY_INDENT_MINOR_NEST))
812            .append(arena.text(")"))
813    }
814}
815
816impl PrettyDoc for CallArg {
817    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
818    where
819        D: DocAllocator<'b, A>,
820        D::Doc: Clone,
821        A: Clone,
822    {
823        match self {
824            CallArg::Star() => arena.text("*"),
825            CallArg::Positional(arg) => arg.pretty_doc(arena),
826            CallArg::PositionalType(_) => {
827                todo!("CallArg::PositionalType")
828            }
829            CallArg::Named(arg) => arg.pretty_doc(arena),
830            CallArg::NamedType(_) => {
831                todo!("CallArg::NamedType")
832            }
833        }
834    }
835}
836
837impl PrettyDoc for CallArgNamed {
838    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
839    where
840        D: DocAllocator<'b, A>,
841        D::Doc: Clone,
842        A: Clone,
843    {
844        let CallArgNamed { name, value } = self;
845        let name = name.pretty_doc(arena);
846        let value = value.pretty_doc(arena);
847        pretty_seperated_doc(":", [name, value], 0, arena)
848    }
849}
850
851impl PrettyDoc for SymbolPrimitive {
852    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
853    where
854        D: DocAllocator<'b, A>,
855        D::Doc: Clone,
856        A: Clone,
857    {
858        let sym = arena.text(self.value.as_str());
859        match self.case {
860            CaseSensitivity::CaseSensitive => arena.text("\"").append(sym).append(arena.text("\"")),
861            CaseSensitivity::CaseInsensitive => sym,
862        }
863    }
864}
865
866impl PrettyDoc for FromClause {
867    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
868    where
869        D: DocAllocator<'b, A>,
870        D::Doc: Clone,
871        A: Clone,
872    {
873        pretty_prefixed_expr("FROM", &self.source, PRETTY_INDENT_MINOR_NEST, arena)
874    }
875}
876
877impl PrettyDoc for FromSource {
878    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
879    where
880        D: DocAllocator<'b, A>,
881        D::Doc: Clone,
882        A: Clone,
883    {
884        match self {
885            FromSource::FromLet(fl) => fl.pretty_doc(arena),
886            FromSource::Join(join) => join.pretty_doc(arena),
887        }
888    }
889}
890
891impl PrettyDoc for FromLet {
892    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
893    where
894        D: DocAllocator<'b, A>,
895        D::Doc: Clone,
896        A: Clone,
897    {
898        let FromLet {
899            expr,
900            kind,
901            as_alias,
902            at_alias,
903            by_alias,
904        } = self;
905
906        let expr = expr.pretty_doc(arena);
907        let as_alias = pretty_as_alias(as_alias.as_ref(), arena);
908        let at_alias = pretty_at_alias(at_alias.as_ref(), arena);
909        let by_alias = pretty_by_alias(by_alias.as_ref(), arena);
910        let aliases: Vec<_> = [as_alias, at_alias, by_alias]
911            .into_iter()
912            .flatten()
913            .collect();
914
915        let clause = match kind {
916            FromLetKind::Scan => expr,
917            FromLetKind::Unpivot => pretty_prefixed_doc("UNPIVOT", expr, arena),
918            FromLetKind::GraphTable => pretty_prefixed_doc("GRAPH_TABLE", expr, arena),
919        };
920
921        if aliases.is_empty() {
922            clause
923        } else {
924            clause.append(arena.concat(aliases).group())
925        }
926        .group()
927    }
928}
929
930impl PrettyDoc for Join {
931    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
932    where
933        D: DocAllocator<'b, A>,
934        D::Doc: Clone,
935        A: Clone,
936    {
937        let Join {
938            kind,
939            left,
940            right,
941            predicate,
942        } = self;
943
944        let arms = [left.as_ref(), right.as_ref()];
945        let kw_join = match kind {
946            JoinKind::Cross => " CROSS JOIN ",
947            JoinKind::Inner => " INNER JOIN ",
948            JoinKind::Left => " LEFT JOIN ",
949            JoinKind::Right => " RIGHT JOIN ",
950            JoinKind::Full => " FULL JOIN ",
951        };
952
953        match (kind, predicate) {
954            (JoinKind::Cross, Some(_)) => {
955                todo!("CROSS JOIN with predicate")
956            }
957            (JoinKind::Cross, None) => pretty_list(arms, 0, arena),
958            (_, None) => pretty_seperated(kw_join, arms, 0, arena),
959            (_, Some(pred)) => match &pred.node {
960                JoinSpec::Natural => {
961                    let kw = arena.text(" NATURAL").append(kw_join);
962                    pretty_seperated(kw, arms, 0, arena)
963                }
964                JoinSpec::On(on) => {
965                    let join = pretty_seperated(kw_join, arms, 0, arena);
966                    let pred = arena
967                        .softline()
968                        .append(arena.text("ON"))
969                        .append(arena.softline())
970                        .append(on.pretty_doc(arena).nest(PRETTY_INDENT_MINOR_NEST));
971                    join.append(pred)
972                }
973                JoinSpec::Using(using) => {
974                    let join = pretty_seperated(kw_join, arms, 0, arena);
975                    let using = pretty_list(using, PRETTY_INDENT_MINOR_NEST, arena);
976                    let pred = arena
977                        .softline()
978                        .append(arena.text("USING"))
979                        .append(arena.softline())
980                        .append(using);
981                    join.append(pred)
982                }
983            },
984        }
985        .group()
986    }
987}
988
989impl PrettyDoc for Let {
990    fn pretty_doc<'b, D, A>(&'b self, _arena: &'b D) -> DocBuilder<'b, D, A>
991    where
992        D: DocAllocator<'b, A>,
993        D::Doc: Clone,
994        A: Clone,
995    {
996        todo!("LET")
997    }
998}
999
1000impl PrettyDoc for WhereClause {
1001    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1002    where
1003        D: DocAllocator<'b, A>,
1004        D::Doc: Clone,
1005        A: Clone,
1006    {
1007        pretty_prefixed_expr("WHERE", &self.expr, PRETTY_INDENT_MINOR_NEST, arena)
1008    }
1009}
1010
1011impl PrettyDoc for GroupByExpr {
1012    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1013    where
1014        D: DocAllocator<'b, A>,
1015        D::Doc: Clone,
1016        A: Clone,
1017    {
1018        let GroupByExpr {
1019            strategy,
1020            keys,
1021            group_as_alias,
1022        } = self;
1023
1024        let mut doc = match strategy {
1025            None => arena.text("GROUP"),
1026            Some(GroupingStrategy::GroupFull) => arena.text("GROUP ALL"),
1027            Some(GroupingStrategy::GroupPartial) => arena.text("GROUP PARTIAL"),
1028        };
1029
1030        if !keys.is_empty() {
1031            doc = doc.append(arena.space()).append(arena.text("BY")).group();
1032            doc = doc.append(arena.softline()).append(pretty_list(
1033                keys,
1034                PRETTY_INDENT_MINOR_NEST,
1035                arena,
1036            ));
1037        }
1038
1039        match group_as_alias {
1040            None => doc,
1041            Some(gas) => {
1042                let gas = pretty_source_as_alias("GROUP", Some(gas), arena);
1043                doc.append(gas)
1044            }
1045        }
1046        .group()
1047    }
1048}
1049
1050impl PrettyDoc for GroupKey {
1051    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1052    where
1053        D: DocAllocator<'b, A>,
1054        D::Doc: Clone,
1055        A: Clone,
1056    {
1057        pretty_source_as_alias(&self.expr, self.as_alias.as_ref(), arena)
1058            .unwrap_or_else(|| self.expr.pretty_doc(arena))
1059    }
1060}
1061
1062impl PrettyDoc for HavingClause {
1063    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1064    where
1065        D: DocAllocator<'b, A>,
1066        D::Doc: Clone,
1067        A: Clone,
1068    {
1069        pretty_prefixed_expr("HAVING", &self.expr, PRETTY_INDENT_MINOR_NEST, arena)
1070    }
1071}
1072
1073impl PrettyDoc for OrderByExpr {
1074    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1075    where
1076        D: DocAllocator<'b, A>,
1077        D::Doc: Clone,
1078        A: Clone,
1079    {
1080        if self.sort_specs.is_empty() {
1081            arena.text("ORDER BY PRESERVE")
1082        } else {
1083            pretty_prefixed_doc(
1084                "ORDER BY",
1085                pretty_list(&self.sort_specs, PRETTY_INDENT_MINOR_NEST, arena),
1086                arena,
1087            )
1088        }
1089        .group()
1090    }
1091}
1092
1093impl PrettyDoc for SortSpec {
1094    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1095    where
1096        D: DocAllocator<'b, A>,
1097        D::Doc: Clone,
1098        A: Clone,
1099    {
1100        let SortSpec {
1101            expr,
1102            ordering_spec,
1103            null_ordering_spec,
1104        } = self;
1105        let mut doc = expr.pretty_doc(arena);
1106        if let Some(os) = ordering_spec {
1107            let os = arena.space().append(os.pretty_doc(arena)).group();
1108            doc = doc.append(os)
1109        };
1110        if let Some(nos) = null_ordering_spec {
1111            let nos = arena.space().append(nos.pretty_doc(arena)).group();
1112            doc = doc.append(nos)
1113        };
1114
1115        doc.group()
1116    }
1117}
1118
1119impl PrettyDoc for OrderingSpec {
1120    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1121    where
1122        D: DocAllocator<'b, A>,
1123        D::Doc: Clone,
1124        A: Clone,
1125    {
1126        arena.text(match self {
1127            OrderingSpec::Asc => "ASC",
1128            OrderingSpec::Desc => "DESC",
1129        })
1130    }
1131}
1132
1133impl PrettyDoc for NullOrderingSpec {
1134    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1135    where
1136        D: DocAllocator<'b, A>,
1137        D::Doc: Clone,
1138        A: Clone,
1139    {
1140        arena.text(match self {
1141            NullOrderingSpec::First => "NULLS FIRST",
1142            NullOrderingSpec::Last => "NULLS LAST",
1143        })
1144    }
1145}
1146
1147impl PrettyDoc for LimitOffsetClause {
1148    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
1149    where
1150        D: DocAllocator<'b, A>,
1151        D::Doc: Clone,
1152        A: Clone,
1153    {
1154        let limit = self
1155            .limit
1156            .as_ref()
1157            .map(|l| pretty_prefixed_expr("LIMIT", l, PRETTY_INDENT_MINOR_NEST, arena));
1158
1159        let offset = self
1160            .offset
1161            .as_ref()
1162            .map(|o| pretty_prefixed_expr("OFFSET", o, PRETTY_INDENT_MINOR_NEST, arena));
1163
1164        match (limit, offset) {
1165            (None, None) => unreachable!(),
1166            (Some(limit), None) => limit,
1167            (None, Some(offset)) => offset,
1168            (Some(limit), Some(offset)) => limit.append(arena.softline()).append(offset),
1169        }
1170    }
1171}
1172
1173fn case_branches<'b, D, A>(
1174    arena: &'b D,
1175    cases: &'b [ExprPair],
1176    default: &'b Option<Box<Expr>>,
1177) -> impl Iterator<Item = DocBuilder<'b, D, A>>
1178where
1179    D: DocAllocator<'b, A>,
1180    D::Doc: Clone,
1181    A: Clone + 'b,
1182{
1183    cases
1184        .iter()
1185        .map(|ExprPair { first, second }| {
1186            let kw_when = arena.text("WHEN");
1187            let test = first.pretty_doc(arena);
1188            let kw_then = arena.text("THEN");
1189            let then = second.pretty_doc(arena);
1190            arena
1191                .intersperse([kw_when, test, kw_then, then], arena.space())
1192                .group()
1193        })
1194        .chain(
1195            default
1196                .iter()
1197                .map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)).group()),
1198        )
1199}
1200
1201fn pretty_prefixed_expr<'b, P, D, A>(
1202    annot: &'static str,
1203    expr: &'b P,
1204    nest: isize,
1205    arena: &'b D,
1206) -> DocBuilder<'b, D, A>
1207where
1208    P: PrettyDoc,
1209    D: DocAllocator<'b, A>,
1210    D::Doc: Clone,
1211    A: Clone,
1212{
1213    pretty_prefixed_doc(annot, expr.pretty_doc(arena).nest(nest), arena)
1214}
1215
1216fn pretty_parenthesized_expr<'b, P, D, A>(
1217    expr: &'b P,
1218    nest: isize,
1219    arena: &'b D,
1220) -> DocBuilder<'b, D, A>
1221where
1222    P: PrettyDoc,
1223    D: DocAllocator<'b, A>,
1224    D::Doc: Clone,
1225    A: Clone,
1226{
1227    pretty_parenthesized_doc(expr.pretty_doc(arena).nest(nest), arena)
1228}
1229
1230fn pretty_alias_helper<'b, D, A>(
1231    kw: &'static str,
1232    sym: Option<&'b SymbolPrimitive>,
1233    arena: &'b D,
1234) -> Option<DocBuilder<'b, D, A>>
1235where
1236    D: DocAllocator<'b, A>,
1237    D::Doc: Clone,
1238    A: Clone,
1239{
1240    sym.map(|sym| {
1241        arena
1242            .space()
1243            .append(arena.text(kw))
1244            .append(arena.space())
1245            .append(sym.pretty_doc(arena))
1246            .group()
1247    })
1248}
1249
1250fn pretty_source_as_alias<'b, S, D, A>(
1251    source: &'b S,
1252    as_alias: Option<&'b SymbolPrimitive>,
1253    arena: &'b D,
1254) -> Option<DocBuilder<'b, D, A>>
1255where
1256    S: PrettyDoc + ?Sized,
1257    D: DocAllocator<'b, A>,
1258    D::Doc: Clone,
1259    A: Clone,
1260{
1261    pretty_as_alias(as_alias, arena).map(|alias| {
1262        let expr = source.pretty_doc(arena);
1263        arena.concat([expr, alias]).group()
1264    })
1265}
1266
1267fn pretty_as_alias<'b, D, A>(
1268    sym: Option<&'b SymbolPrimitive>,
1269    arena: &'b D,
1270) -> Option<DocBuilder<'b, D, A>>
1271where
1272    D: DocAllocator<'b, A>,
1273    D::Doc: Clone,
1274    A: Clone,
1275{
1276    pretty_alias_helper("AS", sym, arena)
1277}
1278
1279fn pretty_at_alias<'b, D, A>(
1280    sym: Option<&'b SymbolPrimitive>,
1281    arena: &'b D,
1282) -> Option<DocBuilder<'b, D, A>>
1283where
1284    D: DocAllocator<'b, A>,
1285    D::Doc: Clone,
1286    A: Clone,
1287{
1288    pretty_alias_helper("AT", sym, arena)
1289}
1290
1291fn pretty_by_alias<'b, D, A>(
1292    sym: Option<&'b SymbolPrimitive>,
1293    arena: &'b D,
1294) -> Option<DocBuilder<'b, D, A>>
1295where
1296    D: DocAllocator<'b, A>,
1297    D::Doc: Clone,
1298    A: Clone,
1299{
1300    pretty_alias_helper("BY", sym, arena)
1301}
partiql_ast/pretty/mod.rs

partiql_ast/pretty/
mod.rs
