Skip to content

Commit 21eb23f

Browse files
KKouldKKould
committed
perf: optimize parse embedding
1 parent 962c24f commit 21eb23f

File tree

4 files changed

+119
-99
lines changed

4 files changed

+119
-99
lines changed

‎src/query/ast/benches/bench.rs‎

Lines changed: 14 additions & 5 deletions
Large diffs are not rendered by default.

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

Lines changed: 92 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,29 @@ use crate::parser::Error;
3434
use crate::parser::ErrorKind;
3535
use crate::Span;
3636

37+
macro_rules! with_span {
38+
($parser:expr) => {
39+
map(consumed($parser), |(span, elem)| WithSpan { span, elem })
40+
};
41+
}
42+
43+
macro_rules! try_dispatch {
44+
($input:expr, $($pat:pat => $body:expr),+ $(,)?) => {{
45+
if let Some(token_0) = $input.tokens.first() {
46+
use TokenKind::*;
47+
48+
if let Some(result) = match token_0.kind {
49+
$($pat => Some($body),)+
50+
_ => None,
51+
} {
52+
if result.is_ok() {
53+
return result;
54+
}
55+
}
56+
}
57+
}};
58+
}
59+
3760
pub fn expr(i: Input) -> IResult<Expr> {
3861
context("expression", subexpr(0)).parse(i)
3962
}
@@ -1214,7 +1237,15 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
12141237
let variable_access = map(variable_ident, ExprElement::VariableAccess);
12151238

12161239
let unary_op = map(unary_op, |op| ExprElement::UnaryOp { op });
1217-
let map_access = map(map_access, |accessor| ExprElement::MapAccess { accessor });
1240+
let bracket_map_access = map(map_access_bracket, |accessor| ExprElement::MapAccess {
1241+
accessor,
1242+
});
1243+
let dot_number_map_access = map(map_access_dot_number, |accessor| ExprElement::MapAccess {
1244+
accessor,
1245+
});
1246+
let colon_map_access = map(map_access_colon, |accessor| ExprElement::MapAccess {
1247+
accessor,
1248+
});
12181249
let dot_access = map(
12191250
rule! {
12201251
"." ~ #column_id
@@ -1531,30 +1562,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
15311562
ExprElement::StageLocation { location }
15321563
});
15331564

