feat: optimize expression parse (#18871)

KKould · web-flow · commit 16c7380bd362 · 2025-10-21T20:37:41.000+08:00
* perf: optimize expression parse

* chore: update ast bench comment

* chore: fix 03_0016_insert_into_values.test

* chore: codefmt
diff --git a/src/query/ast/benches/bench.rs b/src/query/ast/benches/bench.rs
@@ -18,11 +18,11 @@ fn main() {
 
 // bench                     fastest       │ slowest       │ median        │ mean          │ samples │ iters
 // ╰─ dummy                                │               │               │               │         │
-// ├─ deep_function_call  1.238 ms      │ 1.781 ms      │ 1.355 ms      │ 1.353 ms      │ 100     │ 100
-// ├─ deep_query          285.6 µs      │ 434.6 µs      │ 306.8 µs      │ 307 µs        │ 100     │ 100
-// ├─ large_query         1.739 ms      │ 1.893 ms      │ 1.795 ms      │ 1.801 ms      │ 100     │ 100
-// ├─ large_statement     1.745 ms      │ 1.885 ms      │ 1.807 ms      │ 1.806 ms      │ 100     │ 100
-// ╰─ wide_expr           562 µs        │ 651.9 µs      │ 588.2 µs      │ 590.5 µs      │ 100     │ 100
+//     ├─ deep_function_call  802.2 µs      │ 1.207 ms      │ 842 µs        │ 850.6 µs      │ 100     │ 100
+//     ├─ deep_query          242.3 µs      │ 426.3 µs      │ 254.2 µs      │ 257.3 µs      │ 100     │ 100
+//     ├─ large_query         1.104 ms      │ 1.264 ms      │ 1.14 ms       │ 1.142 ms      │ 100     │ 100
+//     ├─ large_statement     1.097 ms      │ 1.2 ms        │ 1.15 ms       │ 1.148 ms      │ 100     │ 100
+//     ╰─ wide_expr           282.4 µs      │ 368.6 µs      │ 298 µs        │ 298.7 µs      │ 100     │ 100
 
 #[divan::bench_group(max_time = 0.5)]
 mod dummy {
diff --git a/src/query/ast/src/parser/common.rs b/src/query/ast/src/parser/common.rs
@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::cell::RefCell;
+use std::rc::Rc;
+
 use nom::branch::alt;
 use nom::combinator::consumed;
 use nom::combinator::map;
@@ -580,20 +583,61 @@ pub fn transform_span(tokens: &[Token]) -> Span {
     })
 }
 
-pub fn run_pratt_parser<'a, I, P, E>(
+pub(crate) trait IterProvider<'a> {
+    type Item;
+    type Iter: Iterator<Item = Self::Item> + ExactSizeIterator;
+
+    fn create_iter(self, span: Rc<RefCell<Option<Input<'a>>>>) -> Self::Iter;
+}
+
+impl<'a, T> IterProvider<'a> for Vec<WithSpan<'a, T>>
+where T: Clone
+{
+    type Item = WithSpan<'a, T>;
+    type Iter = ErrorSpan<'a, T, std::vec::IntoIter<WithSpan<'a, T>>>;
+
+    fn create_iter(self, span: Rc<RefCell<Option<Input<'a>>>>) -> Self::Iter {
+        ErrorSpan::new(self.into_iter(), span)
+    }
+}
+
+pub(crate) struct ErrorSpan<'a, T, I: Iterator<Item = WithSpan<'a, T>>> {
+    iter: I,
+    span: Rc<RefCell<Option<Input<'a>>>>,
+}
+
+impl<'a, T, I: Iterator<Item = WithSpan<'a, T>>> ErrorSpan<'a, T, I> {
+    fn new(iter: I, span: Rc<RefCell<Option<Input<'a>>>>) -> Self {
+        Self { iter, span }
+    }
+}
+
+impl<'a, T, I: Iterator<Item = WithSpan<'a, T>>> Iterator for ErrorSpan<'a, T, I> {
+    type Item = WithSpan<'a, T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter
+            .next()
+            .inspect(|item| *self.span.borrow_mut() = Some(item.span))
+    }
+}
+
+impl<'a, T, I: Iterator<Item = WithSpan<'a, T>>> ExactSizeIterator for ErrorSpan<'a, T, I> {}
+
+pub fn run_pratt_parser<'a, I, P, E, T>(
     mut parser: P,
