egglog/ast/
parse.rs

1//! Parse a string into egglog.
2
3use crate::util::INTERNAL_SYMBOL_PREFIX;
4use crate::*;
5use egglog_ast::generic_ast::*;
6use egglog_ast::span::{EgglogSpan, Span, SrcFile};
7use ordered_float::OrderedFloat;
8
9#[macro_export]
10macro_rules! span {
11    () => {
12        Span::Rust(std::sync::Arc::new(RustSpan {
13            file: file!(),
14            line: line!(),
15            column: column!(),
16        }))
17    };
18}
19
20// We do an unidiomatic thing here by using a struct instead of an enum.
21// This is okay because we don't expect client code to respond
22// differently to different parse errors. The benefit of this is that
23// error messages are defined in the same place that they are created,
24// making it easier to improve errors over time.
25#[derive(Debug, Error)]
26pub struct ParseError(pub Span, pub String);
27
28impl std::fmt::Display for ParseError {
29    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
30        write!(f, "{}\nparse error: {}", self.0, self.1)
31    }
32}
33
34macro_rules! error {
35    ($span:expr, $($fmt:tt)*) => {
36        Err(ParseError($span, format!($($fmt)*)))
37    };
38}
39
40pub enum Sexp {
41    // Will never contain `Literal::Unit`, as this
42    // will be parsed as an empty `Sexp::List`.
43    Literal(Literal, Span),
44    Atom(String, Span),
45    List(Vec<Sexp>, Span),
46}
47
48impl Sexp {
49    pub fn span(&self) -> Span {
50        match self {
51            Sexp::Literal(_, span) => span.clone(),
52            Sexp::Atom(_, span) => span.clone(),
53            Sexp::List(_, span) => span.clone(),
54        }
55    }
56
57    pub fn expect_uint<UInt: TryFrom<u64>>(&self, e: &'static str) -> Result<UInt, ParseError> {
58        if let Sexp::Literal(Literal::Int(x), _) = self {
59            if *x >= 0 {
60                if let Ok(v) = (*x as u64).try_into() {
61                    return Ok(v);
62                }
63            }
64        }
65        error!(
66            self.span(),
67            "expected {e} to be a nonnegative integer literal"
68        )
69    }
70
71    pub fn expect_string(&self, e: &'static str) -> Result<String, ParseError> {
72        if let Sexp::Literal(Literal::String(x), _) = self {
73            return Ok(x.to_string());
74        }
75        error!(self.span(), "expected {e} to be a string literal")
76    }
77
78    pub fn expect_atom(&self, e: &'static str) -> Result<String, ParseError> {
79        if let Sexp::Atom(symbol, _) = self {
80            return Ok(symbol.clone());
81        }
82        error!(self.span(), "expected {e}")
83    }
84
85    pub fn expect_list(&self, e: &'static str) -> Result<&[Sexp], ParseError> {
86        if let Sexp::List(sexps, _) = self {
87            return Ok(sexps);
88        }
89        error!(self.span(), "expected {e}")
90    }
91
92    pub fn expect_call(&self, e: &'static str) -> Result<(String, &[Sexp], Span), ParseError> {
93        if let Sexp::List(sexps, span) = self {
94            if let [Sexp::Atom(func, _), args @ ..] = sexps.as_slice() {
95                return Ok((func.clone(), args, span.clone()));
96            }
97        }
98        error!(self.span(), "expected {e}")
99    }
100}
101
102// helper for mapping a function that returns `Result`
103fn map_fallible<T>(
104    slice: &[Sexp],
105    parser: &mut Parser,
106    func: impl Fn(&mut Parser, &Sexp) -> Result<T, ParseError>,
107) -> Result<Vec<T>, ParseError> {
108    slice
109        .iter()
110        .map(|sexp| func(parser, sexp))
111        .collect::<Result<_, _>>()
112}
113
114pub trait Macro<T>: Send + Sync {
115    fn name(&self) -> &str;
116    fn parse(&self, args: &[Sexp], span: Span, parser: &mut Parser) -> Result<T, ParseError>;
117}
118
119pub struct SimpleMacro<T, F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError> + Send + Sync>(
120    String,
121    F,
122);
123
124impl<T, F> SimpleMacro<T, F>
125where
126    F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError> + Send + Sync,
127{
128    pub fn new(head: &str, f: F) -> Self {
129        Self(head.to_owned(), f)
130    }
131}
132
133impl<T, F> Macro<T> for SimpleMacro<T, F>
134where
135    F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError> + Send + Sync,
136{
137    fn name(&self) -> &str {
138        &self.0
139    }
140
141    fn parse(&self, args: &[Sexp], span: Span, parser: &mut Parser) -> Result<T, ParseError> {
142        self.1(args, span, parser)
143    }
144}
145
146#[derive(Clone)]
147pub struct Parser {
148    commands: HashMap<String, Arc<dyn Macro<Vec<Command>>>>,
149    actions: HashMap<String, Arc<dyn Macro<Vec<Action>>>>,
150    exprs: HashMap<String, Arc<dyn Macro<Expr>>>,
151    user_defined: HashSet<String>,
152    pub symbol_gen: SymbolGen,
153}
154
155impl Default for Parser {
156    fn default() -> Self {
157        Self {
158            commands: Default::default(),
159            actions: Default::default(),
160            exprs: Default::default(),
161            user_defined: Default::default(),
162            symbol_gen: SymbolGen::new(INTERNAL_SYMBOL_PREFIX.to_string()),
163        }
164    }
165}
166
167impl Parser {
168    fn ensure_symbol_not_reserved(&self, symbol: &str, span: &Span) -> Result<(), ParseError> {
169        if self.symbol_gen.is_reserved(symbol) {
170            return error!(
171                span.clone(),
172                "symbols starting with '{}' are reserved for egglog internals",
173                self.symbol_gen.reserved_prefix()
174            );
175        }
176        Ok(())
177    }
178
179    pub fn get_program_from_string(
180        &mut self,
181        filename: Option<String>,
182        input: &str,
183    ) -> Result<Vec<Command>, ParseError> {
184        let sexps = all_sexps(SexpParser::new(filename, input))?;
185        let nested: Vec<Vec<_>> = map_fallible(&sexps, self, Self::parse_command)?;
186        Ok(nested.into_iter().flatten().collect())
187    }
188
189    // currently only used for testing, but no reason it couldn't be used elsewhere later
190    pub fn get_expr_from_string(
191        &mut self,
192        filename: Option<String>,
193        input: &str,
194    ) -> Result<Expr, ParseError> {
195        let sexp = sexp(&mut SexpParser::new(filename, input))?;
196        self.parse_expr(&sexp)
197    }
198
199    pub fn add_command_macro(&mut self, ma: Arc<dyn Macro<Vec<Command>>>) {
200        self.commands.insert(ma.name().to_owned(), ma);
201    }
202
203    pub fn add_action_macro(&mut self, ma: Arc<dyn Macro<Vec<Action>>>) {
204        self.actions.insert(ma.name().to_owned(), ma);
205    }
206
207    pub fn add_expr_macro(&mut self, ma: Arc<dyn Macro<Expr>>) {
208        self.exprs.insert(ma.name().to_owned(), ma);
209    }
210
211    pub(crate) fn add_user_defined(&mut self, name: String) -> Result<(), Error> {
212        if self.actions.contains_key(&name)
213            || self.exprs.contains_key(&name)
214            || self.commands.contains_key(&name)
215        {
216            use egglog_ast::span::{RustSpan, Span};
217            return Err(Error::CommandAlreadyExists(name, span!()));
218        }
219        self.user_defined.insert(name);
220        Ok(())
221    }
222
223    pub fn parse_command(&mut self, sexp: &Sexp) -> Result<Vec<Command>, ParseError> {
224        let (head, tail, span) = sexp.expect_call("command")?;
225
226        if let Some(macr0) = self.commands.get(&head).cloned() {
227            return macr0.parse(tail, span, self);
228        }
229
230        // This prevents user-defined commands from being parsed as built-in commands.
231        if self.user_defined.contains(&head) {
232            let args = map_fallible(tail, self, Self::parse_expr)?;
233            return Ok(vec![Command::UserDefined(span, head, args)]);
234        }
235
236        Ok(match head.as_str() {
237            "sort" => match tail {
238                [name] => vec![Command::Sort(span, name.expect_atom("sort name")?, None)],
239                [name, call] => {
240                    let (func, args, _) = call.expect_call("container sort declaration")?;
241                    vec![Command::Sort(
242                        span,
243                        name.expect_atom("sort name")?,
244                        Some((func, map_fallible(args, self, Self::parse_expr)?)),
245                    )]
246                }
247                _ => {
248                    return error!(
249                        span,
250                        "usages:\n(sort <name>)\n(sort <name> (<container sort> <argument sort>*))"
251                    );
252                }
253            },
254            "datatype" => match tail {
255                [name, variants @ ..] => vec![Command::Datatype {
256                    span,
257                    name: name.expect_atom("sort name")?,
258                    variants: map_fallible(variants, self, Self::variant)?,
259                }],
260                _ => return error!(span, "usage: (datatype <name> <variant>*)"),
261            },
262            "datatype*" => vec![Command::Datatypes {
263                span,
264                datatypes: map_fallible(tail, self, Self::rec_datatype)?,
265            }],
266            "function" => match tail {
267                [name, inputs, output, rest @ ..] => vec![Command::Function {
268                    name: name.expect_atom("function name")?,
269                    schema: self.parse_schema(inputs, output)?,
270                    merge: match self.parse_options(rest)?.as_slice() {
271                        [(":no-merge", [])] => None,
272                        [(":merge", [e])] => Some(self.parse_expr(e)?),
273                        [] => {
274                            return error!(
275                                span,
276                                "functions are required to specify merge behaviour"
277                            );
278                        }
279                        _ => return error!(span, "could not parse function options"),
280                    },
281                    span,
282                }],
283                _ => {
284                    let a = "(function <name> (<input sort>*) <output sort> :merge <expr>)";
285                    let b = "(function <name> (<input sort>*) <output sort> :no-merge)";
286                    return error!(span, "usages:\n{a}\n{b}");
287                }
288            },
289            "constructor" => match tail {
290                [name, inputs, output, rest @ ..] => {
291                    let mut cost = None;
292                    let mut unextractable = false;
293                    match self.parse_options(rest)?.as_slice() {
294                        [] => {}
295                        [(":unextractable", [])] => unextractable = true,
296                        [(":cost", [c])] => cost = Some(c.expect_uint("cost")?),
297                        _ => return error!(span, "could not parse constructor options"),
298                    }
299
300                    vec![Command::Constructor {
301                        span,
302                        name: name.expect_atom("constructor name")?,
303                        schema: self.parse_schema(inputs, output)?,
304                        cost,
305                        unextractable,
306                    }]
307                }
308                _ => {
309                    let a = "(constructor <name> (<input sort>*) <output sort>)";
310                    let b = "(constructor <name> (<input sort>*) <output sort> :cost <cost>)";
311                    let c = "(constructor <name> (<input sort>*) <output sort> :unextractable)";
312                    return error!(span, "usages:\n{a}\n{b}\n{c}");
313                }
314            },
315            "relation" => match tail {
316                [name, inputs] => vec![Command::Relation {
317                    span,
318                    name: name.expect_atom("relation name")?,
319                    inputs: map_fallible(inputs.expect_list("input sorts")?, self, |_, sexp| {
320                        sexp.expect_atom("input sort")
321                    })?,
322                }],
323                _ => return error!(span, "usage: (relation <name> (<input sort>*))"),
324            },
325            "ruleset" => match tail {
326                [name] => vec![Command::AddRuleset(span, name.expect_atom("ruleset name")?)],
327                _ => return error!(span, "usage: (ruleset <name>)"),
328            },
329            "unstable-combined-ruleset" => match tail {
330                [name, subrulesets @ ..] => vec![Command::UnstableCombinedRuleset(
331                    span,
332                    name.expect_atom("combined ruleset name")?,
333                    map_fallible(subrulesets, self, |_, sexp| {
334                        sexp.expect_atom("subruleset name")
335                    })?,
336                )],
337                _ => {
338                    return error!(
339                        span,
340                        "usage: (unstable-combined-ruleset <name> <child ruleset>*)"
341                    );
342                }
343            },
344            "rule" => match tail {
345                [lhs, rhs, rest @ ..] => {
346                    let body =
347                        map_fallible(lhs.expect_list("rule query")?, self, Self::parse_fact)?;
348                    let head: Vec<Vec<_>> =
349                        map_fallible(rhs.expect_list("rule actions")?, self, Self::parse_action)?;
350                    let head = GenericActions(head.into_iter().flatten().collect());
351
352                    let mut ruleset = String::new();
353                    let mut name = String::new();
354                    for option in self.parse_options(rest)? {
355                        match option {
356                            (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
357                            (":name", [s]) => name = s.expect_string("rule name")?,
358                            _ => return error!(span, "could not parse rule option"),
359                        }
360                    }
361
362                    vec![Command::Rule {
363                        rule: Rule {
364                            span,
365                            head,
366                            body,
367                            name,
368                            ruleset,
369                        },
370                    }]
371                }
372                _ => return error!(span, "usage: (rule (<fact>*) (<action>*) <option>*)"),
373            },
374            "rewrite" => match tail {
375                [lhs, rhs, rest @ ..] => {
376                    let lhs = self.parse_expr(lhs)?;
377                    let rhs = self.parse_expr(rhs)?;
378
379                    let mut ruleset = String::new();
380                    let mut conditions = Vec::new();
381                    let mut subsume = false;
382                    for option in self.parse_options(rest)? {
383                        match option {
384                            (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
385                            (":subsume", []) => subsume = true,
386                            (":when", [w]) => {
387                                conditions = map_fallible(
388                                    w.expect_list("rewrite conditions")?,
389                                    self,
390                                    Self::parse_fact,
391                                )?
392                            }
393                            _ => return error!(span, "could not parse rewrite options"),
394                        }
395                    }
396
397                    vec![Command::Rewrite(
398                        ruleset,
399                        Rewrite {
400                            span,
401                            lhs,
402                            rhs,
403                            conditions,
404                        },
405                        subsume,
406                    )]
407                }
408                _ => return error!(span, "usage: (rewrite <expr> <expr> <option>*)"),
409            },
410            "birewrite" => match tail {
411                [lhs, rhs, rest @ ..] => {
412                    let lhs = self.parse_expr(lhs)?;
413                    let rhs = self.parse_expr(rhs)?;
414
415                    let mut ruleset = String::new();
416                    let mut conditions = Vec::new();
417                    for option in self.parse_options(rest)? {
418                        match option {
419                            (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
420                            (":when", [w]) => {
421                                conditions = map_fallible(
422                                    w.expect_list("rewrite conditions")?,
423                                    self,
424                                    Self::parse_fact,
425                                )?
426                            }
427                            _ => return error!(span, "could not parse birewrite options"),
428                        }
429                    }
430
431                    vec![Command::BiRewrite(
432                        ruleset,
433                        Rewrite {
434                            span,
435                            lhs,
436                            rhs,
437                            conditions,
438                        },
439                    )]
440                }
441                _ => return error!(span, "usage: (birewrite <expr> <expr> <option>*)"),
442            },
443            "run" => {
444                if tail.is_empty() {
445                    return error!(span, "usage: (run <ruleset>? <uint> <:until (<fact>*)>?)");
446                }
447
448                let has_ruleset = tail.len() >= 2 && tail[1].expect_uint::<u32>("").is_ok();
449
450                let (ruleset, limit, rest) = if has_ruleset {
451                    (
452                        tail[0].expect_atom("ruleset name")?,
453                        tail[1].expect_uint("number of iterations")?,
454                        &tail[2..],
455                    )
456                } else {
457                    (
458                        String::new(),
459                        tail[0].expect_uint("number of iterations")?,
460                        &tail[1..],
461                    )
462                };
463
464                let until = match self.parse_options(rest)?.as_slice() {
465                    [] => None,
466                    [(":until", facts)] => Some(map_fallible(facts, self, Self::parse_fact)?),
467                    _ => return error!(span, "could not parse run options"),
468                };
469
470                vec![Command::RunSchedule(Schedule::Repeat(
471                    span.clone(),
472                    limit,
473                    Box::new(Schedule::Run(span, RunConfig { ruleset, until })),
474                ))]
475            }
476            "run-schedule" => vec![Command::RunSchedule(Schedule::Sequence(
477                span,
478                map_fallible(tail, self, Self::schedule)?,
479            ))],
480            "extract" => match tail {
481                [e] => vec![Command::Extract(
482                    span.clone(),
483                    self.parse_expr(e)?,
484                    Expr::Lit(span, Literal::Int(0)),
485                )],
486                [e, v] => vec![Command::Extract(
487                    span,
488                    self.parse_expr(e)?,
489                    self.parse_expr(v)?,
490                )],
491                _ => return error!(span, "usage: (extract <expr> <number of variants>?)"),
492            },
493            "check" => vec![Command::Check(
494                span,
495                map_fallible(tail, self, Self::parse_fact)?,
496            )],
497            "push" => match tail {
498                [] => vec![Command::Push(1)],
499                [n] => vec![Command::Push(n.expect_uint("number of times to push")?)],
500                _ => return error!(span, "usage: (push <uint>?)"),
501            },
502            "pop" => match tail {
503                [] => vec![Command::Pop(span, 1)],
504                [n] => vec![Command::Pop(span, n.expect_uint("number of times to pop")?)],
505                _ => return error!(span, "usage: (pop <uint>?)"),
506            },
507            "print-stats" => match tail {
508                [] => vec![Command::PrintOverallStatistics(span, None)],
509                [Sexp::Atom(o, _), file] if o == ":file" => vec![Command::PrintOverallStatistics(
510                    span,
511                    Some(file.expect_string("file name")?),
512                )],
513                _ => {
514                    return error!(
515                        span,
516                        "usages: (print-stats)\n(print-stats :file \"<filename>\")"
517                    );
518                }
519            },
520            "print-function" => match tail {
521                [name] => vec![Command::PrintFunction(
522                    span,
523                    name.expect_atom("table name")?,
524                    None,
525                    None,
526                    PrintFunctionMode::Default,
527                )],
528                [name, rest @ ..] => {
529                    let rows: Option<usize> = rest[0].expect_uint("number of rows").ok();
530                    let rest = if rows.is_some() { &rest[1..] } else { rest };
531
532                    let mut file = None;
533                    let mut mode = PrintFunctionMode::Default;
534                    for opt in self.parse_options(rest)? {
535                        match opt {
536                            (":file", [file_name]) => {
537                                file = Some(file_name.expect_string("file name")?);
538                            }
539                            (":mode", [Sexp::Atom(mode_str, _)]) => {
540                                mode = match mode_str.as_str() {
541                                    "default" => PrintFunctionMode::Default,
542                                    "csv" => PrintFunctionMode::CSV,
543                                    _ => {
544                                        return error!(
545                                            span,
546                                            "Unknown print-function mode. Supported modes are `default` and `csv`."
547                                        );
548                                    }
549                                };
550                            }
551                            _ => {
552                                return error!(
553                                    span,
554                                    "Unknown option to print-function. Supported options are `:mode csv|default` and `:file \"<filename>\"`."
555                                );
556                            }
557                        }
558                    }
559                    vec![Command::PrintFunction(
560                        span,
561                        name.expect_atom("table name")?,
562                        rows,
563                        file,
564                        mode,
565                    )]
566                }
567                _ => {
568                    return error!(
569                        span,
570                        "usage: (print-function <table name> <number of rows>? <option>*)"
571                    );
572                }
573            },
574            "print-size" => match tail {
575                [] => vec![Command::PrintSize(span, None)],
576                [name] => vec![Command::PrintSize(
577                    span,
578                    Some(name.expect_atom("table name")?),
579                )],
580                _ => return error!(span, "usage: (print-size <table name>?)"),
581            },
582            "input" => match tail {
583                [name, file] => vec![Command::Input {
584                    span,
585                    name: name.expect_atom("table name")?,
586                    file: file.expect_string("file name")?,
587                }],
588                _ => return error!(span, "usage: (input <table name> \"<file name>\")"),
589            },
590            "output" => match tail {
591                [file, exprs @ ..] => vec![Command::Output {
592                    span,
593                    file: file.expect_string("file name")?,
594                    exprs: map_fallible(exprs, self, Self::parse_expr)?,
595                }],
596                _ => return error!(span, "usage: (output <file name> <expr>+)"),
597            },
598            "include" => match tail {
599                [file] => vec![Command::Include(span, file.expect_string("file name")?)],
600                _ => return error!(span, "usage: (include <file name>)"),
601            },
602            "fail" => match tail {
603                [subcommand] => {
604                    let mut cs = self.parse_command(subcommand)?;
605                    if cs.len() != 1 {
606                        todo!("extend Fail to work with multiple parsed commands")
607                    }
608                    vec![Command::Fail(span, Box::new(cs.remove(0)))]
609                }
610                _ => return error!(span, "usage: (fail <command>)"),
611            },
612            _ => self
613                .parse_action(sexp)?
614                .into_iter()
615                .map(Command::Action)
616                .collect(),
617        })
618    }
619
620    pub fn schedule(&mut self, sexp: &Sexp) -> Result<Schedule, ParseError> {
621        if let Sexp::Atom(ruleset, span) = sexp {
622            return Ok(Schedule::Run(
623                span.clone(),
624                RunConfig {
625                    ruleset: ruleset.clone(),
626                    until: None,
627                },
628            ));
629        }
630
631        let (head, tail, span) = sexp.expect_call("schedule")?;
632
633        Ok(match head.as_str() {
634            "saturate" => Schedule::Saturate(
635                span.clone(),
636                Box::new(Schedule::Sequence(
637                    span,
638                    map_fallible(tail, self, Self::schedule)?,
639                )),
640            ),
641            "seq" => Schedule::Sequence(span, map_fallible(tail, self, Self::schedule)?),
642            "repeat" => match tail {
643                [limit, tail @ ..] => Schedule::Repeat(
644                    span.clone(),
645                    limit.expect_uint("number of iterations")?,
646                    Box::new(Schedule::Sequence(
647                        span,
648                        map_fallible(tail, self, Self::schedule)?,
649                    )),
650                ),
651                _ => return error!(span, "usage: (repeat <number of iterations> <schedule>*)"),
652            },
653            "run" => {
654                let has_ruleset = match tail.first() {
655                    None => false,
656                    Some(Sexp::Atom(o, _)) if *o == ":until" => false,
657                    _ => true,
658                };
659
660                let (ruleset, rest) = if has_ruleset {
661                    (tail[0].expect_atom("ruleset name")?, &tail[1..])
662                } else {
663                    (String::new(), tail)
664                };
665
666                let until = match self.parse_options(rest)?.as_slice() {
667                    [] => None,
668                    [(":until", facts)] => Some(map_fallible(facts, self, Self::parse_fact)?),
669                    _ => return error!(span, "could not parse run options"),
670                };
671
672                Schedule::Run(span, RunConfig { ruleset, until })
673            }
674            _ => return error!(span, "expected either saturate, seq, repeat, or run"),
675        })
676    }
677
678    pub fn parse_action(&mut self, sexp: &Sexp) -> Result<Vec<Action>, ParseError> {
679        let (head, tail, span) = sexp.expect_call("action")?;
680
681        if let Some(func) = self.actions.get(&head).cloned() {
682            return func.parse(tail, span, self);
683        }
684
685        Ok(match head.as_str() {
686            "let" => match tail {
687                [name, value] => {
688                    let binding_span = name.span();
689                    let binding = name.expect_atom("binding name")?;
690                    self.ensure_symbol_not_reserved(&binding, &binding_span)?;
691                    vec![Action::Let(span, binding, self.parse_expr(value)?)]
692                }
693                _ => return error!(span, "usage: (let <name> <expr>)"),
694            },
695            "set" => match tail {
696                [call, value] => {
697                    let (func, args, _) = call.expect_call("table lookup")?;
698                    let args = map_fallible(args, self, Self::parse_expr)?;
699                    let value = self.parse_expr(value)?;
700                    vec![Action::Set(span, func, args, value)]
701                }
702                _ => return error!(span, "usage: (set (<table name> <expr>*) <expr>)"),
703            },
704            "delete" => match tail {
705                [call] => {
706                    let (func, args, _) = call.expect_call("table lookup")?;
707                    let args = map_fallible(args, self, Self::parse_expr)?;
708                    vec![Action::Change(span, Change::Delete, func, args)]
709                }
710                _ => return error!(span, "usage: (delete (<table name> <expr>*))"),
711            },
712            "subsume" => match tail {
713                [call] => {
714                    let (func, args, _) = call.expect_call("table lookup")?;
715                    let args = map_fallible(args, self, Self::parse_expr)?;
716                    vec![Action::Change(span, Change::Subsume, func, args)]
717                }
718                _ => return error!(span, "usage: (subsume (<table name> <expr>*))"),
719            },
720            "union" => match tail {
721                [e1, e2] => vec![Action::Union(
722                    span,
723                    self.parse_expr(e1)?,
724                    self.parse_expr(e2)?,
725                )],
726                _ => return error!(span, "usage: (union <expr> <expr>)"),
727            },
728            "panic" => match tail {
729                [message] => vec![Action::Panic(span, message.expect_string("error message")?)],
730                _ => return error!(span, "usage: (panic <string>)"),
731            },
732            _ => vec![Action::Expr(span, self.parse_expr(sexp)?)],
733        })
734    }
735
736    pub fn parse_fact(&mut self, sexp: &Sexp) -> Result<Fact, ParseError> {
737        let (head, tail, span) = sexp.expect_call("fact")?;
738
739        Ok(match head.as_str() {
740            "=" => match tail {
741                [e1, e2] => Fact::Eq(span, self.parse_expr(e1)?, self.parse_expr(e2)?),
742                _ => return error!(span, "usage: (= <expr> <expr>)"),
743            },
744            _ => Fact::Fact(self.parse_expr(sexp)?),
745        })
746    }
747
748    pub fn parse_expr(&mut self, sexp: &Sexp) -> Result<Expr, ParseError> {
749        Ok(match sexp {
750            Sexp::Literal(literal, span) => Expr::Lit(span.clone(), literal.clone()),
751            Sexp::Atom(symbol, span) => Expr::Var(
752                span.clone(),
753                if *symbol == "_" {
754                    self.symbol_gen.fresh(symbol)
755                } else {
756                    self.ensure_symbol_not_reserved(symbol, span)?;
757                    symbol.clone()
758                },
759            ),
760            Sexp::List(list, span) => match list.as_slice() {
761                [] => Expr::Lit(span.clone(), Literal::Unit),
762                _ => {
763                    let (head, tail, span) = sexp.expect_call("call expression")?;
764
765                    if let Some(func) = self.exprs.get(&head).cloned() {
766                        return func.parse(tail, span, self);
767                    }
768
769                    Expr::Call(
770                        span.clone(),
771                        head,
772                        map_fallible(tail, self, Self::parse_expr)?,
773                    )
774                }
775            },
776        })
777    }
778
779    pub fn rec_datatype(
780        &mut self,
781        sexp: &Sexp,
782    ) -> Result<(Span, String, Subdatatypes), ParseError> {
783        let (head, tail, span) = sexp.expect_call("datatype")?;
784
785        Ok(match head.as_str() {
786            "sort" => match tail {
787                [name, call] => {
788                    let name = name.expect_atom("sort name")?;
789                    let (func, args, _) = call.expect_call("container sort declaration")?;
790                    let args = map_fallible(args, self, Self::parse_expr)?;
791                    (span, name, Subdatatypes::NewSort(func, args))
792                }
793                _ => {
794                    return error!(
795                        span,
796                        "usage: (sort <name> (<container sort> <argument sort>*))"
797                    );
798                }
799            },
800            _ => {
801                let variants = map_fallible(tail, self, Self::variant)?;
802                (span, head, Subdatatypes::Variants(variants))
803            }
804        })
805    }
806
807    pub fn variant(&mut self, sexp: &Sexp) -> Result<Variant, ParseError> {
808        let (name, tail, span) = sexp.expect_call("datatype variant")?;
809
810        let (types, cost, unextractable) = match tail {
811            [types @ .., Sexp::Atom(o, _)] if *o == ":unextractable" => (types, None, true),
812            [types @ .., Sexp::Atom(o, _), c] if *o == ":cost" => {
813                (types, Some(c.expect_uint("cost")?), false)
814            }
815            types => (types, None, false),
816        };
817
818        Ok(Variant {
819            span,
820            name,
821            types: map_fallible(types, self, |_, sexp| {
822                sexp.expect_atom("variant argument type")
823            })?,
824            cost,
825            unextractable,
826        })
827    }
828
829    // helper for parsing a list of options
830    pub fn parse_options<'a>(
831        &self,
832        sexps: &'a [Sexp],
833    ) -> Result<Vec<(&'a str, &'a [Sexp])>, ParseError> {
834        fn option_name(sexp: &Sexp) -> Option<&str> {
835            if let Sexp::Atom(s, _) = sexp {
836                if let Some(':') = s.chars().next() {
837                    return Some(s);
838                }
839            }
840            None
841        }
842
843        let mut out = Vec::new();
844        let mut i = 0;
845        while i < sexps.len() {
846            let Some(key) = option_name(&sexps[i]) else {
847                return error!(sexps[i].span(), "option key must start with ':'");
848            };
849            i += 1;
850
851            let start = i;
852            while i < sexps.len() && option_name(&sexps[i]).is_none() {
853                i += 1;
854            }
855            out.push((key, &sexps[start..i]));
856        }
857        Ok(out)
858    }
859
860    pub fn parse_schema(&self, input: &Sexp, output: &Sexp) -> Result<Schema, ParseError> {
861        Ok(Schema {
862            input: input
863                .expect_list("input sorts")?
864                .iter()
865                .map(|sexp| sexp.expect_atom("input sort"))
866                .collect::<Result<_, _>>()?,
867            output: output.expect_atom("output sort")?,
868        })
869    }
870}
871
872#[derive(Clone, Debug)]
873pub(crate) struct SexpParser {
874    source: Arc<SrcFile>,
875    index: usize,
876}
877
878impl SexpParser {
879    pub(crate) fn new(name: Option<String>, contents: &str) -> SexpParser {
880        SexpParser {
881            source: Arc::new(SrcFile {
882                name,
883                contents: contents.to_string(),
884            }),
885            index: 0,
886        }
887    }
888
889    fn current_char(&self) -> Option<char> {
890        self.source.contents[self.index..].chars().next()
891    }
892
893    fn advance_char(&mut self) {
894        assert!(self.index < self.source.contents.len());
895        loop {
896            self.index += 1;
897            if self.source.contents.is_char_boundary(self.index) {
898                break;
899            }
900        }
901    }
902
903    fn advance_past_whitespace(&mut self) {
904        let mut in_comment = false;
905        loop {
906            match self.current_char() {
907                None => break,
908                Some(';') => in_comment = true,
909                Some('\n') => in_comment = false,
910                Some(c) if c.is_whitespace() => {}
911                Some(_) if in_comment => {}
912                Some(_) => break,
913            }
914            self.advance_char();
915        }
916    }
917
918    fn is_at_end(&self) -> bool {
919        self.index == self.source.contents.len()
920    }
921
922    fn next(&mut self) -> Result<(Token, EgglogSpan), ParseError> {
923        self.advance_past_whitespace();
924        let mut span = EgglogSpan {
925            file: self.source.clone(),
926            i: self.index,
927            j: self.index,
928        };
929
930        let Some(c) = self.current_char() else {
931            return error!(s(span), "unexpected end of file");
932        };
933        self.advance_char();
934
935        let token = match c {
936            '(' => Token::Open,
937            ')' => Token::Close,
938            '"' => {
939                let mut in_escape = false;
940                let mut string = String::new();
941
942                loop {
943                    span.j = self.index;
944                    match self.current_char() {
945                        None => return error!(s(span), "string is missing end quote"),
946                        Some('"') if !in_escape => break,
947                        Some('\\') if !in_escape => in_escape = true,
948                        Some(c) => {
949                            string.push(match (in_escape, c) {
950                                (false, c) => c,
951                                (true, 'n') => '\n',
952                                (true, 't') => '\t',
953                                (true, '\\') => '\\',
954                                (true, '\"') => '\"',
955                                (true, c) => {
956                                    return error!(s(span), "unrecognized escape character {c}");
957                                }
958                            });
959                            in_escape = false;
960                        }
961                    }
962                    self.advance_char();
963                }
964                self.advance_char();
965
966                Token::String(string)
967            }
968            _ => {
969                loop {
970                    match self.current_char() {
971                        Some(c) if c.is_whitespace() => break,
972                        Some(';' | '(' | ')') => break,
973                        None => break,
974                        Some(_) => self.advance_char(),
975                    }
976                }
977                Token::Other
978            }
979        };
980
981        span.j = self.index;
982        self.advance_past_whitespace();
983
984        Ok((token, span))
985    }
986}
987
988fn s(span: EgglogSpan) -> Span {
989    Span::Egglog(Arc::new(span))
990}
991
992enum Token {
993    Open,
994    Close,
995    String(String),
996    Other,
997}
998
999fn sexp(ctx: &mut SexpParser) -> Result<Sexp, ParseError> {
1000    let mut stack: Vec<(EgglogSpan, Vec<Sexp>)> = vec![];
1001
1002    loop {
1003        let (token, span) = ctx.next()?;
1004
1005        let sexp = match token {
1006            Token::Open => {
1007                stack.push((span, vec![]));
1008                continue;
1009            }
1010            Token::Close => {
1011                if stack.is_empty() {
1012                    return error!(s(span), "unexpected `)`");
1013                }
1014                let (mut list_span, list) = stack.pop().unwrap();
1015                list_span.j = span.j;
1016                Sexp::List(list, s(list_span))
1017            }
1018            Token::String(sym) => Sexp::Literal(Literal::String(sym), s(span)),
1019            Token::Other => {
1020                let span = s(span);
1021                let s = span.string();
1022
1023                if s == "true" {
1024                    Sexp::Literal(Literal::Bool(true), span)
1025                } else if s == "false" {
1026                    Sexp::Literal(Literal::Bool(false), span)
1027                } else if let Ok(int) = s.parse::<i64>() {
1028                    Sexp::Literal(Literal::Int(int), span)
1029                } else if s == "NaN" {
1030                    Sexp::Literal(Literal::Float(OrderedFloat(f64::NAN)), span)
1031                } else if s == "inf" {
1032                    Sexp::Literal(Literal::Float(OrderedFloat(f64::INFINITY)), span)
1033                } else if s == "-inf" {
1034                    Sexp::Literal(Literal::Float(OrderedFloat(f64::NEG_INFINITY)), span)
1035                } else if let Ok(float) = s.parse::<f64>() {
1036                    if float.is_finite() {
1037                        Sexp::Literal(Literal::Float(OrderedFloat(float)), span)
1038                    } else {
1039                        Sexp::Atom(s.into(), span)
1040                    }
1041                } else {
1042                    Sexp::Atom(s.into(), span)
1043                }
1044            }
1045        };
1046
1047        if stack.is_empty() {
1048            return Ok(sexp);
1049        } else {
1050            stack.last_mut().unwrap().1.push(sexp);
1051        }
1052    }
1053}
1054
1055pub(crate) fn all_sexps(mut ctx: SexpParser) -> Result<Vec<Sexp>, ParseError> {
1056    let mut sexps = Vec::new();
1057    ctx.advance_past_whitespace();
1058    while !ctx.is_at_end() {
1059        sexps.push(sexp(&mut ctx)?);
1060        ctx.advance_past_whitespace();
1061    }
1062    Ok(sexps)
1063}
1064
1065#[cfg(test)]
1066mod tests {
1067    use super::*;
1068
1069    #[test]
1070    fn test_parser_display_roundtrip() {
1071        let s = r#"(f (g a 3) 4.0 (H "hello"))"#;
1072        let e = Parser::default().get_expr_from_string(None, s).unwrap();
1073        assert_eq!(format!("{}", e), s);
1074    }
1075
1076    #[test]
1077    #[rustfmt::skip]
1078    fn rust_span_display() {
1079        let actual = format!("{}", span!()).replace('\\', "/");
1080        assert!(actual.starts_with("At "));
1081        assert!(actual.contains(":"));
1082        assert!(actual.ends_with("src/ast/parse.rs"));
1083    }
1084
1085    #[test]
1086    fn test_parser_macros() {
1087        let mut parser = Parser::default();
1088        let y = "xxxx";
1089        parser.add_expr_macro(Arc::new(SimpleMacro::new("qqqq", |tail, span, macros| {
1090            Ok(Expr::Call(
1091                span,
1092                y.into(),
1093                map_fallible(tail, macros, Parser::parse_expr)?,
1094            ))
1095        })));
1096        let s = r#"(f (qqqq a 3) 4.0 (H "hello"))"#;
1097        let t = r#"(f (xxxx a 3) 4.0 (H "hello"))"#;
1098        let e = parser.get_expr_from_string(None, s).unwrap();
1099        assert_eq!(format!("{}", e), t);
1100    }
1101}
egglog/ast/parse.rs

egglog/ast/
parse.rs
