Skip to content

Commit 6dc088c

Browse files
jpschorrjpschorr
authored
Improvements to pretty-printing for AND/OR/CASE/PIVOT (#482)
1 parent 2da9cdb commit 6dc088c

18 files changed

+485
-257
lines changed

‎CHANGELOG.md‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77

88
## [Unreleased]
99
### Changed
10-
10+
- partiql-ast: fixed pretty-printing of `PIVOT`
11+
- partiql-ast: improved pretty-printing of `CASE` and various clauses
12+
-
1113
### Added
1214

1315
### Fixed

‎partiql-ast/src/pretty.rs‎

Lines changed: 101 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -213,14 +213,24 @@ impl PrettyDoc for Select {
213213
D::Doc: Clone,
214214
A: Clone,
215215
{
216+
fn format<'b, C, D, A>(child: &'b C, arena: &'b D) -> DocBuilder<'b, D, A>
217+
where
218+
D: DocAllocator<'b, A>,
219+
D::Doc: Clone,
220+
A: Clone,
221+
C: PrettyDoc,
222+
{
223+
child.pretty_doc(arena).group()
224+
}
225+
216226
fn delegate<'b, C, D, A>(child: &'b Option<C>, arena: &'b D) -> Option<DocBuilder<'b, D, A>>
217227
where
218228
D: DocAllocator<'b, A>,
219229
D::Doc: Clone,
220230
A: Clone,
221231
C: PrettyDoc,
222232
{
223-
child.as_ref().map(|inner| inner.pretty_doc(arena).group())
233+
child.as_ref().map(|inner| format(inner, arena))
224234
}
225235

226236
let Select {
@@ -232,25 +242,32 @@ impl PrettyDoc for Select {
232242
group_by,
233243
having,
234244
} = self;
235-
let clauses = [
236-
Some(project.pretty_doc(arena).group()),
245+
let mut clauses = [
246+
Some(format(project, arena)),
237247
delegate(exclude, arena),
238-
from.as_ref().map(|inner| inner.pretty_doc(arena).group()),
239-
from_let
240-
.as_ref()
241-
.map(|inner| inner.pretty_doc(arena).group()),
242-
where_clause
243-
.as_ref()
244-
.map(|inner| inner.pretty_doc(arena).group()),
245-
group_by
246-
.as_ref()
247-
.map(|inner| inner.pretty_doc(arena).group()),
248-
having.as_ref().map(|inner| inner.pretty_doc(arena).group()),
248+
delegate(from, arena),
249+
delegate(from_let, arena),
250+
delegate(where_clause, arena),
251+
delegate(group_by, arena),
252+
delegate(having, arena),
249253
]
250254
.into_iter()
251255
.flatten();
252256

253-
arena.intersperse(clauses, arena.softline()).group()
257+
let mut result = arena.nil();
258+
let separator = arena.line();
259+
if let Some(first) = clauses.next() {
260+
let mut curr = first;
261+
262+
for clause in clauses {
263+
result = result.append(curr.append(separator.clone()).group());
264+
curr = clause;
265+
}
266+
267+
result = result.append(curr);
268+
}
269+
270+
result
254271
}
255272
}
256273

@@ -282,9 +299,9 @@ impl PrettyDoc for ProjectionKind {
282299
}
283300
ProjectionKind::ProjectPivot(ProjectPivot { key, value }) => {
284301
let parts = [
285-
key.pretty_doc(arena),
286-
arena.text("AT"),
287302
value.pretty_doc(arena),
303+
arena.text("AT"),
304+
key.pretty_doc(arena),
288305
];
289306
let decl = arena.intersperse(parts, arena.space()).group();
290307
pretty_annotated_doc("PIVOT", decl, arena)
@@ -401,6 +418,7 @@ impl PrettyDoc for Expr {
401418
unreachable!();
402419
}
403420
}
421+
.group()
404422
}
405423
}
406424

@@ -570,10 +588,13 @@ impl PrettyDoc for BinOp {
570588
let op = arena.text(sym);
571589
let lhs = lhs.pretty_doc(arena).nest(nest);
572590
let rhs = rhs.pretty_doc(arena).nest(nest);
573-
let sep = arena.space();
591+
let sep = if nest == 0 {
592+
arena.space()
593+
} else {
594+
arena.softline()
595+
};
574596
let expr = arena.intersperse([lhs, op, rhs], sep).group();
575-
let paren_expr = [arena.text("("), expr, arena.text(")")];
576-
arena.concat(paren_expr).group()
597+
pretty_parenthesized_doc(expr, arena).group()
577598
}
578599
}
579600

