Add Comments::find_before

src/ast/comments.rs

@@ -101,6 +101,39 @@ impl Comments {
         })
     }
 
+    /// Finds comments immediately before the token starting at `location`, in source order.
+    ///
+    /// # Example
+    /// ```rust
+    /// use sqlparser::{dialect::GenericDialect, parser::Parser, tokenizer::Location};
+    ///
+    /// let sql = r#"SELECT 1; -- one
+    /// -- two
+    /// SELECT 2"#;
+    /// let (_, comments) = Parser::parse_sql_with_comments(&GenericDialect, sql).unwrap();
+    ///
+    /// let second = Location::new(3, 1);
+    /// assert_eq!(
+    ///    &comments.find_before(second).map(|c| c.as_str()).collect::<Vec<_>>(),
+    ///    &[" one", " two"]);
+    /// ```
+    pub fn find_before(&self, location: Location) -> Iter<'_> {
+        let end = self.end_index(Bound::Excluded(&location));
+        let start = self.0[..end].partition_point(|c| c.next_token_start != Some(location));
+        Iter(self.0[start..end].iter())
+    }
+
+    pub(crate) fn set_next_token_start(&mut self, next_token_start: Location) {
+        for comment in self
+            .0
+            .iter_mut()
+            .rev()
+            .take_while(|c| c.next_token_start.is_none())
+        {
+            comment.next_token_start = Some(next_token_start);
+        }
+    }
+
     /// Find the index of the first comment starting "before" the given location.
     ///
     /// The returned index is _inclusive_ and within the range of `0..=self.0.len()`.
@@ -157,6 +190,10 @@ pub struct CommentWithSpan {
     pub comment: Comment,
     /// The span of the comment including its markers
     pub span: Span,
+    /// End of the previous non-whitespace source token, or `None` before the first token.
+    pub prev_token_end: Option<Location>,
+    /// Start of the next non-whitespace source token, or `None` after the last token.
+    pub next_token_start: Option<Location>,
 }
 
 impl Deref for CommentWithSpan {
@@ -208,7 +245,7 @@ impl Deref for Comment {
     }
 }
 
-/// An opaque iterator implementation over comments served by [Comments::find].
+/// An opaque iterator over comments returned by [`Comments::find`] and [`Comments::find_before`].
 pub struct Iter<'a>(slice::Iter<'a, CommentWithSpan>);
 
 impl<'a> Iterator for Iter<'a> {
@@ -241,21 +278,29 @@ mod tests {
                     prefix: "--".into(),
                 },
                 span: Span::new((1, 1).into(), (1, 7).into()),
+                prev_token_end: None,
+                next_token_start: None,
             });
             c.offer(CommentWithSpan {
                 comment: Comment::MultiLine(" hello ".into()),
                 span: Span::new((2, 3).into(), (2, 14).into()),
+                prev_token_end: None,
+                next_token_start: None,
             });
             c.offer(CommentWithSpan {
                 comment: Comment::SingleLine {
                     content: ", world".into(),
                     prefix: "--".into(),
                 },
                 span: Span::new((2, 14).into(), (2, 21).into()),
+                prev_token_end: None,
+                next_token_start: None,
             });
             c.offer(CommentWithSpan {
                 comment: Comment::MultiLine(" def\n ghi\n jkl\n".into()),
                 span: Span::new((3, 3).into(), (7, 1).into()),
+                prev_token_end: None,
+                next_token_start: None,
             });
             c
         };
@@ -335,4 +380,79 @@ mod tests {
             vec![" abc", " hello ", ", world", " def\n ghi\n jkl\n"]
         );
     }
