Add APIs to reuse token buffers in Tokenizer (apache#1094)

Author: 0rphon

src/parser/mod.rs

@@ -8711,6 +8711,11 @@ impl<'a> Parser<'a> {
         self.expect_token(&Token::RParen)?;
         Ok(partitions)
     }
+
+    /// Consume the parser and return its underlying token buffer
+    pub fn into_tokens(self) -> Vec<TokenWithLocation> {
+        self.tokens
+    }
 }
 
 impl Word {

src/tokenizer.rs

@@ -543,21 +543,30 @@ 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<TokenWithLocation>, TokenizerError> {
+        let mut tokens: Vec<TokenWithLocation> = vec![];
+        self.tokenize_with_location_into_buf(&mut tokens)
+            .map(|_| tokens)
+    }
+
+    /// Tokenize the statement and append tokens with location information into the provided buffer.
+    /// If an error is thrown, the buffer will contain all tokens that were successfully parsed before the error.
+    pub fn tokenize_with_location_into_buf(
+        &mut self,
+        buf: &mut Vec<TokenWithLocation>,
+    ) -> Result<(), TokenizerError> {
         let mut state = State {
             peekable: self.query.chars().peekable(),
             line: 1,
             col: 1,
         };
 
-        let mut tokens: Vec<TokenWithLocation> = vec![];
-
         let mut location = state.location();
         while let Some(token) = self.next_token(&mut state)? {
-            tokens.push(TokenWithLocation { token, location });
+            buf.push(TokenWithLocation { token, location });
 
             location = state.location();
         }
-        Ok(tokens)
+        Ok(())
     }
 
     // Tokenize the identifer or keywords in `ch`

tests/sqlparser_common.rs

@@ -31,6 +31,7 @@ use sqlparser::dialect::{
 };
 use sqlparser::keywords::ALL_KEYWORDS;
 use sqlparser::parser::{Parser, ParserError, ParserOptions};
+use sqlparser::tokenizer::Tokenizer;
 use test_utils::{
     all_dialects, alter_table_op, assert_eq_vec, expr_from_projection, join, number, only, table,
     table_alias, TestedDialects,
@@ -8080,3 +8081,16 @@ fn test_release_savepoint() {
 
     one_statement_parses_to("RELEASE test1", "RELEASE SAVEPOINT test1");
 }
+
+#[test]
+fn test_buffer_reuse() {
+    let d = GenericDialect {};
+    let q = "INSERT INTO customer WITH foo AS (SELECT 1) SELECT * FROM foo UNION VALUES (1)";
+    let mut buf = Vec::new();
+    Tokenizer::new(&d, q)
+        .tokenize_with_location_into_buf(&mut buf)
+        .unwrap();
+    let mut p = Parser::new(&d).with_tokens_with_locations(buf);
+    p.parse_statements().unwrap();
+    let _ = p.into_tokens();
+}