1534-
macro_rules! with_span {
1535-
($parser:expr) => {
1536-
map(consumed($parser), |(span, elem)| WithSpan { span, elem })
1537-
};
1538-
}
1539-
1540-
macro_rules! try_dispatch {
1541-
($($pat:pat => $body:expr),+ $(,)?) => {{
1542-
if let Some(token_0) = i.tokens.first() {
1543-
use TokenKind::*;
1544-
1545-
if let Some(result) = match token_0.kind {
1546-
$($pat => Some($body),)+
1547-
_ => None,
1548-
} {
1549-
if result.is_ok() {
1550-
return result;
1551-
}
1552-
}
1553-
}
1554-
}};
1555-
}
1556-
1557-
try_dispatch!(
1565+
try_dispatch!(i,
15581566
IS => with_span!(rule!(#is_null | #is_distinct_from)).parse(i),
15591567
NOT => with_span!(rule!(
15601568
#in_list
@@ -1584,13 +1592,17 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
15841592
LParen => with_span!(rule!(#tuple | #subquery)).parse(i),
15851593
ANY | SOME | ALL => with_span!(subquery).parse(i),
15861594
Dot => {
1587-
return with_span!(rule!(#chain_function_call | #dot_access | #map_access)).parse(i);
1595+
return with_span!(rule!(#chain_function_call | #dot_access | #dot_number_map_access))
1596+
.parse(i);
15881597
},
15891598
Colon => {
1590-
return with_span!(map_access).parse(i);
1599+
return with_span!(colon_map_access).parse(i);
15911600
},
15921601
LBracket => {
1593-
return with_span!(rule!(#list_comprehensions | #map_access | #array)).parse(i);
1602+
return with_span!(rule!(
1603+
#list_comprehensions | #bracket_map_access | #array
1604+
))
1605+
.parse(i);
15941606
},
15951607
LBrace => with_span!(map_expr).parse(i),
15961608
LiteralAtString => with_span!(stage_location).parse(i),
@@ -1669,7 +1681,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
16691681

16701682
with_span!(alt((rule!(
16711683
#column_ref : "<column>"
1672-
| #map_access : "[<key>] | .<key> | :<key>"
1684+
| #dot_number_map_access : ".<key>"
16731685
| #literal : "<literal>"
16741686
),)))
16751687
.parse(i)
@@ -1796,19 +1808,37 @@ pub fn json_op(i: Input) -> IResult<JsonOperator> {
17961808
}
17971809

17981810
pub fn literal(i: Input) -> IResult<Literal> {
1799-
let string = map(literal_string, Literal::String);
1800-
let code_string = map(code_string, Literal::String);
1801-
let boolean = map(literal_bool, Literal::Boolean);
1802-
let null = value(Literal::Null, rule! { NULL });
1811+
let mut string = map(literal_string, Literal::String);
1812+
let mut code_string = map(code_string, Literal::String);
1813+
let mut boolean = map(literal_bool, Literal::Boolean);
1814+
let mut null = value(Literal::Null, rule! { NULL });
1815+
let mut decimal_uint = map_res(
1816+
rule! {
1817+
LiteralInteger
1818+
},
1819+
|token| parse_uint(token.text(), 10).map_err(nom::Err::Failure),
1820+
);
1821+
let mut hex_uint = map_res(literal_hex_str, |str| {
1822+
parse_uint(str, 16).map_err(nom::Err::Failure)
1823+
});
1824+
let mut decimal_float = map_res(
1825+
rule! {
1826+
LiteralFloat
1827+
},
1828+
|token| parse_float(token.text()).map_err(nom::Err::Failure),
1829+
);
18031830

1804-
rule!(
1805-
#string
1806-
| #code_string
1807-
| #boolean
1808-
| #literal_number
1809-
| #null
1810-
)
1811-
.parse(i)
1831+
try_dispatch!(i,
1832+
LiteralString => string.parse(i),
1833+
LiteralCodeString => code_string.parse(i),
1834+
LiteralInteger => decimal_uint.parse(i),
1835+
LiteralFloat => decimal_float.parse(i),
1836+
MySQLLiteralHex | PGLiteralHex => hex_uint(i),
1837+
TRUE | FALSE => boolean.parse(i),
1838+
NULL => null.parse(i),
1839+
);
1840+
1841+
Err(nom::Err::Error(Error::from_error_kind(i, ErrorKind::Other("expecting `LiteralString`, 'LiteralCodeString', 'LiteralInteger', 'LiteralFloat', 'TRUE', 'FALSE', or more ..."))))
18121842
}
18131843

18141844
pub fn literal_hex_str(i: Input) -> IResult<&str> {
@@ -1872,33 +1902,6 @@ pub fn literal_i64(i: Input) -> IResult<i64> {
18721902
.parse(i)
18731903
}
18741904

1875-
pub fn literal_number(i: Input) -> IResult<Literal> {
1876-
let decimal_uint = map_res(
1877-
rule! {
1878-
LiteralInteger
1879-
},
1880-
|token| parse_uint(token.text(), 10).map_err(nom::Err::Failure),
1881-
);
1882-
1883-
let hex_uint = map_res(literal_hex_str, |str| {
1884-
parse_uint(str, 16).map_err(nom::Err::Failure)
1885-
});
1886-
1887-
let decimal_float = map_res(
1888-
rule! {
1889-
LiteralFloat
1890-
},
1891-
|token| parse_float(token.text()).map_err(nom::Err::Failure),
1892-
);
1893-
1894-
rule!(
1895-
#decimal_uint
1896-
| #decimal_float
1897-
| #hex_uint
1898-
)
1899-
.parse(i)
1900-
}
1901-
19021905
pub fn literal_bool(i: Input) -> IResult<bool> {
19031906
alt((value(true, rule! { TRUE }), value(false, rule! { FALSE }))).parse(i)
19041907
}
@@ -2406,37 +2409,34 @@ pub fn interval_kind(i: Input) -> IResult<IntervalKind> {
24062409
.parse(i)
24072410
}
24082411

2409-
pub fn map_access(i: Input) -> IResult<MapAccessor> {
2410-
let bracket = map(
2412+
fn map_access_bracket(i: Input) -> IResult<MapAccessor> {
2413+
map(
24112414
rule! {
2412-
"[" ~ #subexpr(0) ~ "]"
2415+
"[" ~ #subexpr(0) ~ "]"
24132416
},
24142417
|(_, key, _)| MapAccessor::Bracket { key: Box::new(key) },
2415-
);
2416-
let dot_number = map_res(
2417-
rule! {
2418-
LiteralFloat
2419-
},
2420-
|key| {
2421-
if key.text().starts_with('.') {
2422-
if let Ok(key) = (key.text()[1..]).parse::<u64>() {
2423-
return Ok(MapAccessor::DotNumber { key });
2424-
}
2418+
)
2419+
.parse(i)
2420+
}
2421+
2422+
fn map_access_dot_number(i: Input) -> IResult<MapAccessor> {
2423+
map_res(rule! { LiteralFloat }, |key| {
2424+
if key.text().starts_with('.') {
2425+
if let Ok(key) = (key.text()[1..]).parse::<u64>() {
2426+
return Ok(MapAccessor::DotNumber { key });
24252427
}
2426-
Err(nom::Err::Error(ErrorKind::ExpectText(".")))
2427-
},
2428-
);
2429-
let colon = map(
2428+
}
2429+
Err(nom::Err::Error(ErrorKind::ExpectText(".")))
2430+
})
2431+
.parse(i)
2432+
}
2433+
2434+
fn map_access_colon(i: Input) -> IResult<MapAccessor> {
2435+
map(
24302436
rule! {
2431-
":" ~ #ident
2437+
":" ~ #ident
24322438
},
24332439
|(_, key)| MapAccessor::Colon { key },
2434-
);
2435-
2436-
rule!(
2437-
#bracket
2438-
| #dot_number
2439-
| #colon
24402440
)
24412441
.parse(i)
24422442
}

‎src/query/ast/src/parser/script.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ use nom_rule::rule;
1717

1818
use crate::ast::*;
1919
use crate::parser::common::*;
20+
use crate::parser::error::Error;
21+
use crate::parser::error::ErrorKind;
2022
use crate::parser::expr::*;
2123
use crate::parser::input::Input;
2224
use crate::parser::statement::*;
@@ -157,6 +159,15 @@ pub fn script_stmts(i: Input) -> IResult<Vec<ScriptStatement>> {
157159
}
158160

159161
pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
162+
if let Some(token) = i.tokens.first() {
163+
let kind = token.kind;
164+
if matches!(kind, END | ELSE | ELSEIF | WHEN | UNTIL) {
165+
return Err(nom::Err::Error(Error::from_error_kind(
166+
i,
167+
ErrorKind::Other("block terminator"),
168+
)));
169+
}
170+
}
160171
let let_var_stmt = map(
161172
rule! {
162173
LET ~ #declare_var

‎src/query/script/tests/it/testdata/script-error.txt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,11 @@ LET x := 'min';
7373
LET y := IDENTIFIER(:x)([1,2]);
7474
---------- Output ----------
7575
error:
76-
--> SQL:2:10
76+
--> SQL:2:1
7777
|
7878
1 | LET x := 'min';
7979
2 | LET y := IDENTIFIER(:x)([1,2]);
80-
| ^^^^^^^^^^^^^^ variable is not allowed in this context
80+
| ^^^ unable to parse rest of the sql, rest tokens: [LET(16..19), Ident(20..21), ColonEqual(22..24), IDENTIFIER(25..35), LParen(35..36), Colon(36..37), Ident(37..38), RParen(38..39), LParen(39..40), LBracket(40..41), LiteralInteger(41..42), Comma(42..43), LiteralInteger(43..44), RBracket(44..45), RParen(45..46), SemiColon(46..47), EOI(47..47)]
8181

8282

8383
---------- Input ----------

0 commit comments

Comments
 (0)