+
+    #[test]
+    fn test_find_before() {
+        let first_token = Location::new(2, 1);
+        let second_token = Location::new(4, 1);
+        let mut comments = Comments::default();
+        for comment in [
+            CommentWithSpan {
+                comment: Comment::SingleLine {
+                    content: " first".into(),
+                    prefix: "--".into(),
+                },
+                span: Span::new((1, 1).into(), (1, 9).into()),
+                prev_token_end: None,
+                next_token_start: Some(first_token),
+            },
+            CommentWithSpan {
+                comment: Comment::SingleLine {
+                    content: " second".into(),
+                    prefix: "--".into(),
+                },
+                span: Span::new((2, 11).into(), (2, 20).into()),
+                prev_token_end: Some(Location::new(2, 10)),
+                next_token_start: Some(second_token),
+            },
+            CommentWithSpan {
+                comment: Comment::SingleLine {
+                    content: " third".into(),
+                    prefix: "--".into(),
+                },
+                span: Span::new((3, 1).into(), (3, 9).into()),
+                prev_token_end: Some(Location::new(2, 10)),
+                next_token_start: Some(second_token),
+            },
+            CommentWithSpan {
+                comment: Comment::SingleLine {
+                    content: " tail".into(),
+                    prefix: "--".into(),
+                },
+                span: Span::new((4, 11).into(), (4, 18).into()),
+                prev_token_end: Some(Location::new(4, 10)),
+                next_token_start: None,
+            },
+        ] {
+            comments.offer(comment);
+        }
+
+        fn find_before(comments: &Comments, location: Location) -> Vec<&str> {
+            comments.find_before(location).map(|c| c.as_str()).collect()
+        }
+
+        assert_eq!(find_before(&comments, first_token), vec![" first"]);
+        assert_eq!(
+            find_before(&comments, second_token),
+            vec![" second", " third"]
+        );
+        for location in [
+            Location::new(0, 0),
+            Location::new(1, 3),
+            Location::new(2, 8),
+            Location::new(3, 12),
+            Location::new(4, 3),
+            Location::new(5, 1),
+        ] {
+            assert_eq!(
+                find_before(&comments, location),
+                Vec::<&str>::new(),
+                "{location:?}"
+            );
+        }
+        assert_eq!(
+            find_before(&Comments::default(), first_token),
+            Vec::<&str>::new()
+        );
+    }
 }

src/parser/mod.rs

@@ -596,25 +596,31 @@ impl<'a> Parser<'a> {
     /// Consumes this parser returning comments from the parsed token stream.
     pub fn into_comments(self) -> comments::Comments {
         let mut comments = comments::Comments::default();
+        let mut prev_token_end = None;
         for t in self.tokens.into_iter() {
-            match t.token {
+            let comment = match t.token {
                 Token::Whitespace(Whitespace::SingleLineComment { comment, prefix }) => {
-                    comments.offer(comments::CommentWithSpan {
-                        comment: comments::Comment::SingleLine {
-                            content: comment,
-                            prefix,
-                        },
-                        span: t.span,
-                    });
+                    comments::Comment::SingleLine {
+                        content: comment,
+                        prefix,
+                    }
                 }
                 Token::Whitespace(Whitespace::MultiLineComment(comment)) => {
-                    comments.offer(comments::CommentWithSpan {
-                        comment: comments::Comment::MultiLine(comment),
-                        span: t.span,
-                    });
+                    comments::Comment::MultiLine(comment)
                 }
-                _ => {}
-            }
+                Token::Whitespace(_) | Token::EOF => continue,
+                _ => {
+                    comments.set_next_token_start(t.span.start);
+                    prev_token_end = Some(t.span.end);
+                    continue;
+                }
+            };
+            comments.offer(comments::CommentWithSpan {
+                comment,
+                span: t.span,
+                prev_token_end,
+                next_token_start: None,
+            });
         }
         comments
     }

tests/sqlparser_comments.rs

@@ -25,7 +25,8 @@ use sqlparser::{
     ast::comments::{Comment, CommentWithSpan},
     dialect::GenericDialect,
     parser::Parser,
-    tokenizer::Span,
+    test_utils::all_dialects,
+    tokenizer::{Location, Span, Token, TokenWithSpan, Tokenizer},
 };
 
 #[test]
@@ -40,36 +41,142 @@ more*/
 
   -- end-of-script-with-no-newline"#;
 
-    let comments = match Parser::parse_sql_with_comments(&GenericDialect, sql) {
-        Ok((_, comments)) => comments,
-        Err(e) => panic!("Invalid sql script: {e}"),
-    };
+    let comments = all_dialects().one_of_identical_results(|dialect| {
+        Vec::from(Parser::parse_sql_with_comments(dialect, sql).unwrap().1)
+    });
 
     assert_eq!(
-        Vec::from(comments),
+        comments,
         vec![
             CommentWithSpan {
                 comment: Comment::SingleLine {
                     content: " second line comment".into(),
                     prefix: "--".into()
                 },
                 span: Span::new((2, 1).into(), (2, 23).into()),
+                prev_token_end: None,
+                next_token_start: Some(Location::new(3, 1)),
             },
             CommentWithSpan {
                 comment: Comment::MultiLine(" inline comment after `from` ".into()),
                 span: Span::new((3, 15).into(), (3, 48).into()),
+                prev_token_end: Some(Location::new(3, 14)),
+                next_token_start: Some(Location::new(3, 49)),
             },
             CommentWithSpan {
                 comment: Comment::MultiLine("select\nsome\nmore".into()),
-                span: Span::new((5, 1).into(), (7, 7).into())
+                span: Span::new((5, 1).into(), (7, 7).into()),
+                prev_token_end: Some(Location::new(3, 54)),
+                next_token_start: None,
             },
             CommentWithSpan {
                 comment: Comment::SingleLine {
                     content: " end-of-script-with-no-newline".into(),
                     prefix: "--".into()
                 },
                 span: Span::new((9, 3).into(), (9, 35).into()),
+                prev_token_end: Some(Location::new(3, 54)),
+                next_token_start: None,
             }
         ]
     );
 }
