Properly skip comments when parsing (#332)

Author: jpschorr

partiql-parser/src/lexer.rs

@@ -830,6 +830,46 @@ impl<'input> fmt::Display for Token<'input> {
     }
 }
 
+/// A lexer that wraps another lexer and skips comments.
+pub(crate) struct CommentSkippingLexer<'input, L>
+where
+    L: Iterator<Item = LexResult<'input>>,
+{
+    lexer: L,
+}
+
+impl<'input, L> CommentSkippingLexer<'input, L>
+where
+    L: Iterator<Item = LexResult<'input>>,
+{
+    /// Creates a new CommentSkippingLexer wrapping `lexer`
+    #[inline]
+    pub fn new(lexer: L) -> Self {
+        Self { lexer }
+    }
+}
+
+impl<'input, L> Iterator for CommentSkippingLexer<'input, L>
+where
+    L: Iterator<Item = LexResult<'input>>,
+{
+    type Item = LexResult<'input>;
+
+    #[inline(always)]
+    fn next(&mut self) -> Option<Self::Item> {
+        'next_tok: loop {
+            let next = self.lexer.next();
+            if matches!(
+                next,
+                Some(Ok((_, Token::CommentBlock(_) | Token::CommentLine(_), _)))
+            ) {
+                continue 'next_tok;
+            }
+            return next;
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

partiql-parser/src/parse/mod.rs

@@ -7,6 +7,7 @@ mod parser_state;
 
 use crate::error::{ParseError, UnexpectedTokenData};
 use crate::lexer;
+use crate::lexer::CommentSkippingLexer;
 use crate::parse::parser_state::{IdGenerator, ParserState};
 use crate::preprocessor::{PreprocessingPartiqlLexer, BUILT_INS};
 use lalrpop_util as lpop;
@@ -63,6 +64,7 @@ fn parse_partiql_with_state<'input, Id: IdGenerator>(
 ) -> AstResult<'input> {
     let mut offsets = LineOffsetTracker::default();
     let lexer = PreprocessingPartiqlLexer::new(s, &mut offsets, &BUILT_INS);
+    let lexer = CommentSkippingLexer::new(lexer);
 
     let result: LalrpopResult = grammar::QueryParser::new().parse(s, &mut state, lexer);
 
@@ -520,6 +522,15 @@ mod tests {
         fn select_with_at_and_cross_join_and_at() {
             parse!(r#"SELECT * FROM a AS a AT b CROSS JOIN c AS c AT q"#);
         }
+
+        #[test]
+        fn multiline_with_comments() {
+            parse!(
+                r#"SELECT * FROM hr.employees               -- T1
+                                  UNION
+                                  SELECT title FROM engineering.employees  -- T2"#
+            );
+        }
     }
 
     mod set_ops {