Parser: fix exponential parse time on speculative prefix parsing

Author: LucaCappelletti94

sqlparser_bench/benches/sqlparser_bench.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use criterion::{criterion_group, criterion_main, Criterion};
-use sqlparser::dialect::{GenericDialect, PostgreSqlDialect};
+use sqlparser::dialect::{GenericDialect, PostgreSqlDialect, SQLiteDialect};
 use sqlparser::keywords::Keyword;
 use sqlparser::parser::Parser;
 use sqlparser::tokenizer::{Span, Word};
@@ -249,6 +249,51 @@ fn parse_not_chain(c: &mut Criterion) {
     group.finish();
 }
 
+/// Benchmark parsing pathological `IF(<keyword-fn>(<keyword-fn>(...x` chains
+/// that previously caused 2^N work in `parse_prefix`. Each nested
+/// `current_time(` segment used to be explored twice at every level (once via
+/// the speculative reserved-word arm, once via the unreserved-word fallback),
+/// doubling work per level. Post-fix the cost is linear in chain length.
+fn parse_prefix_keyword_call_chain(c: &mut Criterion) {
+    let mut group = c.benchmark_group("parse_prefix_keyword_call_chain");
+    let dialect = PostgreSqlDialect {};
+
+    for &n in &[10usize, 20, 30] {
+        let sql = String::from("if(") + &"current_time(".repeat(n) + "x";
+
+        group.bench_function(format!("chain_{n}"), |b| {
+            b.iter(|| {
+                let _ = Parser::parse_sql(&dialect, std::hint::black_box(&sql));
+            });
+        });
+    }
+
+    group.finish();
+}
+
+/// Benchmark parsing pathological `case-case-case-...c` chains that
+/// previously caused 2^N work in `parse_prefix`. Each `case` token used to
+/// trigger a speculative `parse_case_expr` that recursively descends the
+/// chain, but the unreserved-word fallback returns `Identifier(case)` so the
+/// overall `parse_prefix` succeeds and the failure cache never fires.
+/// Post-fix the per-arm cache short-circuits the speculative descent.
+fn parse_prefix_case_chain(c: &mut Criterion) {
+    let mut group = c.benchmark_group("parse_prefix_case_chain");
+    let dialect = SQLiteDialect {};
+
+    for &n in &[10usize, 20, 30] {
+        let sql = "case\t-".repeat(n) + "c";
+
+        group.bench_function(format!("chain_{n}"), |b| {
+            b.iter(|| {
+                let _ = Parser::parse_sql(&dialect, std::hint::black_box(&sql));
+            });
+        });
+    }
+
+    group.finish();
+}
+
 criterion_group!(
     benches,
     basic_queries,
@@ -257,6 +302,8 @@ criterion_group!(
     parse_compound_chain,
     parse_named_arg_chain,
     parse_compound_keyword_chain,
-    parse_not_chain
+    parse_not_chain,
+    parse_prefix_keyword_call_chain,
+    parse_prefix_case_chain
 );
 criterion_main!(benches);

src/parser/mod.rs

@@ -17,6 +17,7 @@ use alloc::collections::BTreeSet;
 #[cfg(not(feature = "std"))]
 use alloc::{
     boxed::Box,
+    collections::BTreeMap,
     format,
     string::{String, ToString},
     vec,
@@ -26,6 +27,9 @@ use core::{
     fmt::{self, Display},
     str::FromStr,
 };
+#[cfg(feature = "std")]
+use std::collections::BTreeMap;
+
 use helpers::attached_token::AttachedToken;
 #[cfg(feature = "std")]
 use std::collections::BTreeSet;
@@ -371,6 +375,12 @@ pub struct Parser<'a> {
     /// `<ident>-NOT-<ident>.` ending in a parse error) trigger 2^N exploration
     /// because each `-NOT-` segment otherwise re-walks the rest of the chain.
     failed_unary_not_positions: BTreeSet<usize>,
+    /// Cached errors from `parse_prefix` calls that returned `Err`. See
+    /// [`Parser::parse_prefix`] for the 2^N patterns this guards.
+    failed_prefix_positions: BTreeMap<usize, ParserError>,
+    /// Cached errors from the speculative reserved-word prefix arm. See
+    /// [`Parser::parse_prefix`] for the 2^N patterns this guards.
+    failed_reserved_word_prefix_positions: BTreeMap<usize, ParserError>,
 }
 
 impl<'a> Parser<'a> {
@@ -398,6 +408,8 @@ impl<'a> Parser<'a> {
             recursion_counter: RecursionCounter::new(DEFAULT_REMAINING_DEPTH),
             options: ParserOptions::new().with_trailing_commas(dialect.supports_trailing_commas()),
             failed_unary_not_positions: BTreeSet::new(),
+            failed_prefix_positions: BTreeMap::new(),
+            failed_reserved_word_prefix_positions: BTreeMap::new(),
         }
     }
 
@@ -460,6 +472,8 @@ impl<'a> Parser<'a> {
         self.tokens = tokens;
         self.index = 0;
         self.failed_unary_not_positions.clear();
+        self.failed_prefix_positions.clear();
+        self.failed_reserved_word_prefix_positions.clear();
         self
     }
 
@@ -1731,6 +1745,23 @@ impl<'a> Parser<'a> {
             return prefix;
         }
 
+        // Memoize parse_prefix failures to break 2^N speculation when both
+        // prefix arms fail at every level (e.g. `IF(current_time(...x`).
+        // The per-arm cache in `parse_prefix_inner` complements this for
+        // chains where the reserved arm fails but the unreserved fallback
+        // succeeds (e.g. `case-case-...c`).
+        let start_index = self.index;
+        if let Some(cached) = self.failed_prefix_positions.get(&start_index) {
+            return Err(cached.clone());
+        }
+        let result = self.parse_prefix_inner();
+        if let Err(ref e) = result {
+            self.failed_prefix_positions.insert(start_index, e.clone());
+        }
+        result
+    }
+
+    fn parse_prefix_inner(&mut self) -> Result<Expr, ParserError> {
         // PostgreSQL allows any string literal to be preceded by a type name, indicating that the
         // string literal represents a literal of that type. Some examples:
         //
@@ -1826,7 +1857,21 @@ impl<'a> Parser<'a> {
                 {
                     return self.parse_expr_prefix_by_unreserved_word(&w, span);
                 }
-                match self.try_parse(|parser| parser.parse_expr_prefix_by_reserved_word(&w, span)) {
+                // Memoize failed speculative reserved-word parses. When
+                // the reserved arm (CASE, CURRENT_TIME, etc.) does
+                // exponential work but the unreserved fallback ultimately
+                // succeeds, the overall `parse_prefix` returns `Ok` and the
+                // outer cache never fires. Chains like `case-case-...c`
+                // need this per-arm cache to break the doubling.
+                let try_parse_result = if let Some(cached) = self
+                    .failed_reserved_word_prefix_positions
+                    .get(&next_token_index)
+                {
+                    Err(cached.clone())
+                } else {
+                    self.try_parse(|parser| parser.parse_expr_prefix_by_reserved_word(&w, span))
+                };
+                match try_parse_result {
                     // This word indicated an expression prefix and parsing was successful
                     Ok(Some(expr)) => Ok(expr),
 
@@ -1843,6 +1888,8 @@ impl<'a> Parser<'a> {
                         if w.keyword == Keyword::NOT {
                             self.failed_unary_not_positions.insert(self.index);
                         }
+                        self.failed_reserved_word_prefix_positions
+                            .insert(next_token_index, e.clone());
                         if !self.dialect.is_reserved_for_identifier(w.keyword) {
                             if let Ok(Some(expr)) = self.maybe_parse(|parser| {
                                 parser.parse_expr_prefix_by_unreserved_word(&w, span)

tests/sqlparser_common.rs

@@ -19082,3 +19082,49 @@ fn parse_not_chain_no_exponential_blowup() {
     rx.recv_timeout(Duration::from_secs(5))
         .expect("parser should reject this quickly, not loop exponentially");
 }
+
+/// Regression test for the 2^N parse-time blowup in `parse_prefix` on inputs
+/// like `IF(current_time(current_time(...x`. Each nested `current_time(` used
+/// to be explored twice at every level (once via the speculative reserved-word
+/// arm, once via the unreserved-word fallback), doubling work per level.
+/// Post-fix the failing parse short-circuits via the position-keyed cache.
+#[test]
+fn parse_prefix_keyword_call_chain_no_exponential_blowup() {
+    use std::sync::mpsc;
+    use std::thread;
+    use std::time::Duration;
+
+    let sql = String::from("if(") + &"current_time(".repeat(30) + "x";
+
+    let (tx, rx) = mpsc::channel();
+    thread::spawn(move || {
+        let _ = Parser::parse_sql(&PostgreSqlDialect {}, &sql);
+        let _ = tx.send(());
+    });
+
+    rx.recv_timeout(Duration::from_secs(5))
+        .expect("parser should reject this quickly, not loop exponentially");
+}
+
+/// Regression test for the 2^N parse-time blowup in `parse_prefix` on inputs
+/// like `case-case-case-...c`. Each `case` token triggers a speculative
+/// `parse_case_expr` that fails, but the unreserved-word fallback returns
+/// `Identifier(case)`, so the outer failure cache never fires. Post-fix the
+/// per-arm cache short-circuits the speculative descent.
+#[test]
+fn parse_prefix_case_chain_no_exponential_blowup() {
+    use std::sync::mpsc;
+    use std::thread;
+    use std::time::Duration;
+
+    let sql = "case\t-".repeat(30) + "c";
+
+    let (tx, rx) = mpsc::channel();
+    thread::spawn(move || {
+        let _ = Parser::parse_sql(&SQLiteDialect {}, &sql);
+        let _ = tx.send(());
+    });
+
+    rx.recv_timeout(Duration::from_secs(5))
+        .expect("parser should reject this quickly, not loop exponentially");
+}