fix(pretty_print): keep comments before CASE and GRANT keywords in place

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<String>,
 }
 
 /// 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<usize, String> {
+    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<SourceComment> {
     let Ok(scan) = pgls_query::scan(sql) else {
@@ -218,6 +264,7 @@ fn collect_comments(sql: &str) -> Vec<SourceComment> {
                     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";

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);

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();

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);
 

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);

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 {

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());

crates/pgls_pretty_print/src/renderer.rs

@@ -309,9 +309,14 @@ impl<W: Write> Renderer<W> {
         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,

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;"
+    );
+}