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

egglog/ast/
parse.rs
