add ability to map tokens during tokenization

This patch adds a method to map tokens with provided mapper during
tokenization. This way tokens could be replaced without an additional
pass.

Author: Альберт Скальт

src/tokenizer.rs

@@ -925,7 +925,18 @@ impl<'a> Tokenizer<'a> {
     /// Tokenize the statement and produce a vector of tokens with location information
     pub fn tokenize_with_location(&mut self) -> Result<Vec<TokenWithSpan>, TokenizerError> {
         let mut tokens: Vec<TokenWithSpan> = vec![];
-        self.tokenize_with_location_into_buf(&mut tokens)
+        self.tokenize_with_location_into_buf(&mut tokens, |token| token)
+            .map(|_| tokens)
+    }
+
+    /// Tokenize the statement and produce a vector of tokens, mapping each token
+    /// with provided `mapper`
+    pub fn tokenize_map(
+        &mut self,
+        mapper: impl FnMut(TokenWithSpan) -> TokenWithSpan,
+    ) -> Result<Vec<TokenWithSpan>, TokenizerError> {
+        let mut tokens: Vec<TokenWithSpan> = vec![];
+        self.tokenize_with_location_into_buf(&mut tokens, mapper)
             .map(|_| tokens)
     }
 
@@ -934,6 +945,7 @@ impl<'a> Tokenizer<'a> {
     pub fn tokenize_with_location_into_buf(
         &mut self,
         buf: &mut Vec<TokenWithSpan>,
+        mut mapper: impl FnMut(TokenWithSpan) -> TokenWithSpan,
     ) -> Result<(), TokenizerError> {
         let mut state = State {
             peekable: self.query.chars().peekable(),
@@ -945,7 +957,8 @@ impl<'a> Tokenizer<'a> {
         while let Some(token) = self.next_token(&mut state, buf.last().map(|t| &t.token))? {
             let span = location.span_to(state.location());
 
-            buf.push(TokenWithSpan { token, span });
+            let token = TokenWithSpan { token, span };
+            buf.push(mapper(token));
 
             location = state.location();
         }
@@ -2590,6 +2603,40 @@ mod tests {
         compare(expected, tokens);
     }
 
+    #[test]
+    fn tokenize_with_mapper() {
+        let sql = String::from("SELECT ?");
+        let dialect = GenericDialect {};
+        let mut param_num = 1;
+
+        let tokens = Tokenizer::new(&dialect, &sql)
+            .tokenize_map(|mut token_span| {
+                token_span.token = match token_span.token {
+                    Token::Placeholder(n) => Token::Placeholder(if n == "?" {
+                        let ret = format!("${}", param_num);
+                        param_num += 1;
+                        ret
+                    } else {
+                        n
+                    }),
+                    token => token,
+                };
+                token_span
+            })
+            .unwrap()
+            .into_iter()
+            .map(|t| t.token)
+            .collect();
+
+        let expected = vec![
+            Token::make_keyword("SELECT"),
+            Token::Whitespace(Whitespace::Space),
+            Token::Placeholder("$1".to_string()),
+        ];
+
+        compare(expected, tokens);
+    }
+
     #[test]
     fn tokenize_clickhouse_double_equal() {
         let sql = String::from("SELECT foo=='1'");