1use itertools::Itertools;
2use rowan::Direction;
3use squawk_syntax::ast::{self, AstNode, LitKind};
4use squawk_syntax::quote::quote_column_alias;
5use squawk_syntax::{SyntaxKind, SyntaxNode, SyntaxToken};
6use tiny_pretty::Doc;
7use tiny_pretty::{PrintOptions, print};
8
9fn build_source_file(source_file: &ast::SourceFile) -> Doc<'_> {
13 let mut doc = Doc::nil();
14 for el in source_file.syntax().children_with_tokens() {
15 match el {
16 rowan::NodeOrToken::Node(node) => {
17 if let Some(stmt) = ast::Stmt::cast(node) {
18 match stmt {
19 ast::Stmt::Select(select) => {
20 doc = doc.append(build_select_doc(&select));
21 }
22 ast::Stmt::CreateTable(create_table) => {
23 doc = doc.append(build_create_table(&create_table));
24 }
25 _ => (),
26 }
27 }
28 }
29 rowan::NodeOrToken::Token(token) => {
30 if token.kind() == SyntaxKind::COMMENT {
31 doc = doc.append(Doc::text(token.text().to_string()));
32 } else if token.kind() == SyntaxKind::WHITESPACE {
33 let lines = token.text().lines().count();
35 if lines >= 2 {
36 doc = doc.append(Doc::empty_line()).append(Doc::empty_line());
37 } else {
38 doc = doc.append(Doc::empty_line());
39 }
40 }
41 }
42 }
43 }
44 doc
45}
46
47fn build_create_table<'a>(create_table: &ast::CreateTable) -> Doc<'a> {
48 let mut doc = Doc::text("create")
49 .append(Doc::space())
50 .append(Doc::text("table"))
51 .append(Doc::space())
52 .append(Doc::text(
53 create_table.table_name().unwrap().syntax().to_string(),
54 ))
55 .append(Doc::text("("))
56 .append(
57 Doc::line_or_nil()
58 .append(Doc::list(
59 Itertools::intersperse(
60 create_table
61 .table_arg_list()
62 .unwrap()
63 .args()
64 .map(build_table_arg),
65 Doc::text(",").append(Doc::hard_line()),
66 )
67 .collect(),
68 ))
69 .nest(2)
70 .append(Doc::line_or_nil())
71 .group(),
72 )
73 .append(Doc::text(")"));
74
75 doc = doc.append(build_semicolon(create_table.semicolon_token()));
76
77 doc
78}
79
80fn build_table_arg<'a>(create_table: ast::TableArg) -> Doc<'a> {
81 match create_table {
82 ast::TableArg::Column(column) => build_column_name(column.name().unwrap())
83 .append(Doc::space())
84 .append(Doc::text(column.ty().unwrap().syntax().to_string())),
85 ast::TableArg::LikeClause(_like_clause) => todo!(),
86 ast::TableArg::TableConstraint(_table_constraint) => todo!(),
87 }
88}
89
90fn build_select_doc<'a>(select: &ast::Select) -> Doc<'a> {
91 let mut doc = Doc::text("select").append(Doc::line_or_space());
92
93 if let Some(select_clause) = select.select_clause() {
94 if let Some(distinct_clause) = select_clause.distinct_clause() {
95 doc = doc.append(leading_comments(distinct_clause.syntax()));
96 doc = doc.append(Doc::text("distinct")).append(Doc::space());
97 }
98 if let Some(all_token) = select_clause.all_token() {
99 doc = doc.append(leading_comments_token(&all_token));
100 doc = doc.append(Doc::text("all")).append(Doc::space());
101 }
102 if let Some(target_list) = select_clause.target_list() {
103 doc = doc.append(leading_comments(target_list.syntax()));
104 doc = doc
105 .append(Doc::list(
106 Itertools::intersperse(
107 target_list.targets().flat_map(build_target),
108 Doc::text(",").append(Doc::line_or_space()),
109 )
110 .collect(),
111 ))
112 .nest(2);
113 }
114 }
115
116 if let Some(from) = &select.from_clause() {
117 doc = doc.append(
118 Doc::line_or_space()
119 .append(Doc::text("from"))
120 .append(Doc::space())
121 .append(Doc::text(
122 from.from_items().next().unwrap().syntax().to_string(),
123 )),
124 );
125 }
126
127 if let Some(group) = &select.group_by_clause() {
128 let mut group_doc = Doc::line_or_space().append(leading_comments(group.syntax()));
129 group_doc = group_doc.append(Doc::text("group")).append(Doc::space());
130 if let Some(by_token) = group.by_token() {
131 group_doc = group_doc.append(leading_comments_token(&by_token));
132 }
133 group_doc = group_doc.append(Doc::text("by")).append(Doc::space());
134 if let Some(list) = group.group_by_list() {
135 group_doc = group_doc.append(build_group_by_list(list));
136 }
137 doc = doc.append(group_doc);
138 }
139
140 doc = doc.append(build_semicolon(select.semicolon_token()));
141
142 doc.group()
143}
144
145fn build_group_by_list<'a>(list: ast::GroupByList) -> Doc<'a> {
146 leading_comments(list.syntax()).append(Doc::text(list.syntax().to_string()))
147}
148
149fn build_semicolon<'a>(semi: Option<SyntaxToken>) -> Doc<'a> {
150 let Some(semi) = semi else {
151 return Doc::nil();
152 };
153 let mut doc = Doc::nil();
154 let mut comments: Vec<SyntaxToken> = vec![];
155 for next in semi.siblings_with_tokens(Direction::Prev).skip(1) {
156 match next {
157 rowan::NodeOrToken::Node(_) => break,
158 rowan::NodeOrToken::Token(token) => {
159 if token.kind() == SyntaxKind::COMMENT {
160 comments.push(token);
161 } else if token.kind() == SyntaxKind::WHITESPACE {
162 continue;
163 } else {
164 break;
165 }
166 }
167 }
168 }
169 for comment in comments.iter().rev() {
170 doc = doc.append(Doc::text(comment.text().to_string()));
171 }
172 doc.append(Doc::text(";"))
173}
174
175fn build_expr<'a>(expr: ast::Expr) -> Doc<'a> {
176 match expr {
177 ast::Expr::ArrayExpr(array_expr) => {
178 let mut doc = Doc::nil();
179
180 if array_expr.array_token().is_some() {
182 doc = doc.append(Doc::text("array"));
183 };
184
185 if let Some(select) = array_expr.select() {
186 doc = doc
187 .append(Doc::text("("))
188 .append(build_select_doc(&select))
189 .append(Doc::text(")"))
190 } else {
191 doc = doc
192 .append(Doc::text("["))
193 .append(Doc::list(
194 Itertools::intersperse(
195 array_expr.exprs().map(build_expr),
196 Doc::text(",").append(Doc::space()),
197 )
198 .collect(),
199 ))
200 .append(Doc::text("]"));
201 }
202
203 doc
204 }
205 ast::Expr::BetweenExpr(between_expr) => {
206 let mut doc = build_expr(between_expr.target().unwrap());
207 if between_expr.not_token().is_some() {
208 doc = doc.append(Doc::space()).append(Doc::text("not"));
209 }
210 doc = doc.append(Doc::space()).append(Doc::text("between"));
211 if between_expr.asymmetric_token().is_some() {
212 doc = doc.append(Doc::space()).append(Doc::text("asymmetric"));
213 } else if between_expr.symmetric_token().is_some() {
214 doc = doc.append(Doc::space()).append(Doc::text("symmetric"));
215 }
216 doc.append(Doc::space())
217 .append(build_expr(between_expr.start().unwrap()))
218 .append(Doc::space())
219 .append(Doc::text("and"))
220 .append(Doc::space())
221 .append(build_expr(between_expr.end().unwrap()))
222 }
223 ast::Expr::BinExpr(bin_expr) => build_expr(bin_expr.lhs().unwrap())
224 .append(Doc::space())
225 .append(build_op(bin_expr.op().unwrap()))
226 .append(Doc::space())
227 .append(build_expr(bin_expr.rhs().unwrap())),
228 ast::Expr::CastExpr(cast_expr) => {
231 let mut doc = Doc::nil();
232 if cast_expr.colon_colon().is_some() {
233 doc = doc
234 .append(build_expr(cast_expr.expr().unwrap()))
235 .append(Doc::text("::"))
236 .append(build_type(cast_expr.ty().unwrap()))
237 } else if cast_expr.as_token().is_some() {
238 if cast_expr.cast_token().is_some() {
239 doc = doc.append(Doc::text("cast"))
240 } else if cast_expr.treat_token().is_some() {
241 doc = doc.append(Doc::text("treat"))
242 }
243 doc = doc
244 .append(Doc::text("("))
245 .append(build_expr(cast_expr.expr().unwrap()))
246 .append(Doc::space())
247 .append(Doc::text("as"))
248 .append(Doc::space())
249 .append(build_type(cast_expr.ty().unwrap()))
250 .append(Doc::text(")"))
251 } else {
252 doc = doc
253 .append(build_type(cast_expr.ty().unwrap()))
254 .append(Doc::space())
255 .append(build_literal(cast_expr.literal().unwrap()))
256 }
257 doc
258 }
259 ast::Expr::Collate(collate) => build_expr(collate.expr().unwrap())
260 .append(Doc::space())
261 .append(Doc::text("collate"))
262 .append(Doc::space())
263 .append(Doc::text(
264 collate.collation_ref().unwrap().syntax().to_string(),
265 )),
266 ast::Expr::Literal(literal) => build_literal(literal),
269 ast::Expr::PostfixExpr(postfix_expr) => {
272 let expr = build_expr(postfix_expr.expr().unwrap());
273 let op = match postfix_expr.op().unwrap() {
274 ast::PostfixOp::AtLocal(_) => Doc::text("at local"),
275 ast::PostfixOp::IsNull(_) => Doc::text("isnull"),
276 ast::PostfixOp::NotNull(_) => Doc::text("notnull"),
277 ast::PostfixOp::IsJson(n) => {
278 let mut doc = Doc::text("is json");
279 if let Some(clause) = n.json_keys_unique_clause() {
280 doc = doc
281 .append(Doc::space())
282 .append(build_json_keys_unique_clause(clause));
283 }
284 doc
285 }
286 ast::PostfixOp::IsJsonArray(n) => {
287 let mut doc = Doc::text("is json array");
288 if let Some(clause) = n.json_keys_unique_clause() {
289 doc = doc
290 .append(Doc::space())
291 .append(build_json_keys_unique_clause(clause));
292 }
293 doc
294 }
295 ast::PostfixOp::IsJsonObject(n) => {
296 let mut doc = Doc::text("is json object");
297 if let Some(clause) = n.json_keys_unique_clause() {
298 doc = doc
299 .append(Doc::space())
300 .append(build_json_keys_unique_clause(clause));
301 }
302 doc
303 }
304 ast::PostfixOp::IsJsonScalar(n) => {
305 let mut doc = Doc::text("is json scalar");
306 if let Some(clause) = n.json_keys_unique_clause() {
307 doc = doc
308 .append(Doc::space())
309 .append(build_json_keys_unique_clause(clause));
310 }
311 doc
312 }
313 ast::PostfixOp::IsJsonValue(n) => {
314 let mut doc = Doc::text("is json value");
315 if let Some(clause) = n.json_keys_unique_clause() {
316 doc = doc
317 .append(Doc::space())
318 .append(build_json_keys_unique_clause(clause));
319 }
320 doc
321 }
322 ast::PostfixOp::IsNormalized(n) => {
323 let mut doc = Doc::text("is");
324 if let Some(form) = n.unicode_normal_form() {
325 doc = doc
326 .append(Doc::space())
327 .append(build_unicode_normal_form(form));
328 }
329 doc.append(Doc::space()).append(Doc::text("normalized"))
330 }
331 ast::PostfixOp::IsNotJson(n) => {
332 let mut doc = Doc::text("is not json");
333 if let Some(clause) = n.json_keys_unique_clause() {
334 doc = doc
335 .append(Doc::space())
336 .append(build_json_keys_unique_clause(clause));
337 }
338 doc
339 }
340 ast::PostfixOp::IsNotJsonArray(n) => {
341 let mut doc = Doc::text("is not json array");
342 if let Some(clause) = n.json_keys_unique_clause() {
343 doc = doc
344 .append(Doc::space())
345 .append(build_json_keys_unique_clause(clause));
346 }
347 doc
348 }
349 ast::PostfixOp::IsNotJsonObject(n) => {
350 let mut doc = Doc::text("is not json object");
351 if let Some(clause) = n.json_keys_unique_clause() {
352 doc = doc
353 .append(Doc::space())
354 .append(build_json_keys_unique_clause(clause));
355 }
356 doc
357 }
358 ast::PostfixOp::IsNotJsonScalar(n) => {
359 let mut doc = Doc::text("is not json scalar");
360 if let Some(clause) = n.json_keys_unique_clause() {
361 doc = doc
362 .append(Doc::space())
363 .append(build_json_keys_unique_clause(clause));
364 }
365 doc
366 }
367 ast::PostfixOp::IsNotJsonValue(n) => {
368 let mut doc = Doc::text("is not json value");
369 if let Some(clause) = n.json_keys_unique_clause() {
370 doc = doc
371 .append(Doc::space())
372 .append(build_json_keys_unique_clause(clause));
373 }
374 doc
375 }
376 ast::PostfixOp::IsNotNormalized(n) => {
377 let mut doc = Doc::text("is not");
378 if let Some(form) = n.unicode_normal_form() {
379 doc = doc
380 .append(Doc::space())
381 .append(build_unicode_normal_form(form));
382 }
383 doc.append(Doc::space()).append(Doc::text("normalized"))
384 }
385 };
386 expr.append(Doc::space()).append(op)
387 }
388 _ => Doc::text(expr.syntax().to_string()),
392 }
393}
394
395fn build_json_keys_unique_clause<'a>(clause: ast::JsonKeysUniqueClause) -> Doc<'a> {
396 let prefix = match clause {
397 ast::JsonKeysUniqueClause::JsonWithoutUniqueKeys(_) => "without",
398 ast::JsonKeysUniqueClause::JsonWithUniqueKeys(_) => "with",
399 };
400 Doc::text(prefix)
401 .append(Doc::space())
402 .append(Doc::text("unique"))
403 .append(Doc::space())
404 .append(Doc::text("keys"))
405}
406
407fn build_unicode_normal_form<'a>(form: ast::UnicodeNormalForm) -> Doc<'a> {
408 if form.nfc_token().is_some() {
409 Doc::text("nfc")
410 } else if form.nfd_token().is_some() {
411 Doc::text("nfd")
412 } else if form.nfkc_token().is_some() {
413 Doc::text("nfkc")
414 } else {
415 Doc::text("nfkd")
416 }
417}
418
419fn build_keyword_node<'a>(node: &SyntaxNode) -> Doc<'a> {
420 let mut docs: Vec<Doc<'a>> = vec![];
421 for el in node.children_with_tokens() {
422 match el {
423 rowan::NodeOrToken::Token(token) => match token.kind() {
424 SyntaxKind::WHITESPACE => continue,
425 SyntaxKind::COMMENT => {
426 if !docs.is_empty() {
427 docs.push(Doc::space());
428 }
429 docs.push(Doc::text(token.text().to_string()));
430 }
431 _ => {
432 if !docs.is_empty() {
433 docs.push(Doc::space());
434 }
435 docs.push(Doc::text(token.text().to_ascii_lowercase()));
436 }
437 },
438 rowan::NodeOrToken::Node(_) => (),
439 }
440 }
441 Doc::list(docs)
442}
443
444fn build_op<'a>(op: ast::BinOp) -> Doc<'a> {
445 match op {
446 ast::BinOp::And(_) => Doc::text("and"),
447 ast::BinOp::AtTimeZone(n) => build_keyword_node(n.syntax()),
448 ast::BinOp::Caret(_) => Doc::text("^"),
449 ast::BinOp::ColonColon(_) => Doc::text("::"),
450 ast::BinOp::ColonEq(_) => Doc::text(":="),
451 ast::BinOp::CustomOp(custom_op) => Doc::text(custom_op.syntax().to_string()),
452 ast::BinOp::Eq(_) => Doc::text("="),
453 ast::BinOp::Escape(_) => Doc::text("escape"),
454 ast::BinOp::FatArrow(_) => Doc::text("=>"),
455 ast::BinOp::Gteq(_) => Doc::text(">="),
456 ast::BinOp::Ilike(_) => Doc::text("ilike"),
457 ast::BinOp::In(_) => Doc::text("in"),
458 ast::BinOp::Is(_) => Doc::text("is"),
459 ast::BinOp::IsDistinctFrom(n) => build_keyword_node(n.syntax()),
460 ast::BinOp::IsNot(n) => build_keyword_node(n.syntax()),
461 ast::BinOp::IsNotDistinctFrom(n) => build_keyword_node(n.syntax()),
462 ast::BinOp::LAngle(_) => Doc::text("<"),
463 ast::BinOp::Like(_) => Doc::text("like"),
464 ast::BinOp::Lteq(_) => Doc::text("<="),
465 ast::BinOp::Minus(_) => Doc::text("-"),
466 ast::BinOp::Neq(_) => Doc::text("!="),
467 ast::BinOp::Neqb(_) => Doc::text("<>"),
468 ast::BinOp::NotIlike(n) => build_keyword_node(n.syntax()),
469 ast::BinOp::NotIn(n) => build_keyword_node(n.syntax()),
470 ast::BinOp::NotLike(n) => build_keyword_node(n.syntax()),
471 ast::BinOp::NotSimilarTo(n) => build_keyword_node(n.syntax()),
472 ast::BinOp::OperatorCall(op) => Doc::text(op.syntax().to_string()),
473 ast::BinOp::Or(_) => Doc::text("or"),
474 ast::BinOp::Overlaps(_) => Doc::text("overlaps"),
475 ast::BinOp::Percent(_) => Doc::text("%"),
476 ast::BinOp::Plus(_) => Doc::text("+"),
477 ast::BinOp::RAngle(_) => Doc::text(">"),
478 ast::BinOp::SimilarTo(n) => build_keyword_node(n.syntax()),
479 ast::BinOp::Slash(_) => Doc::text("/"),
480 ast::BinOp::Star(_) => Doc::text("*"),
481 }
482}
483
484fn build_literal<'a>(lit: ast::Literal) -> Doc<'a> {
485 let Some(kind) = lit.kind() else {
486 return Doc::nil();
487 };
488 match kind {
489 LitKind::Default(_) => Doc::text("default"),
490 LitKind::False(_) => Doc::text("false"),
491 LitKind::IntNumber(t) => Doc::text(t.text().to_string()),
492 LitKind::Null(_) => Doc::text("null"),
493 LitKind::NumericNumber(t) => Doc::text(t.text().to_string()),
494 LitKind::PositionalParam(t) => Doc::text(t.text().to_string()),
495 LitKind::True(_) => Doc::text("true"),
496 LitKind::BitString(_)
497 | LitKind::ByteString(_)
498 | LitKind::DollarQuotedString(_)
499 | LitKind::EscString(_)
500 | LitKind::NationalString(_)
501 | LitKind::String(_)
502 | LitKind::UnicodeEscString(_) => build_string_literal(&lit),
503 }
504}
505
506fn build_string_literal<'a>(lit: &ast::Literal) -> Doc<'a> {
507 let parts: Vec<Doc<'a>> = lit
508 .syntax()
509 .children_with_tokens()
510 .filter_map(|el| match el {
511 rowan::NodeOrToken::Token(t) if t.kind() != SyntaxKind::WHITESPACE => {
512 Some(Doc::text(format_string_token(&t)))
513 }
514 _ => None,
515 })
516 .collect();
517 Doc::list(Itertools::intersperse(parts.into_iter(), Doc::hard_line()).collect())
518}
519
520fn format_string_token(t: &SyntaxToken) -> String {
521 let text = t.text();
522 if matches!(
523 t.kind(),
524 SyntaxKind::STRING | SyntaxKind::DOLLAR_QUOTED_STRING
525 ) {
526 return text.to_string();
527 }
528 match text.find('\'') {
529 Some(idx) => {
530 let (prefix, rest) = text.split_at(idx);
531 let mut s = String::with_capacity(text.len());
532 s.push_str(&prefix.to_ascii_lowercase());
533 s.push_str(rest);
534 s
535 }
536 None => text.to_string(),
537 }
538}
539
540fn build_column_name<'a>(name: ast::ColumnName) -> Doc<'a> {
541 Doc::text(quote_column_alias(&name.text()))
542}
543
544fn build_type<'a>(ty: ast::Type) -> Doc<'a> {
545 Doc::text(ty.syntax().to_string())
546}
547
548fn leading_comments_token<'a>(node: &SyntaxToken) -> Doc<'a> {
549 let mut doc = Doc::nil();
550 for next in node.siblings_with_tokens(Direction::Prev).skip(1) {
551 match next {
552 rowan::NodeOrToken::Node(_node) => {
553 break;
554 }
555 rowan::NodeOrToken::Token(token) => {
556 if token.kind() == SyntaxKind::COMMENT {
557 doc = doc
558 .append(Doc::text(token.text().to_string()))
559 .append(Doc::space());
560 } else if token.kind() == SyntaxKind::WHITESPACE {
561 continue;
562 } else {
563 break;
564 }
565 }
566 }
567 }
568 doc
569}
570
571fn leading_comments<'a>(node: &SyntaxNode) -> Doc<'a> {
572 let mut doc = Doc::nil();
573 for next in node.siblings_with_tokens(Direction::Prev).skip(1) {
574 match next {
575 rowan::NodeOrToken::Node(_node) => {
576 break;
577 }
578 rowan::NodeOrToken::Token(token) => {
579 if token.kind() == SyntaxKind::COMMENT {
580 let is_block = token.text().starts_with("--");
581 doc = doc
582 .append(Doc::text(token.text().to_string()))
583 .append(if is_block {
584 Doc::hard_line()
585 } else {
586 Doc::space()
587 });
588 } else if token.kind() == SyntaxKind::WHITESPACE {
589 continue;
590 } else {
591 break;
592 }
593 }
594 }
595 }
596 doc
597}
598
599fn trailing_comments<'a>(node: &SyntaxNode) -> Doc<'a> {
600 let mut doc = Doc::nil();
601 for next in node.siblings_with_tokens(Direction::Next).skip(1) {
602 match next {
603 rowan::NodeOrToken::Node(_node) => {
604 break;
605 }
606 rowan::NodeOrToken::Token(token) => {
607 if token.kind() == SyntaxKind::COMMENT {
608 doc = doc
609 .append(Doc::space())
610 .append(Doc::text(token.text().to_string()));
611 } else if token.kind() == SyntaxKind::WHITESPACE {
612 continue;
613 } else {
614 break;
615 }
616 }
617 }
618 }
619 doc
620}
621
622fn build_target<'a>(target: ast::Target) -> Option<Doc<'a>> {
623 let mut doc = leading_comments(target.syntax());
624
625 if target.star_token().is_some() {
626 return Some(doc.append(Doc::text("*")));
627 }
628 let expr = target.expr()?;
629 doc = doc.append(build_expr(expr));
630
631 if let Some(as_name) = target.as_name() {
632 if as_name.as_token().is_some() {
633 doc = doc.append(Doc::space()).append(Doc::text("as"))
634 }
635
636 if let Some(column_name) = as_name.name() {
637 doc = doc
638 .append(Doc::space())
639 .append(build_column_name(column_name));
640 }
641 }
642
643 doc = doc.append(trailing_comments(target.syntax()));
644
645 Some(doc)
646}
647
648pub fn fmt(text: &str) -> String {
649 let parse = ast::SourceFile::parse(text);
650 let file = parse.tree();
651 println!("{text}");
652 println!("---");
653 println!("{:#?}", file.syntax());
654 println!("---");
655 debug_assert_eq!(
656 parse.errors(),
657 vec![],
658 "should bail out when there's parse errors"
659 );
660 let doc = build_source_file(&file);
661 print(&doc, &PrintOptions::default())
662}