tract_data/dim/
parse.rs

1use super::*;
2use nom::branch::alt;
3use nom::bytes::complete::tag;
4use nom::character::complete::{alpha1, alphanumeric1, digit1, one_of};
5use nom::combinator::{all_consuming, map, map_res, recognize};
6use nom::multi::{fold, many0, separated_list0};
7use nom::sequence::{delimited, pair, preceded, separated_pair};
8use nom::{IResult, Parser};
9use nom_language::error::VerboseError;
10
11type R<'i, O> = IResult<&'i str, O, VerboseError<&'i str>>;
12
13pub fn parse_tdim(symbol_table: &SymbolScope, input: &str) -> TractResult<TDim> {
14    match all_consuming(|i| expr(symbol_table, i)).parse(input) {
15        Ok(pair) => Ok(pair.1),
16        Err(e) => bail!("Failed to parse {:?}, {:?}", input, e),
17    }
18}
19
20pub fn parse_assertion(symbol_table: &SymbolScope, input: &str) -> TractResult<Assertion> {
21    match all_consuming(|i| assertion(symbol_table, i)).parse(input) {
22        Ok(pair) => Ok(pair.1),
23        Err(e) => bail!("Failed to parse {:?}, {:?}", input, e),
24    }
25}
26
27fn assertion<'i>(s: &SymbolScope, i: &'i str) -> R<'i, Assertion> {
28    delimited(
29        spaces,
30        alt((
31            map(separated_pair(|i| expr(s, i), stag("=="), |i| expr(s, i)), |(a, b)| {
32                Assertion::Eq(a, b)
33            }),
34            map(separated_pair(|i| expr(s, i), stag("<="), |i| expr(s, i)), |(a, b)| {
35                Assertion::LTE(a, b)
36            }),
37            map(separated_pair(|i| expr(s, i), stag(">="), |i| expr(s, i)), |(a, b)| {
38                Assertion::GTE(a, b)
39            }),
40            map(separated_pair(|i| expr(s, i), stag("<"), |i| expr(s, i)), |(a, b)| {
41                Assertion::LT(a, b)
42            }),
43            map(separated_pair(|i| expr(s, i), stag(">"), |i| expr(s, i)), |(a, b)| {
44                Assertion::GT(a, b)
45            }),
46        )),
47        spaces,
48    )
49    .parse(i)
50}
51
52fn expr<'i>(symbol_table: &SymbolScope, i: &'i str) -> R<'i, TDim> {
53    broadcast(symbol_table, i)
54}
55
56fn broadcast<'i>(symbol_table: &SymbolScope, input: &'i str) -> R<'i, TDim> {
57    let s = symbol_table;
58    let (mut input, mut result) = add(s, input)?;
59    while let Ok((i, _)) = stag("#").parse(input) {
60        let (i, next) = map_res(|i| add(s, i), |v| result.clone().broadcast(v)).parse(i)?;
61        (input, result) = (i, next);
62    }
63    Ok((input, result))
64}
65
66macro_rules! bin {
67    ($name: ident, $left: expr, $right: expr, $op: expr, $builder: expr) => {
68        fn $name<'i>(symbol_table: &SymbolScope, input: &'i str) -> R<'i, TDim> {
69            let s = symbol_table;
70            let (input, result) = $left(s, input)?;
71            fold(0.., preceded(stag($op), |i| $right(s, i)), move || result.clone(), $builder)
72                .parse(input)
73        }
74    };
75}
76
77bin!(add, sub, sub, "+", |a, b| a + b);
78bin!(sub, mul, mul, "-", |a, b| a - b);
79bin!(mul, div, div, "*", |a, b| a * b);
80bin!(div, atom, |_s, i| numeric(i), "/", |a, b| a / b);
81
82fn atom<'i>(symbol_table: &SymbolScope, i: &'i str) -> R<'i, TDim> {
83    alt((
84        map(numeric, TDim::Val),
85        map(|i| func(symbol_table, "min", i), TDim::Min),
86        map(|i| func(symbol_table, "max", i), TDim::Max),
87        map(|i| func(symbol_table, "broadcast", i), TDim::Broadcast),
88        // No `floor` arm: TDim is integral, so accepting one here reads a
89        // rational ONNX expression as an integer one. Those are parsed in
90        // tract-onnx's `dim_expr`.
91        map(|i| identifier(symbol_table, i), TDim::Sym),
92        map(pair(recognize(stag("-")), |i| atom(symbol_table, i)), |(_, dim)| dim * -1),
93        delimited(stag("("), |i| expr(symbol_table, i), stag(")")),
94    ))
95    .parse(i)
96}
97
98fn func<'i>(symbol_table: &SymbolScope, name: &'static str, i: &'i str) -> R<'i, Vec<TDim>> {
99    preceded(
100        stag(name),
101        delimited(stag("("), separated_list0(stag(","), |i| expr(symbol_table, i)), stag(")")),
102    )
103    .parse(i)
104}
105
106fn identifier<'i>(symbol_table: &SymbolScope, i: &'i str) -> R<'i, Symbol> {
107    map(
108        recognize(pair(
109            alt((alpha1, tag("_"))),
110            many0(alt((alphanumeric1, tag("_"), tag("."), recognize(pair(tag("/"), alpha1))))),
111        )),
112        |s| symbol_table.sym(s),
113    )
114    .parse(i)
115}
116
117fn numeric(i: &str) -> R<'_, i64> {
118    map_res(digit1, std::str::FromStr::from_str).parse(i)
119}
120
121fn spaces(i: &str) -> R<'_, ()> {
122    map(many0(one_of(" \t\n\r")), |_| ()).parse(i)
123}
124
125fn spaced<'s, O, P>(it: P) -> impl Parser<&'s str, Output = O, Error = VerboseError<&'s str>>
126where
127    P: Parser<&'s str, Output = O, Error = VerboseError<&'s str>>,
128{
129    delimited(spaces, it, spaces)
130}
131
132pub(super) fn stag<'s>(
133    t: &'static str,
134) -> impl Parser<&'s str, Output = &'s str, Error = VerboseError<&'s str>> {
135    spaced(tag(t))
136}
137
138#[cfg(test)]
139mod test {
140    use super::*;
141
142    #[test]
143    fn parse_int() {
144        let table = SymbolScope::default();
145        assert_eq!(parse_tdim(&table, "12").unwrap(), TDim::Val(12));
146        assert_eq!(parse_tdim(&table, "-12").unwrap(), TDim::Val(-12));
147    }
148
149    #[test]
150    fn parse_sym() {
151        let table = SymbolScope::default();
152        assert_eq!(parse_tdim(&table, "x").unwrap(), TDim::Sym(table.sym("x")));
153        assert_eq!(
154            parse_tdim(&table, "-y").unwrap(),
155            TDim::MulInt(-1, Box::new(table.sym("y").into()))
156        );
157    }
158
159    #[test]
160    fn parse_bin() {
161        let table = SymbolScope::default();
162        assert_eq!(parse_tdim(&table, "1+2").unwrap(), 3.into());
163        assert_eq!(parse_tdim(&table, "1-2").unwrap(), (-1).into());
164        assert_eq!(parse_tdim(&table, "1*2").unwrap(), 2.into());
165        assert_eq!(parse_tdim(&table, "1/2").unwrap(), 0.into());
166    }
167
168    #[test]
169    fn parse_prio() {
170        let table = SymbolScope::default();
171        assert_eq!(parse_tdim(&table, "1+2*3").unwrap(), 7.into());
172        assert_eq!(parse_tdim(&table, "1*2+3").unwrap(), 5.into());
173    }
174
175    #[test]
176    fn parse_min() {
177        let table = SymbolScope::default();
178        assert_eq!(
179            parse_tdim(&table, "min(P,S)").unwrap(),
180            TDim::Min(vec!(table.sym("P").into(), table.sym("S").into()))
181        );
182    }
183
184    #[test]
185    fn parse_broadcast_func() {
186        let table = SymbolScope::default();
187        assert_eq!(
188            parse_tdim(&table, "broadcast(P,S)").unwrap(),
189            TDim::Broadcast(vec!(table.sym("P").into(), table.sym("S").into()))
190        );
191    }
192
193    #[test]
194    fn parse_broadcast_display_roundtrip() {
195        let table = SymbolScope::default();
196        let original = TDim::Broadcast(vec![table.sym("P").into(), table.sym("S").into()]);
197        let printed = format!("{original}");
198        let reparsed = parse_tdim(&table, &printed).unwrap();
199        assert_eq!(reparsed, original);
200    }
201
202    #[test]
203    fn parse_inequality_0() {
204        let table = SymbolScope::default();
205        assert_eq!(
206            parse_assertion(&table, "P+S<4096").unwrap(),
207            Assertion::LT(parse_tdim(&table, "P+S").unwrap(), 4096.to_dim())
208        );
209    }
210
211    #[test]
212    fn parse_dot_ids() {
213        let table = SymbolScope::default();
214        assert_eq!(parse_tdim(&table, "dot.0").unwrap(), table.sym("dot.0").into());
215    }
216
217    #[test]
218    fn parse_dot_ids_arith() {
219        let table = SymbolScope::default();
220        assert_eq!(parse_tdim(&table, "dot.0/2").unwrap(), table.sym("dot.0").to_dim() / 2);
221    }
222
223    #[test]
224    fn floor_is_refused_rather_than_silently_dropped() {
225        let table = SymbolScope::default();
226        assert!(parse_tdim(&table, "floor(a)").is_err());
227    }
228
229    #[test]
230    fn parse_slash_ids() {
231        let table = SymbolScope::default();
232        assert_eq!(parse_tdim(&table, "foo/bar").unwrap(), table.sym("foo/bar").into());
233        assert_eq!(parse_tdim(&table, "foo/bar/baz").unwrap(), table.sym("foo/bar/baz").into());
234    }
235
236    #[test]
237    fn parse_slash_ids_arith() {
238        let table = SymbolScope::default();
239        assert_eq!(parse_tdim(&table, "foo/bar/2").unwrap(), table.sym("foo/bar").to_dim() / 2);
240    }
241
242    #[test]
243    fn parse_slash_display_roundtrip() {
244        let table = SymbolScope::default();
245        let original: TDim = table.sym("foo/bar").into();
246        let reparsed = parse_tdim(&table, &format!("{original}")).unwrap();
247        assert_eq!(reparsed, original);
248    }
249}
tract_data/dim/parse.rs

tract_data/dim/
parse.rs
