fix: error recovery for extra comma with distinct (#25)

Author: tanzaku

crates/postgresql-cst-parser/src/cst/lr_parse_state.rs

@@ -47,4 +47,8 @@ impl<'a> LRParseState<'a> {
     pub fn last_syntax_kind(&self) -> Option<SyntaxKind> {
         self.stack.last().map(|(_, node)| node.into())
     }
+
+    pub fn iter_syntax_kind_rev(&'a self) -> impl Iterator<Item = SyntaxKind> + 'a {
+        self.stack.iter().rev().map(|(_, node)| node.into())
+    }
 }

crates/postgresql-cst-parser/src/lib.rs

@@ -67,6 +67,30 @@ mod tests {
         let s = r#"select /*param*/ as A
 from /*#foo*/
 ;
+"#;
+        assert!(parse_2way(s).is_ok());
+    }
+
+    #[test]
+    fn test_2way_error_recovery_distinct_on() {
+        let s = r#"select distinct on (col1) , t.* from tbl t;
+;
+"#;
+        assert!(parse_2way(s).is_ok());
+    }
+
+    #[test]
+    fn test_2way_error_recovery_distinct() {
+        let s = r#"select distinct , t.* from tbl t;
+;
+"#;
+        assert!(parse_2way(s).is_ok());
+    }
+
+    #[test]
+    fn test_2way_error_recovery_all() {
+        let s = r#"select all , t.* from tbl t;
+;
 "#;
         assert!(parse_2way(s).is_ok());
     }

crates/postgresql-cst-parser/src/transform/skip_extra_comma.rs

@@ -8,7 +8,30 @@ pub struct SkipExtraComma;
 impl SkipExtraComma {
     fn allow_extra_comma<'a>(lr_parse_state: &LRParseState<'a>) -> bool {
         match lr_parse_state.last_syntax_kind() {
-            Some(SyntaxKind::FROM) | Some(SyntaxKind::BY) | Some(SyntaxKind::SELECT) => true,
+            Some(SyntaxKind::FROM)
+            | Some(SyntaxKind::BY)
+            | Some(SyntaxKind::SELECT)
+            | Some(SyntaxKind::ALL)
+            | Some(SyntaxKind::DISTINCT) => true,
+
+            Some(SyntaxKind::RParen) => {
+                // select distinct on
+                let expected_select_distinct_on_kinds = [
+                    SyntaxKind::RParen,
+                    SyntaxKind::expr_list,
+                    SyntaxKind::LParen,
+                    SyntaxKind::ON,
+                    SyntaxKind::DISTINCT,
+                    SyntaxKind::SELECT,
+                ];
+
+                let actual_kinds: Vec<_> = lr_parse_state
+                    .iter_syntax_kind_rev()
+                    .take(expected_select_distinct_on_kinds.len())
+                    .collect();
+
+                expected_select_distinct_on_kinds.as_ref() == actual_kinds.as_slice()
+            }
             _ => false,
         }
     }