From 681f65b2d018c64abe6d846bcfce10466998fe33 Mon Sep 17 00:00:00 2001
From: "xudong.w" <wxd963996380@gmail.com>
Date: Sun, 15 Jan 2023 16:06:38 +0800
Subject: [PATCH] Fix parser with # (#147)

Co-authored-by: xxchan <xxchan22f@gmail.com>
---
 sqllogictest-bin/src/engines/external.rs |  2 +-
 sqllogictest/src/parser.rs               | 38 +++++++++++++++++-------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/sqllogictest-bin/src/engines/external.rs b/sqllogictest-bin/src/engines/external.rs
index bb9c6b6..7f0c843 100644
--- a/sqllogictest-bin/src/engines/external.rs
+++ b/sqllogictest-bin/src/engines/external.rs
@@ -104,7 +104,7 @@ impl AsyncDB for ExternalDriver {
         };
         match output {
             Output::Success { result } => Ok(DBOutput::Rows {
-                types: vec![], // FIXME: Fix it after https://github.com/risinglightdb/sqllogictest-rs/issues/36 is resolved.
+                types: vec![], /* FIXME: Fix it after https://github.com/risinglightdb/sqllogictest-rs/issues/36 is resolved. */
                 rows: result,
             }),
             Output::Failed { err } => Err(ExternalDriverError::Sql(err)),
diff --git a/sqllogictest/src/parser.rs b/sqllogictest/src/parser.rs
index 1003743..596db88 100644
--- a/sqllogictest/src/parser.rs
+++ b/sqllogictest/src/parser.rs
@@ -397,24 +397,24 @@ pub fn parse(script: &str) -> Result<Vec<Record>, ParseError> {
 
 #[allow(clippy::collapsible_match)]
 fn parse_inner(loc: &Location, script: &str) -> Result<Vec<Record>, ParseError> {
-    let mut lines = script.lines().enumerate();
+    let mut lines = script.lines().enumerate().peekable();
     let mut records = vec![];
     let mut conditions = vec![];
+    let mut comments = vec![];
 
-    while let Some((mut num, mut line)) = lines.next() {
+    while let Some((num, line)) = lines.next() {
         if let Some(text) = line.strip_prefix('#') {
-            let mut comments = vec![text.to_string()];
-            for (num_, line_) in lines.by_ref() {
-                num = num_;
-                line = line_;
-                if let Some(text) = line.strip_prefix('#') {
-                    comments.push(text.to_string());
-                } else {
-                    break;
-                }
+            comments.push(text.to_string());
+            if lines.peek().is_none() {
+                records.push(Record::Comment(comments));
+                comments = vec![];
             }
+            continue;
+        }
 
+        if !comments.is_empty() {
             records.push(Record::Comment(comments));
+            comments = vec![];
         }
 
         if line.is_empty() {
@@ -635,6 +635,22 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn test_trailing_comment() {
+        let script = "\
+# comment 1
+#  comment 2
+";
+        let records = parse(script).unwrap();
+        assert_eq!(
+            records,
+            vec![Record::Comment(vec![
+                " comment 1".to_string(),
+                "  comment 2".to_string(),
+            ]),]
+        );
+    }
+
     #[test]
     fn test_include_glob() {
         let records = parse_file("../examples/include/include_1.slt").unwrap();