perf: use try_dispatch to categorize the statement_body, reducing the number of branches and stack usage (otherwise, stack overflow is extremely likely).

Author: KKould

src/query/ast/benches/bench.rs

@@ -18,12 +18,12 @@ fn main() {
 
 // bench                     fastest       │ slowest       │ median        │ mean          │ samples │ iters
 // ╰─ dummy                                │               │               │               │         │
-// ├─ deep_function_call  132.5 µs      │ 470.9 µs      │ 162.3 µs      │ 166.2 µs      │ 100     │ 100
-// ├─ deep_query          232.3 µs      │ 495 µs        │ 283.6 µs      │ 283.4 µs      │ 100     │ 100
-// ├─ large_query         196.4 µs      │ 348.2 µs      │ 240.1 µs      │ 240.1 µs      │ 100     │ 100
-// ├─ large_statement     197.6 µs      │ 286.6 µs      │ 234.6 µs      │ 236.1 µs      │ 100     │ 100
-// ├─ wide_embedding      3.964 ms      │ 7.229 ms      │ 4.3 ms        │ 4.397 ms      │ 100     │ 100
-// ╰─ wide_expr           40.59 µs      │ 56.32 µs      │ 43.06 µs      │ 43.68 µs      │ 100     │ 100
+// ├─ deep_function_call  117.2 µs      │ 407.3 µs      │ 127.5 µs      │ 132.2 µs      │ 100     │ 100
+// ├─ deep_query          229.7 µs      │ 395.8 µs      │ 242.6 µs      │ 244.6 µs      │ 100     │ 100
+// ├─ large_query         203.8 µs      │ 275.7 µs      │ 210.9 µs      │ 212.8 µs      │ 100     │ 100
+// ├─ large_statement     206.7 µs      │ 238.3 µs      │ 212.9 µs      │ 213.7 µs      │ 100     │ 100
+// ├─ wide_embedding      3.628 ms      │ 4.088 ms      │ 3.761 ms      │ 3.762 ms      │ 100     │ 100
+// ╰─ wide_expr           39.76 µs      │ 46.04 µs      │ 40.18 µs      │ 40.43 µs      │ 100     │ 100
 
 #[divan::bench_group(max_time = 0.5)]
 mod dummy {

src/query/ast/src/parser/expr.rs

@@ -40,23 +40,6 @@ macro_rules! with_span {
     };
 }
 
-macro_rules! try_dispatch {
-    ($input:expr, $($pat:pat => $body:expr),+ $(,)?) => {{
-        if let Some(token_0) = $input.tokens.first() {
-            use TokenKind::*;
-
-            if let Some(result) = match token_0.kind {
-                $($pat => Some($body),)+
-                _ => None,
-            } {
-                if result.is_ok() {
-                    return result;
-                }
-            }
-        }
-    }};
-}
-
 pub fn expr(i: Input) -> IResult<Expr> {
     context("expression", subexpr(0)).parse(i)
 }
@@ -1562,7 +1545,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
         ExprElement::StageLocation { location }
     });
 
-    try_dispatch!(i,
+    try_dispatch!(i, true,
         IS => with_span!(rule!(#is_null | #is_distinct_from)).parse(i),
         NOT => with_span!(rule!(
             #in_list
@@ -1828,7 +1811,7 @@ pub fn literal(i: Input) -> IResult<Literal> {
         |token| parse_float(token.text()).map_err(nom::Err::Failure),
     );
 
-    try_dispatch!(i,
+    try_dispatch!(i, true,
         LiteralString => string.parse(i),
         LiteralCodeString => code_string.parse(i),
         LiteralInteger => decimal_uint.parse(i),

src/query/ast/src/parser/mod.rs

@@ -12,6 +12,23 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+macro_rules! try_dispatch {
+    ($input:expr, $return_if_ok:literal, $($pat:pat => $body:expr),+ $(,)?) => {{
+        if let Some(token_0) = $input.tokens.first() {
+            use TokenKind::*;
+
+            if let Some(result) = match token_0.kind {
+                $($pat => Some($body),)+
+                _ => None,
+            } {
+                if !$return_if_ok || result.is_ok() {
+                    return result;
+                }
+            }
+        }
+    }};
+}
+
 mod comment;
 mod common;
 mod copy;