+
+#[test]
+fn parse_sql_with_comments_between_tokens() {
+    let sql = r#"SELECT 1; /* a */
+
+-- b
+/* c */ SELECT 2"#;
+    let comments = all_dialects().one_of_identical_results(|dialect| {
+        Vec::from(Parser::parse_sql_with_comments(dialect, sql).unwrap().1)
+    });
+
+    assert_eq!(
+        comments,
+        vec![
+            CommentWithSpan {
+                comment: Comment::MultiLine(" a ".into()),
+                span: Span::new((1, 11).into(), (1, 18).into()),
+                prev_token_end: Some(Location::new(1, 10)),
+                next_token_start: Some(Location::new(4, 9)),
+            },
+            CommentWithSpan {
+                comment: Comment::SingleLine {
+                    content: " b".into(),
+                    prefix: "--".into(),
+                },
+                span: Span::new((3, 1).into(), (3, 5).into()),
+                prev_token_end: Some(Location::new(1, 10)),
+                next_token_start: Some(Location::new(4, 9)),
+            },
+            CommentWithSpan {
+                comment: Comment::MultiLine(" c ".into()),
+                span: Span::new((4, 1).into(), (4, 8).into()),
+                prev_token_end: Some(Location::new(1, 10)),
+                next_token_start: Some(Location::new(4, 9)),
+            },
+        ]
+    );
+}
+
+#[test]
+fn parse_sql_with_comments_only() {
+    for (sql, expected) in [
+        ("", vec![]),
+        (
+            r#"-- a
+/* b */"#,
+            vec![
+                CommentWithSpan {
+                    comment: Comment::SingleLine {
+                        content: " a".into(),
+                        prefix: "--".into(),
+                    },
+                    span: Span::new((1, 1).into(), (1, 5).into()),
+                    prev_token_end: None,
+                    next_token_start: None,
+                },
+                CommentWithSpan {
+                    comment: Comment::MultiLine(" b ".into()),
+                    span: Span::new((2, 1).into(), (2, 8).into()),
+                    prev_token_end: None,
+                    next_token_start: None,
+                },
+            ],
+        ),
+    ] {
+        let comments = all_dialects().one_of_identical_results(|dialect| {
+            Vec::from(Parser::parse_sql_with_comments(dialect, sql).unwrap().1)
+        });
+        assert_eq!(comments, expected);
+    }
+}
+
+#[test]
+fn into_comments_with_eof() {
+    let mut tokens = Tokenizer::new(&GenericDialect, "SELECT 1; -- tail")
+        .tokenize_with_location()
+        .unwrap();
+    let eof = tokens.last().unwrap().span.end;
+    tokens.push(TokenWithSpan {
+        token: Token::EOF,
+        span: Span::new(eof, eof),
+    });
+    let comments = Parser::new(&GenericDialect)
+        .with_tokens_with_locations(tokens)
+        .into_comments();
+
+    assert_eq!(
+        Vec::from(comments),
+        vec![CommentWithSpan {
+            comment: Comment::SingleLine {
+                content: " tail".into(),
+                prefix: "--".into(),
+            },
+            span: Span::new((1, 11).into(), (1, 18).into()),
+            prev_token_end: Some(Location::new(1, 10)),
+            next_token_start: None,
+        }]
+    );
+}