@@ -698,30 +719,9 @@ impl PrettyDoc for SimpleCase {
698719
default,
699720
} = self;
700721

701-
let kw_case = arena.text("CASE");
702722
let search = expr.pretty_doc(arena);
703-
let branches = cases.iter().map(|ExprPair { first, second }| {
704-
let kw_when = arena.text("WHEN");
705-
let test = first.pretty_doc(arena);
706-
let kw_then = arena.text("THEN");
707-
let then = second.pretty_doc(arena);
708-
arena
709-
.intersperse([kw_when, test, kw_then, then], arena.space())
710-
.group()
711-
});
712-
let branches = arena
713-
.intersperse(branches, arena.softline())
714-
.group()
715-
.nest(MINOR_NEST_INDENT);
716-
let default = default
717-
.as_ref()
718-
.map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)));
719-
720-
if let Some(default) = default {
721-
arena.intersperse([kw_case, search, branches, default], arena.softline())
722-
} else {
723-
arena.intersperse([kw_case, search, branches], arena.softline())
724-
}
723+
let branches = case_branches(arena, cases, default);
724+
pretty_seq_doc(branches, "CASE", Some(search), "END", " ", arena)
725725
}
726726
}
727727

@@ -734,30 +734,37 @@ impl PrettyDoc for SearchedCase {
734734
{
735735
let SearchedCase { cases, default } = self;
736736

737-
let kw_case = arena.text("CASE");
738-
let branches = cases.iter().map(|ExprPair { first, second }| {
737+
let branches = case_branches(arena, cases, default);
738+
pretty_seq_doc(branches, "CASE", None, "END", " ", arena)
739+
}
740+
}
741+
742+
fn case_branches<'b, D, A>(
743+
arena: &'b D,
744+
cases: &'b [ExprPair],
745+
default: &'b Option<Box<Expr>>,
746+
) -> impl Iterator<Item = DocBuilder<'b, D, A>>
747+
where
748+
D: DocAllocator<'b, A>,
749+
D::Doc: Clone,
750+
A: Clone + 'b,
751+
{
752+
cases
753+
.iter()
754+
.map(|ExprPair { first, second }| {
739755
let kw_when = arena.text("WHEN");
740756
let test = first.pretty_doc(arena);
741757
let kw_then = arena.text("THEN");
742758
let then = second.pretty_doc(arena);
743759
arena
744760
.intersperse([kw_when, test, kw_then, then], arena.space())
745761
.group()
746-
});
747-
let branches = arena
748-
.intersperse(branches, arena.softline())
749-
.group()
750-
.nest(MINOR_NEST_INDENT);
751-
let default = default
752-
.as_ref()
753-
.map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)));
754-
755-
if let Some(default) = default {
756-
arena.intersperse([kw_case, branches, default], arena.softline())
757-
} else {
758-
arena.intersperse([kw_case, branches], arena.softline())
759-
}
760-
}
762+
})
763+
.chain(
764+
default
765+
.iter()
766+
.map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)).group()),
767+
)
761768
}
762769

