diff --git a/src/ast/comments.rs b/src/ast/comments.rs index a0c25ad162..6b562b10d3 100644 --- a/src/ast/comments.rs +++ b/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::>(), + /// &[" 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, + /// Start of the next non-whitespace source token, or `None` after the last token. + pub next_token_start: Option, } 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,10 +278,14 @@ 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 { @@ -252,10 +293,14 @@ mod tests { 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() + ); + } } diff --git a/src/parser/mod.rs b/src/parser/mod.rs index db54296fce..82159a47b0 100644 --- a/src/parser/mod.rs +++ b/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 } diff --git a/tests/sqlparser_comments.rs b/tests/sqlparser_comments.rs index 8b41643870..ed78a77f8e 100644 --- a/tests/sqlparser_comments.rs +++ b/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,13 +41,12 @@ 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 { @@ -54,14 +54,20 @@ more*/ 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 { @@ -69,7 +75,108 @@ more*/ 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, + }] + ); +}