-    iter: &I,
+    parsers: T,
     rest: Input<'a>,
     input: Input<'a>,
 ) -> IResult<'a, P::Output>
 where
     E: std::fmt::Debug,
     P: PrattParser<I, Input = WithSpan<'a, E>, Error = &'static str>,
-    I: Iterator<Item = P::Input> + ExactSizeIterator + Clone,
+    I: Iterator<Item = P::Input> + ExactSizeIterator,
+    T: IterProvider<'a, Item = P::Input, Iter = I>,
 {
-    let mut iter_cloned = iter.clone();
-    let mut iter = iter.clone().peekable();
-    let len = iter.len();
+    let span = Rc::new(RefCell::new(None));
+    let mut iter = parsers.create_iter(span.clone()).peekable();
     let expr = parser
         .parse_input(&mut iter, Precedence(0))
         .map_err(|err| {
@@ -602,22 +646,27 @@ where
 
             let err_kind = match err {
                 PrattError::EmptyInput => ErrorKind::Other("expecting an operand"),
-                PrattError::UnexpectedNilfix(_) => ErrorKind::Other("unable to parse the element"),
-                PrattError::UnexpectedPrefix(_) => {
+                PrattError::UnexpectedNilfix(i) => {
+                    *span.borrow_mut() = Some(i.span);
+                    ErrorKind::Other("unable to parse the element")
+                }
+                PrattError::UnexpectedPrefix(i) => {
+                    *span.borrow_mut() = Some(i.span);
                     ErrorKind::Other("unable to parse the prefix operator")
                 }
-                PrattError::UnexpectedInfix(_) => {
+                PrattError::UnexpectedInfix(i) => {
+                    *span.borrow_mut() = Some(i.span);
                     ErrorKind::Other("missing lhs or rhs for the binary operator")
                 }
-                PrattError::UnexpectedPostfix(_) => {
+                PrattError::UnexpectedPostfix(i) => {
+                    *span.borrow_mut() = Some(i.span);
                     ErrorKind::Other("unable to parse the postfix operator")
                 }
                 PrattError::UserError(err) => ErrorKind::Other(err),
             };
 
-            let span = iter_cloned
-                .nth(len - iter.len() - 1)
-                .map(|elem| elem.span)
+            let span = span
+                .take()
                 // It's safe to slice one more token because input must contain EOI.
                 .unwrap_or_else(|| rest.slice(..1));
 
diff --git a/src/query/ast/src/parser/expr.rs b/src/query/ast/src/parser/expr.rs
@@ -19,6 +19,7 @@ use nom::combinator::consumed;
 use nom::combinator::map;
 use nom::combinator::value;
 use nom::error::context;
+use nom::Slice;
 use nom_rule::rule;
 use pratt::Affix;
 use pratt::Associativity;
@@ -126,7 +127,7 @@ pub fn subexpr(min_precedence: u32) -> impl FnMut(Input) -> IResult<Expr> {
             }
         }
 
-        run_pratt_parser(ExprParser, &expr_elements.into_iter(), rest, i)
+        run_pratt_parser(ExprParser, expr_elements, rest, i)
     }
 }
 
@@ -1567,76 +1568,127 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
     )(i)
 }
 
+#[inline]
+fn return_op<T>(i: Input, start: usize, op: T) -> IResult<T> {
+    Ok((i.slice(start..), op))
+}
+
+macro_rules! op_branch {
+    ($i:ident, $token_0:ident, $($kind:ident => $op:expr),+ $(,)?) => {
+        match $token_0.kind {
+            $(
+                TokenKind::$kind => return return_op($i, 1, $op),
+            )+
+            _ => (),
+        }
+    };
+}
+
 pub fn unary_op(i: Input) -> IResult<UnaryOperator> {
     // Plus and Minus are parsed as binary op at first.
-    alt((
-        value(UnaryOperator::Not, rule! { NOT }),
-        value(UnaryOperator::Factorial, rule! { Factorial }),
-        value(UnaryOperator::SquareRoot, rule! { SquareRoot }),
-        value(UnaryOperator::BitwiseNot, rule! { BitWiseNot }),
-        value(UnaryOperator::CubeRoot, rule! { CubeRoot }),
-        value(UnaryOperator::Abs, rule! { Abs }),
-    ))(i)
+    if let Some(token_0) = i.tokens.first() {
+        op_branch!(
+            i, token_0,
+            NOT => UnaryOperator::Not,
+            Factorial => UnaryOperator::Factorial,
+            SquareRoot => UnaryOperator::SquareRoot,
+            BitWiseNot => UnaryOperator::BitwiseNot,
+            CubeRoot => UnaryOperator::CubeRoot,
+            Abs => UnaryOperator::Abs,
+        );
+    }
+    Err(nom::Err::Error(Error::from_error_kind(
+        i,
+        ErrorKind::Other("expecting `NOT`, '!', '|/', '~', '||/', '@', or more ..."),
+    )))
 }
 
 pub fn binary_op(i: Input) -> IResult<BinaryOperator> {
-    alt((
-        alt((
-            value(BinaryOperator::Plus, rule! { "+" }),
-            value(BinaryOperator::Minus, rule! { "-" }),
-            value(BinaryOperator::Multiply, rule! { "*" }),
-            value(BinaryOperator::Divide, rule! { "/" }),
-            value(BinaryOperator::IntDiv, rule! { "//" }),
-            value(BinaryOperator::Div, rule! { DIV }),
-            value(BinaryOperator::Modulo, rule! { "%" }),
-            value(BinaryOperator::StringConcat, rule! { "||" }),
-            value(BinaryOperator::CosineDistance, rule! { "<=>" }),
-            value(BinaryOperator::L1Distance, rule! { "<+>" }),
-            value(BinaryOperator::L2Distance, rule! { "<->" }),
-            value(BinaryOperator::Gt, rule! { ">" }),
-            value(BinaryOperator::Lt, rule! { "<" }),
-            value(BinaryOperator::Gte, rule! { ">=" }),
-            value(BinaryOperator::Lte, rule! { "<=" }),
-            value(BinaryOperator::Eq, rule! { "=" }),
-            value(BinaryOperator::NotEq, rule! { "<>" | "!=" }),
-            value(BinaryOperator::Caret, rule! { "^" }),
-        )),
-        alt((
-            value(BinaryOperator::And, rule! { AND }),
-            value(BinaryOperator::Or, rule! { OR }),
-            value(BinaryOperator::Xor, rule! { XOR }),
-            value(BinaryOperator::LikeAny(None), rule! { LIKE ~ ANY }),
-            value(BinaryOperator::Like(None), rule! { LIKE }),
-            value(BinaryOperator::NotLike(None), rule! { NOT ~ LIKE }),
-            value(BinaryOperator::Regexp, rule! { REGEXP }),
-            value(BinaryOperator::NotRegexp, rule! { NOT ~ REGEXP }),
-            value(BinaryOperator::RLike, rule! { RLIKE }),
-            value(BinaryOperator::NotRLike, rule! { NOT ~ RLIKE }),
-            value(BinaryOperator::SoundsLike, rule! { SOUNDS ~ LIKE }),
-            value(BinaryOperator::BitwiseOr, rule! { BitWiseOr }),
-            value(BinaryOperator::BitwiseAnd, rule! { BitWiseAnd }),
-            value(BinaryOperator::BitwiseXor, rule! { BitWiseXor }),
-            value(BinaryOperator::BitwiseShiftLeft, rule! { ShiftLeft }),
-            value(BinaryOperator::BitwiseShiftRight, rule! { ShiftRight }),
-        )),
-    ))(i)
+    if let Some(token_0) = i.tokens.first() {
+        op_branch!(
+            i, token_0,
+            Plus => BinaryOperator::Plus,
+            Minus => BinaryOperator::Minus,
+            Multiply => BinaryOperator::Multiply,
+            Divide => BinaryOperator::Divide,
+            IntDiv => BinaryOperator::IntDiv,
+            DIV => BinaryOperator::Div,
+            Modulo => BinaryOperator::Modulo,
+            StringConcat => BinaryOperator::StringConcat,
+            Spaceship => BinaryOperator::CosineDistance,
+            L1DISTANCE => BinaryOperator::L1Distance,
+            L2DISTANCE => BinaryOperator::L2Distance,
+            Gt => BinaryOperator::Gt,
+            Lt => BinaryOperator::Lt,
+            Gte => BinaryOperator::Gte,
+            Lte => BinaryOperator::Lte,
+            Eq => BinaryOperator::Eq,
+            NotEq => BinaryOperator::NotEq,
+            Caret => BinaryOperator::Caret,
+            AND => BinaryOperator::And,
+            OR => BinaryOperator::Or,
+            XOR => BinaryOperator::Xor,
+            REGEXP => BinaryOperator::Regexp,
+            RLIKE => BinaryOperator::RLike,
+            BitWiseOr => BinaryOperator::BitwiseOr,
+            BitWiseAnd => BinaryOperator::BitwiseAnd,
+            BitWiseXor => BinaryOperator::BitwiseXor,
+            ShiftLeft => BinaryOperator::BitwiseShiftLeft,
+            ShiftRight => BinaryOperator::BitwiseShiftRight,
+        );
+        match token_0.kind {
+            TokenKind::LIKE => {
+                return if matches!(
+                    i.tokens.get(1).map(|first| first.kind == TokenKind::ANY),
+                    Some(true)
+                ) {
+                    return_op(i, 2, BinaryOperator::LikeAny(None))
+                } else {
+                    return_op(i, 1, BinaryOperator::Like(None))
+                }
+            }
+            TokenKind::NOT => match i.tokens.get(1).map(|first| first.kind) {
+                Some(TokenKind::LIKE) => {
+                    return return_op(i, 2, BinaryOperator::NotLike(None));
+                }
+                Some(TokenKind::REGEXP) => {
+                    return return_op(i, 2, BinaryOperator::NotRegexp);
+                }
+                Some(TokenKind::RLIKE) => {
+                    return return_op(i, 2, BinaryOperator::NotRLike);
+                }
+                _ => (),
+            },
+            TokenKind::SOUNDS => {
+                if let Some(TokenKind::LIKE) = i.tokens.get(1).map(|first| first.kind) {
+                    return return_op(i, 2, BinaryOperator::SoundsLike);
+                }
+            }
+            _ => (),
+        }
+    }
+    Err(nom::Err::Error(Error::from_error_kind(i, ErrorKind::Other("expecting `IS`, `IN`, `LIKE`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<=>`, `<+>`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `NOT`, `REGEXP`, `RLIKE`, `SOUNDS`, or more ..."))))
 }
 
-pub fn json_op(i: Input) -> IResult<JsonOperator> {
-    alt((
-        value(JsonOperator::Arrow, rule! { "->" }),
-        value(JsonOperator::LongArrow, rule! { "->>" }),
-        value(JsonOperator::HashArrow, rule! { "#>" }),
-        value(JsonOperator::HashLongArrow, rule! { "#>>" }),
-        value(JsonOperator::Question, rule! { "?" }),
-        value(JsonOperator::QuestionOr, rule! { "?|" }),
-        value(JsonOperator::QuestionAnd, rule! { "?&" }),
-        value(JsonOperator::AtArrow, rule! { "@>" }),
-        value(JsonOperator::ArrowAt, rule! { "<@" }),
-        value(JsonOperator::AtQuestion, rule! { "@?" }),
-        value(JsonOperator::AtAt, rule! { "@@" }),
-        value(JsonOperator::HashMinus, rule! { "#-" }),
-    ))(i)
+pub(crate) fn json_op(i: Input) -> IResult<JsonOperator> {
+    if let Some(token_0) = i.tokens.first() {
+        op_branch!(
+            i, token_0,
+            RArrow => JsonOperator::Arrow,
+            LongRArrow => JsonOperator::LongArrow,
+            HashRArrow => JsonOperator::HashArrow,
+            HashLongRArrow => JsonOperator::HashLongArrow,
+            Placeholder => JsonOperator::Question,
+            QuestionOr => JsonOperator::QuestionOr,
+            QuestionAnd => JsonOperator::QuestionAnd,
+            AtArrow => JsonOperator::AtArrow,
+            ArrowAt => JsonOperator::ArrowAt,
+            AtQuestion => JsonOperator::AtQuestion,
+            AtAt => JsonOperator::AtAt,
+            HashMinus => JsonOperator::HashMinus,
+        );
+    }
+    Err(nom::Err::Error(Error::from_error_kind(i, ErrorKind::Other("expecting `->`, '->>', '#>', '#>>', '?', '?|', '?&', '@>', '<@', '@?', '@@', '#-', or more ..."))))
 }
 
 pub fn literal(i: Input) -> IResult<Literal> {
diff --git a/src/query/ast/src/parser/query.rs b/src/query/ast/src/parser/query.rs
@@ -48,8 +48,7 @@ pub fn query(i: Input) -> IResult<Query> {
 
 pub fn set_operation(i: Input) -> IResult<SetExpr> {
     let (rest, set_operation_elements) = rule! { #set_operation_element+ }(i)?;
-    let iter = &mut set_operation_elements.into_iter();
-    run_pratt_parser(SetOperationParser, iter, rest, i)
+    run_pratt_parser(SetOperationParser, set_operation_elements, rest, i)
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -711,8 +710,7 @@ pub fn order_by_expr(i: Input) -> IResult<OrderByExpr> {
 
 pub fn table_reference(i: Input) -> IResult<TableReference> {
     let (rest, table_reference_elements) = rule! { #table_reference_element+ }(i)?;
-    let iter = &mut table_reference_elements.into_iter();
-    run_pratt_parser(TableReferenceParser, iter, rest, i)
+    run_pratt_parser(TableReferenceParser, table_reference_elements, rest, i)
 }
 
 #[derive(Debug, Clone, PartialEq)]
diff --git a/src/query/ast/tests/it/testdata/dialect.txt b/src/query/ast/tests/it/testdata/dialect.txt
@@ -237,7 +237,7 @@ error:
 1 | "a\"b"
   | ^^^^^^
   | |
-  | invalid identifier
+  | expecting `IS`, `IN`, `LIKE`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<=>`, `<+>`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `NOT`, `REGEXP`, `RLIKE`, `SOUNDS`, or more ...
   | while parsing expression
 
 
diff --git a/src/query/ast/tests/it/testdata/expr-error.txt b/src/query/ast/tests/it/testdata/expr-error.txt
@@ -52,7 +52,7 @@ error:
   --> SQL:1:10
   |
 1 | CAST(col1)
-  | ----     ^ unexpected `)`, expecting `AS`, `,`, `(`, `IS`, `NOT`, `IN`, `LIKE`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<=>`, `<+>`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `::`, `POSITION`, or 49 more ...
+  | ----     ^ expecting `IS`, `IN`, `LIKE`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<=>`, `<+>`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `NOT`, `REGEXP`, `RLIKE`, `SOUNDS`, or more ...
   | |         
   | while parsing `CAST(... AS ...)`
   | while parsing expression
@@ -81,7 +81,7 @@ error:
 1 | $ abc + 3
   | ^
   | |
-  | unexpected `$`, expecting `IS`, `IN`, `LIKE`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<=>`, `<+>`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `NOT`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `::`, `POSITION`, `IdentVariable`, `DATEADD`, `DATE_ADD`, or 47 more ...
+  | expecting `IS`, `IN`, `LIKE`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<=>`, `<+>`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `NOT`, `REGEXP`, `RLIKE`, `SOUNDS`, or more ...
   | while parsing expression
 
 
diff --git a/src/query/ast/tests/it/testdata/stmt-error.txt b/src/query/ast/tests/it/testdata/stmt-error.txt
diff --git a/tests/sqllogictests/suites/base/03_common/03_0016_insert_into_values.test b/tests/sqllogictests/suites/base/03_common/03_0016_insert_into_values.test
