Unify numeric separator validation

Author: fyrsta7

src/tokenizer.rs

@@ -1360,36 +1360,12 @@ impl<'a> Tokenizer<'a> {
                         );
                     }
 
-                    // Some dialects support underscore as number separator
-                    // There can only be one at a time and it must be followed by another digit
-                    let is_number_separator = |ch: char, next_char: Option<char>| {
-                        self.dialect.supports_numeric_literal_underscores()
-                            && ch == '_'
-                            && next_char.is_some_and(|next_ch| next_ch.is_ascii_hexdigit())
-                    };
-
-                    let mut s = peeking_next_take_while(chars, |ch, next_ch| {
-                        ch.is_ascii_digit() || is_number_separator(ch, next_ch)
-                    });
-
-                    if self.dialect.supports_numeric_literal_underscores()
-                        && chars.peek() == Some(&'_')
-                    {
-                        return self.tokenizer_error(chars.location(), "Unexpected character '_'");
-                    }
+                    let mut s = self.tokenize_number_part(chars, |ch| ch.is_ascii_digit())?;
 
                     // match binary literal that starts with 0x
                     if s == "0" && chars.peek() == Some(&'x') {
                         chars.next();
-                        let s2 = peeking_next_take_while(chars, |ch, next_ch| {
-                            ch.is_ascii_hexdigit() || is_number_separator(ch, next_ch)
-                        });
-                        if self.dialect.supports_numeric_literal_underscores()
-                            && chars.peek() == Some(&'_')
-                        {
-                            return self
-                                .tokenizer_error(chars.location(), "Unexpected character '_'");
-                        }
+                        let s2 = self.tokenize_number_part(chars, |ch| ch.is_ascii_hexdigit())?;
                         return Ok(Some(Token::HexStringLiteral(s2)));
                     }
 
@@ -1410,23 +1386,8 @@ impl<'a> Tokenizer<'a> {
                         }
                     }
 
-                    if s != "."
-                        && self.dialect.supports_numeric_literal_underscores()
-                        && chars.peek() == Some(&'_')
-                    {
-                        return self.tokenizer_error(chars.location(), "Unexpected character '_'");
-                    }
-
                     // Consume fractional digits.
-                    s += &peeking_next_take_while(chars, |ch, next_ch| {
-                        ch.is_ascii_digit() || is_number_separator(ch, next_ch)
-                    });
-
-                    if self.dialect.supports_numeric_literal_underscores()
-                        && chars.peek() == Some(&'_')
-                    {
-                        return self.tokenizer_error(chars.location(), "Unexpected character '_'");
-                    }
+                    s += &self.tokenize_number_part(chars, |ch| ch.is_ascii_digit())?;
 
                     // No fraction -> Token::Period
                     if s == "." {
@@ -1450,12 +1411,16 @@ impl<'a> Tokenizer<'a> {
 
                         match char_clone.peek() {
                             // Definitely an exponent, get original iterator up to speed and use it
-                            Some(&c) if c.is_ascii_digit() => {
+                            Some(&c)
+                                if c.is_ascii_digit()
+                                    || (c == '_'
+                                        && self.dialect.supports_numeric_literal_underscores()) =>
+                            {
                                 for _ in 0..exponent_part.len() {
                                     chars.next();
                                 }
                                 exponent_part +=
-                                    &peeking_take_while(chars, |ch| ch.is_ascii_digit());
+                                    &self.tokenize_number_part(chars, |ch| ch.is_ascii_digit())?;
                                 s += exponent_part.as_str();
                             }
                             // Not an exponent, discard the work done
@@ -2060,6 +2025,33 @@ impl<'a> Tokenizer<'a> {
         })
     }
 
+    fn tokenize_number_part(
+        &self,
+        chars: &mut State,
+        is_digit: impl Fn(char) -> bool,
+    ) -> Result<String, TokenizerError> {
+        let supports_separator = self.dialect.supports_numeric_literal_underscores();
+        let mut s = String::new();
+
+        while let Some(&ch) = chars.peek() {
+            if is_digit(ch) {
+                chars.next();
+                s.push(ch);
+            } else if supports_separator && ch == '_' {
+                let next_char = chars.peekable.clone().nth(1);
+                if s.is_empty() || !next_char.is_some_and(&is_digit) {
+                    return self.tokenizer_error(chars.location(), "Unexpected character '_'");
+                }
+                chars.next();
+                s.push(ch);
+            } else {
+                break;
+            }
+        }
+
+        Ok(s)
+    }
+
     // Consume characters until newline
     fn tokenize_single_line_comment(&self, chars: &mut State) -> String {
         peeking_take_while(chars, |ch| match ch {
@@ -2436,24 +2428,6 @@ fn peeking_take_while(chars: &mut State, mut predicate: impl FnMut(char) -> bool
     s
 }
 
-/// Same as peeking_take_while, but also passes the next character to the predicate.
-fn peeking_next_take_while(
-    chars: &mut State,
-    mut predicate: impl FnMut(char, Option<char>) -> bool,
-) -> String {
-    let mut s = String::new();
-    while let Some(&ch) = chars.peek() {
-        let next_char = chars.peekable.clone().nth(1);
-        if predicate(ch, next_char) {
-            chars.next(); // consume
-            s.push(ch);
-        } else {
-            break;
-        }
-    }
-    s
-}
-
 fn unescape_single_quoted_string(chars: &mut State<'_>) -> Option<String> {
     Unescape::new(chars).unescape()
 }
@@ -2785,7 +2759,7 @@ mod tests {
             all_dialects_where(|dialect| dialect.supports_numeric_literal_underscores());
 
         numeric_underscore_dialects.tokenizes_to(
-            "SELECT 10_000, _10_000, 1_000.123, 1_000.123_456",
+            "SELECT 10_000, _10_000, 1_000.123, 1_000.123_456, 1e1_0",
             vec![
                 Token::make_keyword("SELECT"),
                 Token::Whitespace(Whitespace::Space),
@@ -2799,15 +2773,29 @@ mod tests {
                 Token::Comma,
                 Token::Whitespace(Whitespace::Space),
                 Token::Number("1_000.123_456".to_string(), false), // with an underscore in the decimal digits
+                Token::Comma,
+                Token::Whitespace(Whitespace::Space),
+                Token::Number("1e1_0".to_string(), false), // with an underscore in the exponent
             ],
         );
 
+        numeric_underscore_dialects.tokenizes_to(
+            "0xFF_FF",
+            vec![Token::HexStringLiteral("FF_FF".to_string())],
+        );
+
         for dialect in &numeric_underscore_dialects.dialects {
             for sql in [
                 "SELECT 10_00_",
                 "SELECT 10___0",
                 "SELECT 1_000.123_",
                 "SELECT 1._000",
+                "SELECT 1_a",
+                "SELECT 1e_1",
+                "SELECT 1e1_",
+                "SELECT 1e1__0",
+                "SELECT 0x_1",
+                "SELECT 0x1_",
             ] {
                 let err = Tokenizer::new(&**dialect, sql).tokenize().unwrap_err();
                 assert_eq!("Unexpected character '_'", err.message);

tests/sqlparser_common.rs

@@ -79,6 +79,9 @@ fn parse_numeric_literal_underscore() {
         ("SELECT 10__00", 10),
         ("SELECT 10_00_", 13),
         ("SELECT 1._000", 10),
+        ("SELECT 1_a", 9),
+        ("SELECT 1e_1", 10),
+        ("SELECT 1e1_", 11),
     ] {
         let err = dialects.parse_sql_statements(sql).unwrap_err();
         assert_eq!(