763770
impl PrettyDoc for Struct {
@@ -1258,41 +1265,61 @@ where
12581265
D::Doc: Clone,
12591266
A: Clone,
12601267
{
1261-
arena
1262-
.text("(")
1263-
.append(arena.space())
1264-
.append(doc)
1265-
.append(arena.space())
1266-
.append(arena.text(")"))
1267-
.group()
1268+
arena.text("(").append(doc).append(arena.text(")")).group()
12681269
}
12691270

1270-
fn pretty_seq<'i, 'b, I, P, D, A>(
1271-
list: I,
1271+
fn pretty_seq_doc<'i, 'b, I, E, D, A>(
1272+
seq: I,
12721273
start: &'static str,
1274+
qualifier: Option<E>,
12731275
end: &'static str,
12741276
sep: &'static str,
12751277
arena: &'b D,
12761278
) -> DocBuilder<'b, D, A>
12771279
where
1278-
I: IntoIterator<Item = &'b P>,
1279-
P: PrettyDoc + 'b,
1280+
E: Pretty<'b, D, A>,
1281+
I: IntoIterator<Item = E>,
12801282
D: DocAllocator<'b, A>,
12811283
D::Doc: Clone,
12821284
A: Clone,
12831285
{
12841286
let start = arena.text(start);
12851287
let end = arena.text(end);
12861288
let sep = arena.text(sep).append(arena.line());
1287-
let seq = list.into_iter().map(|l| l.pretty_doc(arena));
1288-
let body = arena.line().append(arena.intersperse(seq, sep)).group();
1289+
let start = if let Some(qual) = qualifier {
1290+
start.append(arena.space()).append(qual)
1291+
} else {
1292+
start
1293+
};
1294+
let body = arena
1295+
.line()
1296+
.append(arena.intersperse(seq, sep))
1297+
.append(arena.line())
1298+
.group();
12891299
start
12901300
.append(body.nest(MINOR_NEST_INDENT))
1291-
.append(arena.line())
12921301
.append(end)
12931302
.group()
12941303
}
12951304

1305+
fn pretty_seq<'i, 'b, I, P, D, A>(
1306+
list: I,
1307+
start: &'static str,
1308+
end: &'static str,
1309+
sep: &'static str,
1310+
arena: &'b D,
1311+
) -> DocBuilder<'b, D, A>
1312+
where
1313+
I: IntoIterator<Item = &'b P>,
1314+
P: PrettyDoc + 'b,
1315+
D: DocAllocator<'b, A>,
1316+
D::Doc: Clone,
1317+
A: Clone,
1318+
{
1319+
let seq = list.into_iter().map(|l| l.pretty_doc(arena));
1320+
pretty_seq_doc(seq, start, None, end, sep, arena)
1321+
}
1322+
12961323
fn pretty_list<'b, I, P, D, A>(list: I, nest: isize, arena: &'b D) -> DocBuilder<'b, D, A>
12971324
where
12981325
I: IntoIterator<Item = &'b P>,

‎partiql/tests/pretty.rs‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use itertools::Itertools;
2+
use partiql_ast::ast::{AstNode, TopLevelQuery};
23
use partiql_ast::pretty::ToPretty;
34
use partiql_parser::ParserResult;
45

@@ -15,13 +16,18 @@ fn pretty_print_test(name: &str, statement: &str) {
1516
assert!(res.is_ok());
1617
let res = res.unwrap();
1718

18-
// TODO https://github.com/partiql/partiql-lang-rust/issues/473
19+
pretty_print_output_test(name, statement, &res.ast);
20+
pretty_print_roundtrip_test(&res.ast);
21+
}
1922

23+
#[track_caller]
24+
fn pretty_print_output_test(name: &str, statement: &str, statement_ast: &AstNode<TopLevelQuery>) {
25+
// TODO https://github.com/partiql/partiql-lang-rust/issues/473
2026
let doc = [180, 120, 80, 40, 30, 20, 10]
2127
.into_iter()
2228
.map(|w| {
2329
let header = format!("{:-<w$}", "");
24-
let ast = format!("{}\n", res.ast.to_pretty_string(w).unwrap());
30+
let ast = format!("{}\n", statement_ast.to_pretty_string(w).unwrap());
2531
format!("{header}\n{ast}")
2632
})
2733
.join("\n");
@@ -33,6 +39,18 @@ fn pretty_print_test(name: &str, statement: &str) {
3339
insta::assert_snapshot!(name, doc)
3440
}
3541

42+
#[track_caller]
43+
fn pretty_print_roundtrip_test(statement_ast: &AstNode<TopLevelQuery>) {
44+
let pretty = statement_ast.to_pretty_string(40).unwrap();
45+
46+
let reparsed = parse(pretty.as_str());
47+
assert!(reparsed.is_ok());
48+
49+
let pretty2 = reparsed.unwrap().ast.to_pretty_string(40).unwrap();
50+
51+
assert_eq!(pretty, pretty2);
52+
}
53+
3654
#[test]
3755
fn pretty() {
3856
pretty_print_test(
@@ -169,6 +187,18 @@ fn pretty_pivot() {
169187
);
170188
}
171189

190+
#[test]
191+
fn pretty_ands_and_ors() {
192+
pretty_print_test(
193+
"pretty_ands_and_ors",
194+
"
195+
SELECT *
196+
FROM test_data AS test_data
197+
WHERE ((((test_data.country_code <> 'Distinctio.') AND ((test_data.* < false) AND (NOT (test_data.description LIKE 'Esse solam.') AND NOT (test_data.transaction_id LIKE 'Esset accusata.')))) OR (test_data.test_address <> 'Potest. Sed.')) AND (test_data.* > -28.146858383543243))
198+
",
199+
);
200+
}
201+
172202
#[test]
173203
fn pretty_exclude() {
174204
pretty_print_test(

0 commit comments

Comments
 (0)