diff --git a/crates/pgls_pretty_print/src/comments.rs b/crates/pgls_pretty_print/src/comments.rs index bc792bb44..816d603b9 100644 --- a/crates/pgls_pretty_print/src/comments.rs +++ b/crates/pgls_pretty_print/src/comments.rs @@ -19,6 +19,13 @@ pub struct Comment { /// True for a line comment written after a list comma, `a, -- note`. It stays after the comma /// rather than moving onto the next item's line. pub after_separator: bool, + /// The keyword written right after the comment, in lower case, when no node starts there, + /// such as `else` in `then 'x' /* why */ else 'y'`. + /// + /// The comment is attached to the node after that keyword, but the parent emits the keyword + /// first. A parent that knows the keyword takes such comments before writing it, so that they + /// don't move past it. A parent that doesn't leaves them where they were. + pub before_keyword: Option, } /// Comments of a statement, indexed by the node they surround. @@ -58,6 +65,7 @@ pub fn attach_comments(sql: &str, ast: &NodeEnum) -> AttachedComments { let mut locations = collect_node_locations(ast); locations.sort_unstable(); + let keyword_after = keywords_after_comments(sql, &locations); for mut source_comment in comments { let line_start = sql[..source_comment.start] @@ -94,6 +102,10 @@ pub fn attach_comments(sql: &str, ast: &NodeEnum) -> AttachedComments { .rev() .find(|location| **location as usize <= source_comment.start) .copied(); + source_comment.comment.before_keyword = keyword_after.get(&source_comment.start).cloned(); + // Nothing inside the CASE follows its END, so a comment before END belongs to the last + // branch rather than to whatever follows the expression. + let leads = leads && source_comment.comment.before_keyword.as_deref() != Some("end"); // The preferred side first, the other one as a fallback. A comment closing a list or a // statement has no node after it, and printing it after the node it already follows in the @@ -147,6 +159,7 @@ fn ends_with_clause_header(line_prefix: &str) -> bool { "USING", "AND", "OR", + "CASE", "WHEN", "THEN", "ELSE", @@ -190,6 +203,39 @@ struct SourceComment { comment: Comment, } +/// For every comment followed by a keyword that no node starts at, the keyword in lower case, by +/// the comment's start. +fn keywords_after_comments(sql: &str, locations: &[i32]) -> HashMap { + use pgls_query::protobuf::KeywordKind; + + let Ok(scan) = pgls_query::scan(sql) else { + return HashMap::new(); + }; + + let mut keywords = HashMap::new(); + for pair in scan.tokens.windows(2) { + let [comment, next] = [&pair[0], &pair[1]]; + if !matches!( + Token::try_from(comment.token), + Ok(Token::SqlComment | Token::CComment) + ) || matches!( + Token::try_from(next.token), + Ok(Token::SqlComment | Token::CComment) + ) { + continue; + } + if next.keyword_kind == KeywordKind::NoKeyword as i32 + || locations.binary_search(&next.start).is_ok() + { + continue; + } + if let Some(keyword) = sql.get(next.start as usize..next.end as usize) { + keywords.insert(comment.start as usize, keyword.to_ascii_lowercase()); + } + } + keywords +} + /// Every comment of the statement, with its source range, in source order. fn collect_comments(sql: &str) -> Vec { let Ok(scan) = pgls_query::scan(sql) else { @@ -218,6 +264,7 @@ fn collect_comments(sql: &str) -> Vec { line_comment, own_line: false, after_separator: false, + before_keyword: None, }, }) }) @@ -384,6 +431,28 @@ mod tests { assert_eq!(attached.unattached[0].text, "-- trailing"); } + #[test] + fn a_comment_before_a_keyword_records_the_keyword() { + let sql = "SELECT CASE WHEN a THEN 'x' /* why */ ELSE 'y' END"; + let attached = attach_comments(sql, &parse(sql)); + + let comments = attached + .leading_by_location + .values() + .next() + .expect("one entry"); + assert_eq!(comments[0].before_keyword.as_deref(), Some("else")); + } + + #[test] + fn a_comment_before_end_stays_with_the_last_branch() { + let sql = "SELECT CASE WHEN a THEN 1 ELSE 2 /* last */ END FROM t"; + let attached = attach_comments(sql, &parse(sql)); + + assert!(attached.leading_by_location.is_empty()); + assert_eq!(attached.trailing_by_location.len(), 1); + } + #[test] fn a_statement_without_comments_produces_an_empty_map() { let sql = "SELECT 1 FROM s.t"; diff --git a/crates/pgls_pretty_print/src/emitter.rs b/crates/pgls_pretty_print/src/emitter.rs index 9d1c5f758..a32449652 100644 --- a/crates/pgls_pretty_print/src/emitter.rs +++ b/crates/pgls_pretty_print/src/emitter.rs @@ -227,6 +227,23 @@ impl EventEmitter { true } + /// Emits and consumes the comments before `location` that were written before `keyword`, + /// the way trailing comments of the previous node are emitted. + pub fn take_comments_before_keyword_at(&mut self, location: i32, keyword: &str) { + let Some(comments) = self.leading_comments.get_mut(&location) else { + return; + }; + let (before_keyword, rest): (Vec<_>, Vec<_>) = std::mem::take(comments) + .into_iter() + .partition(|comment| comment.before_keyword.as_deref() == Some(keyword)); + if rest.is_empty() { + self.leading_comments.remove(&location); + } else { + *comments = rest; + } + self.emit_trailing_comments(before_keyword); + } + /// Emits and consumes comments that follow `location`, if any. pub fn take_trailing_comments_at(&mut self, location: i32) { let comments = self.remove_trailing_comments_at(location); diff --git a/crates/pgls_pretty_print/src/nodes/case_expr.rs b/crates/pgls_pretty_print/src/nodes/case_expr.rs index 204649eca..b8ed85b53 100644 --- a/crates/pgls_pretty_print/src/nodes/case_expr.rs +++ b/crates/pgls_pretty_print/src/nodes/case_expr.rs @@ -31,6 +31,7 @@ pub(super) fn emit_case_expr(e: &mut EventEmitter, n: &CaseExpr) { // ELSE clause if let Some(ref defresult) = n.defresult { + super::emit_comments_before_keyword(e, "else", defresult); super::emit_layout_break(e); e.token(TokenKind::ELSE_KW); e.space(); diff --git a/crates/pgls_pretty_print/src/nodes/case_when.rs b/crates/pgls_pretty_print/src/nodes/case_when.rs index 2d2a45736..2ce18c85a 100644 --- a/crates/pgls_pretty_print/src/nodes/case_when.rs +++ b/crates/pgls_pretty_print/src/nodes/case_when.rs @@ -23,6 +23,9 @@ pub(super) fn emit_case_when(e: &mut EventEmitter, n: &CaseWhen) { if expanded { e.indent_start(); } + if let Some(ref result) = n.result { + super::emit_comments_before_keyword(e, "then", result); + } super::emit_layout_break(e); e.token(TokenKind::THEN_KW); diff --git a/crates/pgls_pretty_print/src/nodes/grant_role_stmt.rs b/crates/pgls_pretty_print/src/nodes/grant_role_stmt.rs index 26d6cf364..45f8530aa 100644 --- a/crates/pgls_pretty_print/src/nodes/grant_role_stmt.rs +++ b/crates/pgls_pretty_print/src/nodes/grant_role_stmt.rs @@ -41,6 +41,9 @@ pub(super) fn emit_grant_role_stmt(e: &mut EventEmitter, n: &GrantRoleStmt) { e.indent_start(); // TO or FROM + if let Some(grantee) = n.grantee_roles.first() { + super::emit_comments_before_keyword(e, if n.is_grant { "to" } else { "from" }, grantee); + } if n.is_grant { e.line(LineType::SoftOrSpace); e.token(TokenKind::TO_KW); diff --git a/crates/pgls_pretty_print/src/nodes/grant_stmt.rs b/crates/pgls_pretty_print/src/nodes/grant_stmt.rs index 8b7f31a42..79e31eed5 100644 --- a/crates/pgls_pretty_print/src/nodes/grant_stmt.rs +++ b/crates/pgls_pretty_print/src/nodes/grant_stmt.rs @@ -50,6 +50,9 @@ fn emit_grant_stmt_impl(e: &mut EventEmitter, n: &GrantStmt, with_semicolon: boo // The clauses after the privileges continue the statement, so a broken one indents them. e.indent_start(); + if let Some(object) = n.objects.first() { + super::emit_comments_before_keyword(e, "on", object); + } e.line(LineType::SoftOrSpace); e.token(TokenKind::ON_KW); e.space(); @@ -188,6 +191,9 @@ fn emit_grant_stmt_impl(e: &mut EventEmitter, n: &GrantStmt, with_semicolon: boo } // TO/FROM - allow line break before this clause + if let Some(grantee) = n.grantees.first() { + super::emit_comments_before_keyword(e, if n.is_grant { "to" } else { "from" }, grantee); + } e.line(LineType::SoftOrSpace); if n.is_grant { diff --git a/crates/pgls_pretty_print/src/nodes/mod.rs b/crates/pgls_pretty_print/src/nodes/mod.rs index 272349d44..72e986b2b 100644 --- a/crates/pgls_pretty_print/src/nodes/mod.rs +++ b/crates/pgls_pretty_print/src/nodes/mod.rs @@ -587,6 +587,21 @@ pub(super) fn emit_with_comments_at( e.take_trailing_comments_at(location); } +/// Emits the comments written before `keyword` whose node is `node`, the expression the parent +/// writes after the keyword. Call it right before the break that precedes the keyword. +/// +/// The comments are attached to `node` because the keyword has no node of its own. Without this, +/// `then 'x' /* why */ else 'y'` would come out as `else /* why */ 'y'`. +pub(super) fn emit_comments_before_keyword(e: &mut EventEmitter, keyword: &str, node: &Node) { + if let Some(location) = node + .node + .as_ref() + .and_then(crate::comments::first_node_location) + { + e.take_comments_before_keyword_at(location, keyword); + } +} + pub(super) fn emit_clause_condition(e: &mut EventEmitter, clause: &Node) { use crate::ClauseBodyStyle; use crate::emitter::LineType; @@ -971,6 +986,7 @@ mod tests { line_comment: true, own_line: true, after_separator: false, + before_keyword: None, }; let leading = HashMap::from([(7, vec![comment])]); let mut e = EventEmitter::with_comments(Default::default(), leading, HashMap::new()); diff --git a/crates/pgls_pretty_print/src/renderer.rs b/crates/pgls_pretty_print/src/renderer.rs index 907dc5755..551459905 100644 --- a/crates/pgls_pretty_print/src/renderer.rs +++ b/crates/pgls_pretty_print/src/renderer.rs @@ -309,9 +309,14 @@ impl Renderer { for event in group_events { match event { LayoutEvent::Token(token) => buffer.push_str(&token.render(&self.config)), + // Two spaces in a row collapse, as in `write_space`. LayoutEvent::Space | LayoutEvent::Line(LineType::SoftOrSpace) - | LayoutEvent::Line(LineType::Fill) => buffer.push(' '), + | LayoutEvent::Line(LineType::Fill) => { + if !buffer.ends_with(' ') { + buffer.push(' '); + } + } // Collapsing would push code behind a `--` comment or flatten a multi-line block // comment into the surrounding expression. LayoutEvent::Line(LineType::Hard) => return None, diff --git a/crates/pgls_pretty_print/tests/comment_layout.rs b/crates/pgls_pretty_print/tests/comment_layout.rs index a7e6cc334..0cc0cffe3 100644 --- a/crates/pgls_pretty_print/tests/comment_layout.rs +++ b/crates/pgls_pretty_print/tests/comment_layout.rs @@ -139,3 +139,71 @@ fn standalone_comment_stays_before_the_conjunction_it_documents() { ); } } + +fn format_both_layouts(sql: &str) -> [String; 2] { + let expanded = FormatConfig { + layout: Layout::Expanded, + ..Default::default() + }; + [ + format_idempotently(sql, &FormatConfig::default()), + format_idempotently(sql, &expanded), + ] +} + +#[test] +fn a_comment_before_else_stays_before_else() { + for output in + format_both_layouts("select case when a then 'x' /* inline */ else 'y' end from t") + { + assert!(output.contains("'x' /* inline */"), "{output}"); + } + for output in + format_both_layouts("select case when a then 'x'\n-- before else\nelse 'y' end from t") + { + assert!(output.contains("-- before else\n"), "{output}"); + let comment = output.find("-- before else").unwrap(); + assert!(comment < output.find("else 'y'").unwrap(), "{output}"); + } +} + +#[test] +fn a_comment_before_then_stays_before_then() { + for output in format_both_layouts("select case when a /* cond */ then 1 else 2 end from t") { + assert!(output.contains("a /* cond */"), "{output}"); + assert!(!output.contains("then /* cond */"), "{output}"); + } +} + +#[test] +fn a_comment_before_end_stays_inside_the_case() { + for output in format_both_layouts("select case when a then 1 else 2 /* last */ end from t") { + assert!(output.contains("2 /* last */"), "{output}"); + assert!(!output.contains("/* last */ t"), "{output}"); + } +} + +#[test] +fn a_comment_after_case_stays_before_the_first_when() { + for output in format_both_layouts("select case -- why\nwhen a then 'x' else 'y' end from t") { + let comment = output.find("-- why").unwrap(); + assert!(comment < output.find("when a").unwrap(), "{output}"); + assert!(!output.contains("'y' -- why"), "{output}"); + } +} + +#[test] +fn a_comment_before_on_in_a_grant_stays_before_on() { + assert_eq!( + format_idempotently("grant select -- c\non t to u", &FormatConfig::default()), + "grant select -- c\n on table t\n to u;" + ); + assert_eq!( + format_idempotently("grant select on t /* c */ to u", &FormatConfig::default()), + "grant select on table t /* c */ to u;" + ); + assert_eq!( + format_idempotently("grant r1 /* c */ to u", &FormatConfig::default()), + "grant r1 /* c */ to u;" + ); +}