From eba0e9d2239d22a4fd13cb33c345567f2f49d2df Mon Sep 17 00:00:00 2001
From: Yihong Zhang <yz489@cs.washington.edu>
Date: Sun, 5 Jan 2025 00:02:31 -0800
Subject: [PATCH 1/6] Polish the macro interface

---
 src/ast/desugar.rs |    2 +-
 src/ast/parse.rs   | 1139 +++++++++++++++++++++++---------------------
 src/lib.rs         |    2 +-
 src/termdag.rs     |    4 +-
 4 files changed, 601 insertions(+), 546 deletions(-)

diff --git a/src/ast/desugar.rs b/src/ast/desugar.rs
index 67aa786a..d73c5c44 100644
--- a/src/ast/desugar.rs
+++ b/src/ast/desugar.rs
@@ -112,7 +112,7 @@ pub(crate) fn desugar_command(
             let s = std::fs::read_to_string(&file)
                 .unwrap_or_else(|_| panic!("{span} Failed to read file {file}"));
             return desugar_program(
-                parse_program(Some(file), &s, parser)?,
+                parser.get_program_from_string(Some(file), &s)?,
                 parser,
                 seminaive_transform,
             );
diff --git a/src/ast/parse.rs b/src/ast/parse.rs
index a95a72c2..3c4db857 100644
--- a/src/ast/parse.rs
+++ b/src/ast/parse.rs
@@ -3,47 +3,6 @@
 use crate::*;
 use ordered_float::OrderedFloat;
 
-pub type Macro<T> = fn(&[Sexp], Span, &Parser) -> Result<T, ParseError>;
-
-#[derive(Clone)]
-pub struct Parser {
-    pub commands: HashMap<Symbol, Macro<Vec<Command>>>,
-    pub actions: HashMap<Symbol, Macro<Vec<Action>>>,
-    pub exprs: HashMap<Symbol, Macro<Expr>>,
-    pub symbol_gen: SymbolGen,
-}
-
-impl Default for Parser {
-    fn default() -> Self {
-        Self {
-            commands: Default::default(),
-            actions: Default::default(),
-            exprs: Default::default(),
-            symbol_gen: SymbolGen::new("$".to_string()),
-        }
-    }
-}
-
-pub fn parse_program(
-    filename: Option<String>,
-    input: &str,
-    parser: &Parser,
-) -> Result<Vec<Command>, ParseError> {
-    let sexps = all_sexps(Context::new(filename, input))?;
-    let nested = map_fallible(&sexps, parser, commands)?;
-    Ok(nested.into_iter().flatten().collect())
-}
-
-// currently only used for testing, but no reason it couldn't be used elsewhere later
-pub fn parse_expr(
-    filename: Option<String>,
-    input: &str,
-    parser: &Parser,
-) -> Result<Expr, ParseError> {
-    let sexp = sexp(&mut Context::new(filename, input))?;
-    expr(&sexp, parser)
-}
-
 /// A [`Span`] contains the file name and a pair of offsets representing the start and the end.
 #[derive(Clone, PartialEq, Eq, Hash)]
 pub enum Span {
@@ -250,11 +209,11 @@ impl Sexp {
 fn map_fallible<T>(
     slice: &[Sexp],
     parser: &Parser,
-    func: impl Fn(&Sexp, &Parser) -> Result<T, ParseError>,
+    func: impl Fn(&Parser, &Sexp) -> Result<T, ParseError>,
 ) -> Result<Vec<T>, ParseError> {
     slice
         .iter()
-        .map(|sexp| func(sexp, parser))
+        .map(|sexp| func(parser, sexp))
         .collect::<Result<_, _>>()
 }
 
@@ -287,548 +246,642 @@ fn options(sexps: &[Sexp]) -> Result<Vec<(&str, &[Sexp])>, ParseError> {
     Ok(out)
 }
 
-fn commands(sexp: &Sexp, parser: &Parser) -> Result<Vec<Command>, ParseError> {
-    let (head, tail, span) = sexp.expect_call("command")?;
+fn schema(input: &Sexp, output: &Sexp) -> Result<Schema, ParseError> {
+    Ok(Schema {
+        input: input
+            .expect_list("input sorts")?
+            .iter()
+            .map(|sexp| sexp.expect_atom("input sort"))
+            .collect::<Result<_, _>>()?,
+        output: output.expect_atom("output sort")?,
+    })
+}
+
+pub trait Macro<T> {
+    fn get_head(&self) -> String;
+    fn parse(&self, args: &[Sexp], span: Span, parser: &Parser) -> Result<T, ParseError>;
+}
+
+pub struct SimpleMacro<T, F: Fn(&[Sexp], Span, &Parser) -> Result<T, ParseError>>(String, F);
 
-    if let Some(func) = parser.commands.get(&head) {
-        return func(tail, span, parser);
+impl<T, F> SimpleMacro<T, F>
+where
+    F: Fn(&[Sexp], Span, &Parser) -> Result<T, ParseError>,
+{
+    pub fn new(head: &str, f: F) -> Self {
+        Self(head.into(), f)
     }
+}
 
-    Ok(match head.into() {
-        "set-option" => match tail {
-            [name, value] => vec![Command::SetOption {
-                name: name.expect_atom("option name")?,
-                value: expr(value, parser)?,
-            }],
-            _ => return error!(span, "usage: (set-option <name> <value>)"),
-        },
-        "sort" => match tail {
-            [name] => vec![Command::Sort(span, name.expect_atom("sort name")?, None)],
-            [name, call] => {
-                let (func, args, _) = call.expect_call("container sort declaration")?;
-                vec![Command::Sort(
-                    span,
-                    name.expect_atom("sort name")?,
-                    Some((func, map_fallible(args, parser, expr)?)),
-                )]
-            }
-            _ => {
-                return error!(
-                    span,
-                    "usages:\n(sort <name>)\n(sort <name> (<container sort> <argument sort>*))"
-                )
-            }
-        },
-        "datatype" => match tail {
-            [name, variants @ ..] => vec![Command::Datatype {
-                span,
-                name: name.expect_atom("sort name")?,
-                variants: map_fallible(variants, parser, variant)?,
-            }],
-            _ => return error!(span, "usage: (datatype <name> <variant>*)"),
-        },
-        "datatype*" => vec![Command::Datatypes {
-            span,
-            datatypes: map_fallible(tail, parser, rec_datatype)?,
-        }],
-        "function" => match tail {
-            [name, inputs, output, rest @ ..] => vec![Command::Function {
-                name: name.expect_atom("function name")?,
-                schema: schema(inputs, output)?,
-                merge: match options(rest)?.as_slice() {
-                    [(":no-merge", [])] => None,
-                    [(":merge", [e])] => Some(expr(e, parser)?),
-                    [] => return error!(span, "functions are required to specify merge behaviour"),
-                    _ => return error!(span, "could not parse function options"),
-                },
-                span,
-            }],
-            _ => {
-                let a = "(function <name> (<input sort>*) <output sort> :merge <expr>)";
-                let b = "(function <name> (<input sort>*) <output sort> :no-merge)";
-                return error!(span, "usages:\n{a}\n{b}");
-            }
-        },
-        "constructor" => match tail {
-            [name, inputs, output, rest @ ..] => {
-                let mut cost = None;
-                let mut unextractable = false;
-                match options(rest)?.as_slice() {
-                    [] => {}
-                    [(":unextractable", [])] => unextractable = true,
-                    [(":cost", [c])] => cost = Some(c.expect_uint("cost")?),
-                    _ => return error!(span, "could not parse constructor options"),
-                }
+impl<T, F> Macro<T> for SimpleMacro<T, F>
+where
+    F: Fn(&[Sexp], Span, &Parser) -> Result<T, ParseError>,
+{
+    fn get_head(&self) -> String {
+        self.0.clone()
+    }
+
+    fn parse(&self, args: &[Sexp], span: Span, parser: &Parser) -> Result<T, ParseError> {
+        self.1(args, span, parser)
+    }
+}
+
+#[derive(Clone)]
+pub struct Parser {
+    commands: HashMap<Symbol, Arc<dyn Macro<Vec<Command>>>>,
+    actions: HashMap<Symbol, Arc<dyn Macro<Vec<Action>>>>,
+    exprs: HashMap<Symbol, Arc<dyn Macro<Expr>>>,
+    pub symbol_gen: SymbolGen,
+}
+
+impl Default for Parser {
+    fn default() -> Self {
+        Self {
+            commands: Default::default(),
+            actions: Default::default(),
+            exprs: Default::default(),
+            symbol_gen: SymbolGen::new("$".to_string()),
+        }
+    }
+}
+
+impl Parser {
+    pub fn get_program_from_string(
+        &self,
+        filename: Option<String>,
+        input: &str,
+    ) -> Result<Vec<Command>, ParseError> {
+        let sexps = all_sexps(Context::new(filename, input))?;
+        let nested = map_fallible(&sexps, self, Self::commands)?;
+        Ok(nested.into_iter().flatten().collect())
+    }
+
+    // currently only used for testing, but no reason it couldn't be used elsewhere later
+    pub fn get_expr_from_string(
+        &self,
+        filename: Option<String>,
+        input: &str,
+    ) -> Result<Expr, ParseError> {
+        let sexp = sexp(&mut Context::new(filename, input))?;
+        self.expr(&sexp)
+    }
+
+    pub fn add_command_macro(&mut self, ma: Arc<dyn Macro<Vec<Command>>>) {
+        self.commands.insert(ma.get_head().into(), ma);
+    }
+
+    pub fn add_action_macro(&mut self, ma: Arc<dyn Macro<Vec<Action>>>) {
+        self.actions.insert(ma.get_head().into(), ma);
+    }
+
+    pub fn add_expr_macro(&mut self, ma: Arc<dyn Macro<Expr>>) {
+        self.exprs.insert(ma.get_head().into(), ma);
+    }
+
+    pub fn commands(&self, sexp: &Sexp) -> Result<Vec<Command>, ParseError> {
+        let (head, tail, span) = sexp.expect_call("command")?;
 
-                vec![Command::Constructor {
+        if let Some(macr0) = self.commands.get(&head) {
+            return macr0.parse(tail, span, self);
+        }
+
+        Ok(match head.into() {
+            "set-option" => match tail {
+                [name, value] => vec![Command::SetOption {
+                    name: name.expect_atom("option name")?,
+                    value: self.expr(value)?,
+                }],
+                _ => return error!(span, "usage: (set-option <name> <value>)"),
+            },
+            "sort" => match tail {
+                [name] => vec![Command::Sort(span, name.expect_atom("sort name")?, None)],
+                [name, call] => {
+                    let (func, args, _) = call.expect_call("container sort declaration")?;
+                    vec![Command::Sort(
+                        span,
+                        name.expect_atom("sort name")?,
+                        Some((func, map_fallible(args, self, Self::expr)?)),
+                    )]
+                }
+                _ => {
+                    return error!(
+                        span,
+                        "usages:\n(sort <name>)\n(sort <name> (<container sort> <argument sort>*))"
+                    )
+                }
+            },
+            "datatype" => match tail {
+                [name, variants @ ..] => vec![Command::Datatype {
                     span,
-                    name: name.expect_atom("constructor name")?,
-                    schema: schema(inputs, output)?,
-                    cost,
-                    unextractable,
-                }]
-            }
-            _ => {
-                let a = "(constructor <name> (<input sort>*) <output sort>)";
-                let b = "(constructor <name> (<input sort>*) <output sort> :cost <cost>)";
-                let c = "(constructor <name> (<input sort>*) <output sort> :unextractable)";
-                return error!(span, "usages:\n{a}\n{b}\n{c}");
-            }
-        },
-        "relation" => match tail {
-            [name, inputs] => vec![Command::Relation {
+                    name: name.expect_atom("sort name")?,
+                    variants: map_fallible(variants, self, Self::variant)?,
+                }],
+                _ => return error!(span, "usage: (datatype <name> <variant>*)"),
+            },
+            "datatype*" => vec![Command::Datatypes {
                 span,
-                name: name.expect_atom("relation name")?,
-                inputs: map_fallible(inputs.expect_list("input sorts")?, parser, |sexp, _| {
-                    sexp.expect_atom("input sort")
-                })?,
+                datatypes: map_fallible(tail, self, Self::rec_datatype)?,
             }],
-            _ => return error!(span, "usage: (relation <name> (<input sort>*))"),
-        },
-        "ruleset" => match tail {
-            [name] => vec![Command::AddRuleset(name.expect_atom("ruleset name")?)],
-            _ => return error!(span, "usage: (ruleset <name>)"),
-        },
-        "unstable-combined-ruleset" => match tail {
-            [name, subrulesets @ ..] => vec![Command::UnstableCombinedRuleset(
-                name.expect_atom("combined ruleset name")?,
-                map_fallible(subrulesets, parser, |sexp, _| {
-                    sexp.expect_atom("subruleset name")
-                })?,
-            )],
-            _ => {
-                return error!(
+            "function" => match tail {
+                [name, inputs, output, rest @ ..] => vec![Command::Function {
+                    name: name.expect_atom("function name")?,
+                    schema: schema(inputs, output)?,
+                    merge: match options(rest)?.as_slice() {
+                        [(":no-merge", [])] => None,
+                        [(":merge", [e])] => Some(self.expr(e)?),
+                        [] => {
+                            return error!(
+                                span,
+                                "functions are required to specify merge behaviour"
+                            )
+                        }
+                        _ => return error!(span, "could not parse function options"),
+                    },
                     span,
-                    "usage: (unstable-combined-ruleset <name> <child ruleset>*)"
-                )
-            }
-        },
-        "rule" => match tail {
-            [lhs, rhs, rest @ ..] => {
-                let body = map_fallible(lhs.expect_list("rule query")?, parser, fact)?;
-                let head = map_fallible(rhs.expect_list("rule actions")?, parser, actions)?;
-                let head = GenericActions(head.into_iter().flatten().collect());
-
-                let mut ruleset = "".into();
-                let mut name = "".into();
-                for option in options(rest)? {
-                    match option {
-                        (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
-                        (":name", [s]) => name = s.expect_string("rule name")?.into(),
-                        _ => return error!(span, "could not parse rule option"),
-                    }
+                }],
+                _ => {
+                    let a = "(function <name> (<input sort>*) <output sort> :merge <expr>)";
+                    let b = "(function <name> (<input sort>*) <output sort> :no-merge)";
+                    return error!(span, "usages:\n{a}\n{b}");
                 }
+            },
+            "constructor" => match tail {
+                [name, inputs, output, rest @ ..] => {
+                    let mut cost = None;
+                    let mut unextractable = false;
+                    match options(rest)?.as_slice() {
+                        [] => {}
+                        [(":unextractable", [])] => unextractable = true,
+                        [(":cost", [c])] => cost = Some(c.expect_uint("cost")?),
+                        _ => return error!(span, "could not parse constructor options"),
+                    }
 
-                vec![Command::Rule {
-                    ruleset,
-                    name,
-                    rule: Rule { span, head, body },
-                }]
-            }
-            _ => return error!(span, "usage: (rule (<fact>*) (<action>*) <option>*)"),
-        },
-        "rewrite" => match tail {
-            [lhs, rhs, rest @ ..] => {
-                let lhs = expr(lhs, parser)?;
-                let rhs = expr(rhs, parser)?;
-
-                let mut ruleset = "".into();
-                let mut conditions = Vec::new();
-                let mut subsume = false;
-                for option in options(rest)? {
-                    match option {
-                        (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
-                        (":subsume", []) => subsume = true,
-                        (":when", [w]) => {
-                            conditions =
-                                map_fallible(w.expect_list("rewrite conditions")?, parser, fact)?
+                    vec![Command::Constructor {
+                        span,
+                        name: name.expect_atom("constructor name")?,
+                        schema: schema(inputs, output)?,
+                        cost,
+                        unextractable,
+                    }]
+                }
+                _ => {
+                    let a = "(constructor <name> (<input sort>*) <output sort>)";
+                    let b = "(constructor <name> (<input sort>*) <output sort> :cost <cost>)";
+                    let c = "(constructor <name> (<input sort>*) <output sort> :unextractable)";
+                    return error!(span, "usages:\n{a}\n{b}\n{c}");
+                }
+            },
+            "relation" => match tail {
+                [name, inputs] => vec![Command::Relation {
+                    span,
+                    name: name.expect_atom("relation name")?,
+                    inputs: map_fallible(inputs.expect_list("input sorts")?, self, |_, sexp| {
+                        sexp.expect_atom("input sort")
+                    })?,
+                }],
+                _ => return error!(span, "usage: (relation <name> (<input sort>*))"),
+            },
+            "ruleset" => match tail {
+                [name] => vec![Command::AddRuleset(name.expect_atom("ruleset name")?)],
+                _ => return error!(span, "usage: (ruleset <name>)"),
+            },
+            "unstable-combined-ruleset" => match tail {
+                [name, subrulesets @ ..] => vec![Command::UnstableCombinedRuleset(
+                    name.expect_atom("combined ruleset name")?,
+                    map_fallible(subrulesets, self, |_, sexp| {
+                        sexp.expect_atom("subruleset name")
+                    })?,
+                )],
+                _ => {
+                    return error!(
+                        span,
+                        "usage: (unstable-combined-ruleset <name> <child ruleset>*)"
+                    )
+                }
+            },
+            "rule" => match tail {
+                [lhs, rhs, rest @ ..] => {
+                    let body = map_fallible(lhs.expect_list("rule query")?, self, Self::fact)?;
+                    let head = map_fallible(rhs.expect_list("rule actions")?, self, Self::actions)?;
+                    let head = GenericActions(head.into_iter().flatten().collect());
+
+                    let mut ruleset = "".into();
+                    let mut name = "".into();
+                    for option in options(rest)? {
+                        match option {
+                            (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
+                            (":name", [s]) => name = s.expect_string("rule name")?.into(),
+                            _ => return error!(span, "could not parse rule option"),
                         }
-                        _ => return error!(span, "could not parse rewrite options"),
                     }
+
+                    vec![Command::Rule {
+                        ruleset,
+                        name,
+                        rule: Rule { span, head, body },
+                    }]
                 }
+                _ => return error!(span, "usage: (rule (<fact>*) (<action>*) <option>*)"),
+            },
+            "rewrite" => match tail {
+                [lhs, rhs, rest @ ..] => {
+                    let lhs = self.expr(lhs)?;
+                    let rhs = self.expr(rhs)?;
+
+                    let mut ruleset = "".into();
+                    let mut conditions = Vec::new();
+                    let mut subsume = false;
+                    for option in options(rest)? {
+                        match option {
+                            (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
+                            (":subsume", []) => subsume = true,
+                            (":when", [w]) => {
+                                conditions = map_fallible(
+                                    w.expect_list("rewrite conditions")?,
+                                    self,
+                                    Self::fact,
+                                )?
+                            }
+                            _ => return error!(span, "could not parse rewrite options"),
+                        }
+                    }
 
-                vec![Command::Rewrite(
-                    ruleset,
-                    Rewrite {
-                        span,
-                        lhs,
-                        rhs,
-                        conditions,
-                    },
-                    subsume,
-                )]
-            }
-            _ => return error!(span, "usage: (rewrite <expr> <expr> <option>*)"),
-        },
-        "birewrite" => match tail {
-            [lhs, rhs, rest @ ..] => {
-                let lhs = expr(lhs, parser)?;
-                let rhs = expr(rhs, parser)?;
-
-                let mut ruleset = "".into();
-                let mut conditions = Vec::new();
-                for option in options(rest)? {
-                    match option {
-                        (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
-                        (":when", [w]) => {
-                            conditions =
-                                map_fallible(w.expect_list("rewrite conditions")?, parser, fact)?
+                    vec![Command::Rewrite(
+                        ruleset,
+                        Rewrite {
+                            span,
+                            lhs,
+                            rhs,
+                            conditions,
+                        },
+                        subsume,
+                    )]
+                }
+                _ => return error!(span, "usage: (rewrite <expr> <expr> <option>*)"),
+            },
+            "birewrite" => match tail {
+                [lhs, rhs, rest @ ..] => {
+                    let lhs = self.expr(lhs)?;
+                    let rhs = self.expr(rhs)?;
+
+                    let mut ruleset = "".into();
+                    let mut conditions = Vec::new();
+                    for option in options(rest)? {
+                        match option {
+                            (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
+                            (":when", [w]) => {
+                                conditions = map_fallible(
+                                    w.expect_list("rewrite conditions")?,
+                                    self,
+                                    Self::fact,
+                                )?
+                            }
+                            _ => return error!(span, "could not parse birewrite options"),
                         }
-                        _ => return error!(span, "could not parse birewrite options"),
                     }
+
+                    vec![Command::BiRewrite(
+                        ruleset,
+                        Rewrite {
+                            span,
+                            lhs,
+                            rhs,
+                            conditions,
+                        },
+                    )]
+                }
+                _ => return error!(span, "usage: (birewrite <expr> <expr> <option>*)"),
+            },
+            "run" => {
+                if tail.is_empty() {
+                    return error!(span, "usage: (run <ruleset>? <uint> <:until (<fact>*)>?)");
                 }
 
-                vec![Command::BiRewrite(
-                    ruleset,
-                    Rewrite {
-                        span,
-                        lhs,
-                        rhs,
-                        conditions,
-                    },
-                )]
-            }
-            _ => return error!(span, "usage: (birewrite <expr> <expr> <option>*)"),
-        },
-        "run" => {
-            if tail.is_empty() {
-                return error!(span, "usage: (run <ruleset>? <uint> <:until (<fact>*)>?)");
-            }
+                let has_ruleset = tail.len() >= 2 && tail[1].expect_uint("").is_ok();
 
-            let has_ruleset = tail.len() >= 2 && tail[1].expect_uint("").is_ok();
+                let (ruleset, limit, rest) = if has_ruleset {
+                    (
+                        tail[0].expect_atom("ruleset name")?,
+                        tail[1].expect_uint("number of iterations")?,
+                        &tail[2..],
+                    )
+                } else {
+                    (
+                        "".into(),
+                        tail[0].expect_uint("number of iterations")?,
+                        &tail[1..],
+                    )
+                };
 
-            let (ruleset, limit, rest) = if has_ruleset {
-                (
-                    tail[0].expect_atom("ruleset name")?,
-                    tail[1].expect_uint("number of iterations")?,
-                    &tail[2..],
-                )
-            } else {
-                (
-                    "".into(),
-                    tail[0].expect_uint("number of iterations")?,
-                    &tail[1..],
-                )
-            };
-
-            let until = match options(rest)?.as_slice() {
-                [] => None,
-                [(":until", facts)] => Some(map_fallible(facts, parser, fact)?),
-                _ => return error!(span, "could not parse run options"),
-            };
-
-            vec![Command::RunSchedule(Schedule::Repeat(
-                span.clone(),
-                limit,
-                Box::new(Schedule::Run(span, RunConfig { ruleset, until })),
-            ))]
-        }
-        "run-schedule" => vec![Command::RunSchedule(Schedule::Sequence(
-            span,
-            map_fallible(tail, parser, schedule)?,
-        ))],
-        "simplify" => match tail {
-            [s, e] => vec![Command::Simplify {
-                span,
-                schedule: schedule(s, parser)?,
-                expr: expr(e, parser)?,
-            }],
-            _ => return error!(span, "usage: (simplify <schedule> <expr>)"),
-        },
-        "query-extract" => match tail {
-            [rest @ .., e] => {
-                let variants = match options(rest)?.as_slice() {
-                    [] => 0,
-                    [(":variants", [v])] => v.expect_uint("number of variants")?,
-                    _ => return error!(span, "could not parse query-extract options"),
+                let until = match options(rest)?.as_slice() {
+                    [] => None,
+                    [(":until", facts)] => Some(map_fallible(facts, self, Self::fact)?),
+                    _ => return error!(span, "could not parse run options"),
                 };
-                vec![Command::QueryExtract {
-                    span,
-                    expr: expr(e, parser)?,
-                    variants,
-                }]
+
+                vec![Command::RunSchedule(Schedule::Repeat(
+                    span.clone(),
+                    limit,
+                    Box::new(Schedule::Run(span, RunConfig { ruleset, until })),
+                ))]
             }
-            _ => return error!(span, "usage: (query-extract <:variants <uint>>? <expr>)"),
-        },
-        "check" => vec![Command::Check(span, map_fallible(tail, parser, fact)?)],
-        "push" => match tail {
-            [] => vec![Command::Push(1)],
-            [n] => vec![Command::Push(n.expect_uint("number of times to push")?)],
-            _ => return error!(span, "usage: (push <uint>?)"),
-        },
-        "pop" => match tail {
-            [] => vec![Command::Pop(span, 1)],
-            [n] => vec![Command::Pop(span, n.expect_uint("number of times to pop")?)],
-            _ => return error!(span, "usage: (pop <uint>?)"),
-        },
-        "print-stats" => match tail {
-            [] => vec![Command::PrintOverallStatistics],
-            _ => return error!(span, "usage: (print-stats)"),
-        },
-        "print-function" => match tail {
-            [name, rows] => vec![Command::PrintFunction(
+            "run-schedule" => vec![Command::RunSchedule(Schedule::Sequence(
                 span,
-                name.expect_atom("table name")?,
-                rows.expect_uint("number of rows to print")?,
-            )],
-            _ => {
-                return error!(
+                map_fallible(tail, self, Self::schedule)?,
+            ))],
+            "simplify" => match tail {
+                [s, e] => vec![Command::Simplify {
                     span,
-                    "usage: (print-function <table name> <number of rows>)"
-                )
-            }
-        },
-        "print-size" => match tail {
-            [] => vec![Command::PrintSize(span, None)],
-            [name] => vec![Command::PrintSize(
-                span,
-                Some(name.expect_atom("table name")?),
-            )],
-            _ => return error!(span, "usage: (print-size <table name>?)"),
-        },
-        "input" => match tail {
-            [name, file] => vec![Command::Input {
-                span,
-                name: name.expect_atom("table name")?,
-                file: file.expect_string("file name")?,
-            }],
-            _ => return error!(span, "usage: (input <table name> \"<file name>\")"),
-        },
-        "output" => match tail {
-            [file, exprs @ ..] => vec![Command::Output {
-                span,
-                file: file.expect_string("file name")?,
-                exprs: map_fallible(exprs, parser, expr)?,
-            }],
-            _ => return error!(span, "usage: (output <file name> <expr>+)"),
-        },
-        "include" => match tail {
-            [file] => vec![Command::Include(span, file.expect_string("file name")?)],
-            _ => return error!(span, "usage: (include <file name>)"),
-        },
-        "fail" => match tail {
-            [subcommand] => {
-                let mut cs = commands(subcommand, parser)?;
-                if cs.len() != 1 {
-                    todo!("extend Fail to work with multiple parsed commands")
+                    schedule: self.schedule(s)?,
+                    expr: self.expr(e)?,
+                }],
+                _ => return error!(span, "usage: (simplify <schedule> <expr>)"),
+            },
+            "query-extract" => match tail {
+                [rest @ .., e] => {
+                    let variants = match options(rest)?.as_slice() {
+                        [] => 0,
+                        [(":variants", [v])] => v.expect_uint("number of variants")?,
+                        _ => return error!(span, "could not parse query-extract options"),
+                    };
+                    vec![Command::QueryExtract {
+                        span,
+                        expr: self.expr(e)?,
+                        variants,
+                    }]
                 }
-                vec![Command::Fail(span, Box::new(cs.remove(0)))]
-            }
-            _ => return error!(span, "usage: (fail <command>)"),
-        },
-        _ => actions(sexp, parser)
-            .into_iter()
-            .flatten()
-            .map(Command::Action)
-            .collect(),
-    })
-}
-
-fn schema(input: &Sexp, output: &Sexp) -> Result<Schema, ParseError> {
-    Ok(Schema {
-        input: input
-            .expect_list("input sorts")?
-            .iter()
-            .map(|sexp| sexp.expect_atom("input sort"))
-            .collect::<Result<_, _>>()?,
-        output: output.expect_atom("output sort")?,
-    })
-}
+                _ => return error!(span, "usage: (query-extract <:variants <uint>>? <expr>)"),
+            },
+            "check" => vec![Command::Check(span, map_fallible(tail, self, Self::fact)?)],
+            "push" => match tail {
+                [] => vec![Command::Push(1)],
+                [n] => vec![Command::Push(n.expect_uint("number of times to push")?)],
+                _ => return error!(span, "usage: (push <uint>?)"),
+            },
+            "pop" => match tail {
+                [] => vec![Command::Pop(span, 1)],
+                [n] => vec![Command::Pop(span, n.expect_uint("number of times to pop")?)],
+                _ => return error!(span, "usage: (pop <uint>?)"),
+            },
+            "print-stats" => match tail {
+                [] => vec![Command::PrintOverallStatistics],
+                _ => return error!(span, "usage: (print-stats)"),
+            },
+            "print-function" => match tail {
+                [name, rows] => vec![Command::PrintFunction(
+                    span,
+                    name.expect_atom("table name")?,
+                    rows.expect_uint("number of rows to print")?,
+                )],
+                _ => {
+                    return error!(
+                        span,
+                        "usage: (print-function <table name> <number of rows>)"
+                    )
+                }
+            },
+            "print-size" => match tail {
+                [] => vec![Command::PrintSize(span, None)],
+                [name] => vec![Command::PrintSize(
+                    span,
+                    Some(name.expect_atom("table name")?),
+                )],
+                _ => return error!(span, "usage: (print-size <table name>?)"),
+            },
+            "input" => match tail {
+                [name, file] => vec![Command::Input {
+                    span,
+                    name: name.expect_atom("table name")?,
+                    file: file.expect_string("file name")?,
+                }],
+                _ => return error!(span, "usage: (input <table name> \"<file name>\")"),
+            },
+            "output" => match tail {
+                [file, exprs @ ..] => vec![Command::Output {
+                    span,
+                    file: file.expect_string("file name")?,
+                    exprs: map_fallible(exprs, self, Self::expr)?,
+                }],
+                _ => return error!(span, "usage: (output <file name> <expr>+)"),
+            },
+            "include" => match tail {
+                [file] => vec![Command::Include(span, file.expect_string("file name")?)],
+                _ => return error!(span, "usage: (include <file name>)"),
+            },
+            "fail" => match tail {
+                [subcommand] => {
+                    let mut cs = self.commands(subcommand)?;
+                    if cs.len() != 1 {
+                        todo!("extend Fail to work with multiple parsed commands")
+                    }
+                    vec![Command::Fail(span, Box::new(cs.remove(0)))]
+                }
+                _ => return error!(span, "usage: (fail <command>)"),
+            },
+            _ => self
+                .actions(sexp)
+                .into_iter()
+                .flatten()
+                .map(Command::Action)
+                .collect(),
+        })
+    }
 
-fn rec_datatype(sexp: &Sexp, parser: &Parser) -> Result<(Span, Symbol, Subdatatypes), ParseError> {
-    let (head, tail, span) = sexp.expect_call("datatype")?;
+    pub fn rec_datatype(&self, sexp: &Sexp) -> Result<(Span, Symbol, Subdatatypes), ParseError> {
+        let (head, tail, span) = sexp.expect_call("datatype")?;
 
-    Ok(match head.into() {
-        "sort" => match tail {
-            [name, call] => {
-                let name = name.expect_atom("sort name")?;
-                let (func, args, _) = call.expect_call("container sort declaration")?;
-                let args = map_fallible(args, parser, expr)?;
-                (span, name, Subdatatypes::NewSort(func, args))
-            }
+        Ok(match head.into() {
+            "sort" => match tail {
+                [name, call] => {
+                    let name = name.expect_atom("sort name")?;
+                    let (func, args, _) = call.expect_call("container sort declaration")?;
+                    let args = map_fallible(args, self, Self::expr)?;
+                    (span, name, Subdatatypes::NewSort(func, args))
+                }
+                _ => {
+                    return error!(
+                        span,
+                        "usage: (sort <name> (<container sort> <argument sort>*))"
+                    )
+                }
+            },
             _ => {
-                return error!(
-                    span,
-                    "usage: (sort <name> (<container sort> <argument sort>*))"
-                )
+                let variants = map_fallible(tail, self, Self::variant)?;
+                (span, head, Subdatatypes::Variants(variants))
             }
-        },
-        _ => {
-            let variants = map_fallible(tail, parser, variant)?;
-            (span, head, Subdatatypes::Variants(variants))
-        }
-    })
-}
-
-fn variant(sexp: &Sexp, parser: &Parser) -> Result<Variant, ParseError> {
-    let (name, tail, span) = sexp.expect_call("datatype variant")?;
+        })
+    }
 
-    let (types, cost) = match tail {
-        [types @ .., Sexp::Atom(o, _), c] if *o == ":cost".into() => {
-            (types, Some(c.expect_uint("cost")?))
-        }
-        types => (types, None),
-    };
+    pub fn variant(&self, sexp: &Sexp) -> Result<Variant, ParseError> {
+        let (name, tail, span) = sexp.expect_call("datatype variant")?;
 
-    Ok(Variant {
-        span,
-        name,
-        types: map_fallible(types, parser, |sexp, _| {
-            sexp.expect_atom("variant argument type")
-        })?,
-        cost,
-    })
-}
+        let (types, cost) = match tail {
+            [types @ .., Sexp::Atom(o, _), c] if *o == ":cost".into() => {
+                (types, Some(c.expect_uint("cost")?))
+            }
+            types => (types, None),
+        };
 
-fn schedule(sexp: &Sexp, parser: &Parser) -> Result<Schedule, ParseError> {
-    if let Sexp::Atom(ruleset, span) = sexp {
-        return Ok(Schedule::Run(
-            span.clone(),
-            RunConfig {
-                ruleset: *ruleset,
-                until: None,
-            },
-        ));
+        Ok(Variant {
+            span,
+            name,
+            types: map_fallible(types, self, |_, sexp| {
+                sexp.expect_atom("variant argument type")
+            })?,
+            cost,
+        })
     }
 
-    let (head, tail, span) = sexp.expect_call("schedule")?;
+    pub fn schedule(&self, sexp: &Sexp) -> Result<Schedule, ParseError> {
+        if let Sexp::Atom(ruleset, span) = sexp {
+            return Ok(Schedule::Run(
+                span.clone(),
+                RunConfig {
+                    ruleset: *ruleset,
+                    until: None,
+                },
+            ));
+        }
+
+        let (head, tail, span) = sexp.expect_call("schedule")?;
 
-    Ok(match head.into() {
-        "saturate" => Schedule::Saturate(
-            span.clone(),
-            Box::new(Schedule::Sequence(
-                span,
-                map_fallible(tail, parser, schedule)?,
-            )),
-        ),
-        "seq" => Schedule::Sequence(span, map_fallible(tail, parser, schedule)?),
-        "repeat" => match tail {
-            [limit, tail @ ..] => Schedule::Repeat(
+        Ok(match head.into() {
+            "saturate" => Schedule::Saturate(
                 span.clone(),
-                limit.expect_uint("number of iterations")?,
                 Box::new(Schedule::Sequence(
                     span,
-                    map_fallible(tail, parser, schedule)?,
+                    map_fallible(tail, self, Self::schedule)?,
                 )),
             ),
-            _ => return error!(span, "usage: (repeat <number of iterations> <schedule>*)"),
-        },
-        "run" => {
-            let has_ruleset = match tail.first() {
-                None => false,
-                Some(Sexp::Atom(o, _)) if *o == ":until".into() => false,
-                _ => true,
-            };
-
-            let (ruleset, rest) = if has_ruleset {
-                (tail[0].expect_atom("ruleset name")?, &tail[1..])
-            } else {
-                ("".into(), tail)
-            };
-
-            let until = match options(rest)?.as_slice() {
-                [] => None,
-                [(":until", facts)] => Some(map_fallible(facts, parser, fact)?),
-                _ => return error!(span, "could not parse run options"),
-            };
+            "seq" => Schedule::Sequence(span, map_fallible(tail, self, Self::schedule)?),
+            "repeat" => match tail {
+                [limit, tail @ ..] => Schedule::Repeat(
+                    span.clone(),
+                    limit.expect_uint("number of iterations")?,
+                    Box::new(Schedule::Sequence(
+                        span,
+                        map_fallible(tail, self, Self::schedule)?,
+                    )),
+                ),
+                _ => return error!(span, "usage: (repeat <number of iterations> <schedule>*)"),
+            },
+            "run" => {
+                let has_ruleset = match tail.first() {
+                    None => false,
+                    Some(Sexp::Atom(o, _)) if *o == ":until".into() => false,
+                    _ => true,
+                };
 
-            Schedule::Run(span, RunConfig { ruleset, until })
-        }
-        _ => return error!(span, "expected either saturate, seq, repeat, or run"),
-    })
-}
+                let (ruleset, rest) = if has_ruleset {
+                    (tail[0].expect_atom("ruleset name")?, &tail[1..])
+                } else {
+                    ("".into(), tail)
+                };
 
-fn actions(sexp: &Sexp, parser: &Parser) -> Result<Vec<Action>, ParseError> {
-    let (head, tail, span) = sexp.expect_call("action")?;
+                let until = match options(rest)?.as_slice() {
+                    [] => None,
+                    [(":until", facts)] => Some(map_fallible(facts, self, Self::fact)?),
+                    _ => return error!(span, "could not parse run options"),
+                };
 
-    if let Some(func) = parser.actions.get(&head) {
-        return func(tail, span, parser);
+                Schedule::Run(span, RunConfig { ruleset, until })
+            }
+            _ => return error!(span, "expected either saturate, seq, repeat, or run"),
+        })
     }
 
-    Ok(match head.into() {
-        "let" => match tail {
-            [name, value] => vec![Action::Let(
-                span,
-                name.expect_atom("binding name")?,
-                expr(value, parser)?,
-            )],
-            _ => return error!(span, "usage: (let <name> <expr>)"),
-        },
-        "set" => match tail {
-            [call, value] => {
-                let (func, args, _) = call.expect_call("table lookup")?;
-                let args = map_fallible(args, parser, expr)?;
-                let value = expr(value, parser)?;
-                vec![Action::Set(span, func, args, value)]
-            }
-            _ => return error!(span, "usage: (set (<table name> <expr>*) <expr>)"),
-        },
-        "delete" => match tail {
-            [call] => {
-                let (func, args, _) = call.expect_call("table lookup")?;
-                let args = map_fallible(args, parser, expr)?;
-                vec![Action::Change(span, Change::Delete, func, args)]
-            }
-            _ => return error!(span, "usage: (delete (<table name> <expr>*))"),
-        },
-        "subsume" => match tail {
-            [call] => {
-                let (func, args, _) = call.expect_call("table lookup")?;
-                let args = map_fallible(args, parser, expr)?;
-                vec![Action::Change(span, Change::Subsume, func, args)]
-            }
-            _ => return error!(span, "usage: (subsume (<table name> <expr>*))"),
-        },
-        "union" => match tail {
-            [e1, e2] => vec![Action::Union(span, expr(e1, parser)?, expr(e2, parser)?)],
-            _ => return error!(span, "usage: (union <expr> <expr>)"),
-        },
-        "panic" => match tail {
-            [message] => vec![Action::Panic(span, message.expect_string("error message")?)],
-            _ => return error!(span, "usage: (panic <string>)"),
-        },
-        "extract" => match tail {
-            [e] => vec![Action::Extract(
-                span.clone(),
-                expr(e, parser)?,
-                Expr::Lit(span, Literal::Int(0)),
-            )],
-            [e, v] => vec![Action::Extract(span, expr(e, parser)?, expr(v, parser)?)],
-            _ => return error!(span, "usage: (extract <expr> <number of variants>?)"),
-        },
-        _ => vec![Action::Expr(span, expr(sexp, parser)?)],
-    })
-}
+    pub fn actions(&self, sexp: &Sexp) -> Result<Vec<Action>, ParseError> {
+        let (head, tail, span) = sexp.expect_call("action")?;
 
-fn fact(sexp: &Sexp, parser: &Parser) -> Result<Fact, ParseError> {
-    let (head, tail, span) = sexp.expect_call("fact")?;
+        if let Some(func) = self.actions.get(&head) {
+            return func.parse(tail, span, self);
+        }
 
-    Ok(match head.into() {
-        "=" => match tail {
-            [e1, e2] => Fact::Eq(span, expr(e1, parser)?, expr(e2, parser)?),
-            _ => return error!(span, "usage: (= <expr> <expr>)"),
-        },
-        _ => Fact::Fact(expr(sexp, parser)?),
-    })
-}
+        Ok(match head.into() {
+            "let" => match tail {
+                [name, value] => vec![Action::Let(
+                    span,
+                    name.expect_atom("binding name")?,
+                    self.expr(value)?,
+                )],
+                _ => return error!(span, "usage: (let <name> <expr>)"),
+            },
+            "set" => match tail {
+                [call, value] => {
+                    let (func, args, _) = call.expect_call("table lookup")?;
+                    let args = map_fallible(args, self, Self::expr)?;
+                    let value = self.expr(value)?;
+                    vec![Action::Set(span, func, args, value)]
+                }
+                _ => return error!(span, "usage: (set (<table name> <expr>*) <expr>)"),
+            },
+            "delete" => match tail {
+                [call] => {
+                    let (func, args, _) = call.expect_call("table lookup")?;
+                    let args = map_fallible(args, self, Self::expr)?;
+                    vec![Action::Change(span, Change::Delete, func, args)]
+                }
+                _ => return error!(span, "usage: (delete (<table name> <expr>*))"),
+            },
+            "subsume" => match tail {
+                [call] => {
+                    let (func, args, _) = call.expect_call("table lookup")?;
+                    let args = map_fallible(args, self, Self::expr)?;
+                    vec![Action::Change(span, Change::Subsume, func, args)]
+                }
+                _ => return error!(span, "usage: (subsume (<table name> <expr>*))"),
+            },
+            "union" => match tail {
+                [e1, e2] => vec![Action::Union(span, self.expr(e1)?, self.expr(e2)?)],
+                _ => return error!(span, "usage: (union <expr> <expr>)"),
+            },
+            "panic" => match tail {
+                [message] => vec![Action::Panic(span, message.expect_string("error message")?)],
+                _ => return error!(span, "usage: (panic <string>)"),
+            },
+            "extract" => match tail {
+                [e] => vec![Action::Extract(
+                    span.clone(),
+                    self.expr(e)?,
+                    Expr::Lit(span, Literal::Int(0)),
+                )],
+                [e, v] => vec![Action::Extract(span, self.expr(e)?, self.expr(v)?)],
+                _ => return error!(span, "usage: (extract <expr> <number of variants>?)"),
+            },
+            _ => vec![Action::Expr(span, self.expr(sexp)?)],
+        })
+    }
 
-fn expr(sexp: &Sexp, parser: &Parser) -> Result<Expr, ParseError> {
-    Ok(match sexp {
-        Sexp::Literal(literal, span) => Expr::Lit(span.clone(), literal.clone()),
-        Sexp::Atom(symbol, span) => Expr::Var(span.clone(), *symbol),
-        Sexp::List(list, span) => match list.as_slice() {
-            [] => Expr::Lit(span.clone(), Literal::Unit),
-            _ => {
-                let (head, tail, span) = sexp.expect_call("call expression")?;
+    fn fact(&self, sexp: &Sexp) -> Result<Fact, ParseError> {
+        let (head, tail, span) = sexp.expect_call("fact")?;
 
-                if let Some(func) = parser.exprs.get(&head) {
-                    return func(tail, span, parser);
-                }
+        Ok(match head.into() {
+            "=" => match tail {
+                [e1, e2] => Fact::Eq(span, self.expr(e1)?, self.expr(e2)?),
+                _ => return error!(span, "usage: (= <expr> <expr>)"),
+            },
+            _ => Fact::Fact(self.expr(sexp)?),
+        })
+    }
 
-                Expr::Call(span.clone(), head, map_fallible(tail, parser, expr)?)
-            }
-        },
-    })
+    fn expr(&self, sexp: &Sexp) -> Result<Expr, ParseError> {
+        Ok(match sexp {
+            Sexp::Literal(literal, span) => Expr::Lit(span.clone(), literal.clone()),
+            Sexp::Atom(symbol, span) => Expr::Var(span.clone(), *symbol),
+            Sexp::List(list, span) => match list.as_slice() {
+                [] => Expr::Lit(span.clone(), Literal::Unit),
+                _ => {
+                    let (head, tail, span) = sexp.expect_call("call expression")?;
+
+                    if let Some(func) = self.exprs.get(&head) {
+                        return func.parse(tail, span, self);
+                    }
+
+                    Expr::Call(span.clone(), head, map_fallible(tail, self, Self::expr)?)
+                }
+            },
+        })
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -1028,28 +1081,30 @@ mod tests {
     #[test]
     fn test_parser_display_roundtrip() {
         let s = r#"(f (g a 3) 4.0 (H "hello"))"#;
-        let e = crate::ast::parse_expr(None, s, &Default::default()).unwrap();
+        let e = Parser::default().get_expr_from_string(None, s).unwrap();
         assert_eq!(format!("{}", e), s);
     }
 
     #[test]
+    #[rustfmt::skip]
     fn rust_span_display() {
-        assert_eq!(format!("{}", span!()), "At 1037:34 of src/ast/parse.rs");
+        assert_eq!(format!("{}", span!()), format!("At {}:34 of src/ast/parse.rs", line!()));
     }
 
     #[test]
     fn test_parser_macros() {
         let mut parser = Parser::default();
-        parser.exprs.insert("qqqq".into(), |tail, span, macros| {
+        let y = "xxxx";
+        parser.add_expr_macro(Arc::new(SimpleMacro::new("qqqq", |tail, span, macros| {
             Ok(Expr::Call(
                 span,
-                "xxxx".into(),
-                map_fallible(tail, macros, expr)?,
+                y.into(),
+                map_fallible(tail, macros, Parser::expr)?,
             ))
-        });
+        })));
         let s = r#"(f (qqqq a 3) 4.0 (H "hello"))"#;
         let t = r#"(f (xxxx a 3) 4.0 (H "hello"))"#;
-        let e = crate::ast::parse_expr(None, s, &parser).unwrap();
+        let e = parser.get_expr_from_string(None, s).unwrap();
         assert_eq!(format!("{}", e), t);
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 0ed949dc..42bfaaac 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1472,7 +1472,7 @@ impl EGraph {
         filename: Option<String>,
         input: &str,
     ) -> Result<Vec<String>, Error> {
-        let parsed = parse_program(filename, input, &self.parser)?;
+        let parsed = self.parser.get_program_from_string(filename, input)?;
         self.run_program(parsed)
     }
 
diff --git a/src/termdag.rs b/src/termdag.rs
index cee64fe9..0ce8d405 100644
--- a/src/termdag.rs
+++ b/src/termdag.rs
@@ -194,7 +194,7 @@ mod tests {
     use crate::{ast::*, span};
 
     fn parse_term(s: &str) -> (TermDag, Term) {
-        let e = parse_expr(None, s, &Default::default()).unwrap();
+        let e = Parser::default().get_expr_from_string(None, s).unwrap();
         let mut td = TermDag::default();
         let t = td.expr_to_term(&e);
         (td, t)
@@ -203,7 +203,7 @@ mod tests {
     #[test]
     fn test_to_from_expr() {
         let s = r#"(f (g x y) x y (g x y))"#;
-        let e = parse_expr(None, s, &Default::default()).unwrap();
+        let e = Parser::default().get_expr_from_string(None, s).unwrap();
         let mut td = TermDag::default();
         assert_eq!(td.size(), 0);
         let t = td.expr_to_term(&e);

From 01a4a614b620fefad95b154d270f607105435fce Mon Sep 17 00:00:00 2001
From: Yihong Zhang <yz489@cs.washington.edu>
Date: Sun, 5 Jan 2025 01:57:31 -0800
Subject: [PATCH 2/6] make more methods public make macro take a mutable parser
 for symbol_gen rename methods

---
 src/ast/parse.rs | 227 ++++++++++++++++++++++++++---------------------
 1 file changed, 125 insertions(+), 102 deletions(-)

diff --git a/src/ast/parse.rs b/src/ast/parse.rs
index 3c4db857..3fd12a02 100644
--- a/src/ast/parse.rs
+++ b/src/ast/parse.rs
@@ -133,6 +133,12 @@ impl Display for Span {
 #[derive(Debug, Error)]
 pub struct ParseError(Span, String);
 
+impl ParseError {
+    pub fn new(span: Span, message: &str) -> Self {
+        Self(span, message.to_string())
+    }
+}
+
 impl std::fmt::Display for ParseError {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         write!(f, "{}\nparse error: {}", self.0, self.1)
@@ -154,7 +160,7 @@ pub enum Sexp {
 }
 
 impl Sexp {
-    fn span(&self) -> Span {
+    pub fn span(&self) -> Span {
         match self {
             Sexp::Literal(_, span) => span.clone(),
             Sexp::Atom(_, span) => span.clone(),
@@ -162,7 +168,7 @@ impl Sexp {
         }
     }
 
-    fn expect_uint(&self, e: &'static str) -> Result<usize, ParseError> {
+    pub fn expect_uint(&self, e: &'static str) -> Result<usize, ParseError> {
         if let Sexp::Literal(Literal::Int(x), _) = self {
             if *x >= 0 {
                 return Ok(*x as usize);
@@ -174,28 +180,28 @@ impl Sexp {
         )
     }
 
-    fn expect_string(&self, e: &'static str) -> Result<String, ParseError> {
+    pub fn expect_string(&self, e: &'static str) -> Result<String, ParseError> {
         if let Sexp::Literal(Literal::String(x), _) = self {
             return Ok(x.to_string());
         }
         error!(self.span(), "expected {e} to be a string literal")
     }
 
-    fn expect_atom(&self, e: &'static str) -> Result<Symbol, ParseError> {
+    pub fn expect_atom(&self, e: &'static str) -> Result<Symbol, ParseError> {
         if let Sexp::Atom(symbol, _) = self {
             return Ok(*symbol);
         }
         error!(self.span(), "expected {e}")
     }
 
-    fn expect_list(&self, e: &'static str) -> Result<&[Sexp], ParseError> {
+    pub fn expect_list(&self, e: &'static str) -> Result<&[Sexp], ParseError> {
         if let Sexp::List(sexps, _) = self {
             return Ok(sexps);
         }
         error!(self.span(), "expected {e}")
     }
 
-    fn expect_call(&self, e: &'static str) -> Result<(Symbol, &[Sexp], Span), ParseError> {
+    pub fn expect_call(&self, e: &'static str) -> Result<(Symbol, &[Sexp], Span), ParseError> {
         if let Sexp::List(sexps, span) = self {
             if let [Sexp::Atom(func, _), args @ ..] = sexps.as_slice() {
                 return Ok((*func, args, span.clone()));
@@ -208,8 +214,8 @@ impl Sexp {
 // helper for mapping a function that returns `Result`
 fn map_fallible<T>(
     slice: &[Sexp],
-    parser: &Parser,
-    func: impl Fn(&Parser, &Sexp) -> Result<T, ParseError>,
+    parser: &mut Parser,
+    func: impl Fn(&mut Parser, &Sexp) -> Result<T, ParseError>,
 ) -> Result<Vec<T>, ParseError> {
     slice
         .iter()
@@ -259,14 +265,14 @@ fn schema(input: &Sexp, output: &Sexp) -> Result<Schema, ParseError> {
 
 pub trait Macro<T> {
     fn get_head(&self) -> String;
-    fn parse(&self, args: &[Sexp], span: Span, parser: &Parser) -> Result<T, ParseError>;
+    fn parse(&self, args: &[Sexp], span: Span, parser: &mut Parser) -> Result<T, ParseError>;
 }
 
-pub struct SimpleMacro<T, F: Fn(&[Sexp], Span, &Parser) -> Result<T, ParseError>>(String, F);
+pub struct SimpleMacro<T, F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError>>(String, F);
 
 impl<T, F> SimpleMacro<T, F>
 where
-    F: Fn(&[Sexp], Span, &Parser) -> Result<T, ParseError>,
+    F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError>,
 {
     pub fn new(head: &str, f: F) -> Self {
         Self(head.into(), f)
@@ -275,13 +281,13 @@ where
 
 impl<T, F> Macro<T> for SimpleMacro<T, F>
 where
-    F: Fn(&[Sexp], Span, &Parser) -> Result<T, ParseError>,
+    F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError>,
 {
     fn get_head(&self) -> String {
         self.0.clone()
     }
 
-    fn parse(&self, args: &[Sexp], span: Span, parser: &Parser) -> Result<T, ParseError> {
+    fn parse(&self, args: &[Sexp], span: Span, parser: &mut Parser) -> Result<T, ParseError> {
         self.1(args, span, parser)
     }
 }
@@ -307,23 +313,23 @@ impl Default for Parser {
 
 impl Parser {
     pub fn get_program_from_string(
-        &self,
+        &mut self,
         filename: Option<String>,
         input: &str,
     ) -> Result<Vec<Command>, ParseError> {
         let sexps = all_sexps(Context::new(filename, input))?;
-        let nested = map_fallible(&sexps, self, Self::commands)?;
+        let nested = map_fallible(&sexps, self, Self::parse_command)?;
         Ok(nested.into_iter().flatten().collect())
     }
 
     // currently only used for testing, but no reason it couldn't be used elsewhere later
     pub fn get_expr_from_string(
-        &self,
+        &mut self,
         filename: Option<String>,
         input: &str,
     ) -> Result<Expr, ParseError> {
         let sexp = sexp(&mut Context::new(filename, input))?;
-        self.expr(&sexp)
+        self.parse_expr(&sexp)
     }
 
     pub fn add_command_macro(&mut self, ma: Arc<dyn Macro<Vec<Command>>>) {
@@ -338,10 +344,10 @@ impl Parser {
         self.exprs.insert(ma.get_head().into(), ma);
     }
 
-    pub fn commands(&self, sexp: &Sexp) -> Result<Vec<Command>, ParseError> {
+    pub fn parse_command(&mut self, sexp: &Sexp) -> Result<Vec<Command>, ParseError> {
         let (head, tail, span) = sexp.expect_call("command")?;
 
-        if let Some(macr0) = self.commands.get(&head) {
+        if let Some(macr0) = self.commands.get(&head).cloned() {
             return macr0.parse(tail, span, self);
         }
 
@@ -349,7 +355,7 @@ impl Parser {
             "set-option" => match tail {
                 [name, value] => vec![Command::SetOption {
                     name: name.expect_atom("option name")?,
-                    value: self.expr(value)?,
+                    value: self.parse_expr(value)?,
                 }],
                 _ => return error!(span, "usage: (set-option <name> <value>)"),
             },
@@ -360,7 +366,7 @@ impl Parser {
                     vec![Command::Sort(
                         span,
                         name.expect_atom("sort name")?,
-                        Some((func, map_fallible(args, self, Self::expr)?)),
+                        Some((func, map_fallible(args, self, Self::parse_expr)?)),
                     )]
                 }
                 _ => {
@@ -388,7 +394,7 @@ impl Parser {
                     schema: schema(inputs, output)?,
                     merge: match options(rest)?.as_slice() {
                         [(":no-merge", [])] => None,
-                        [(":merge", [e])] => Some(self.expr(e)?),
+                        [(":merge", [e])] => Some(self.parse_expr(e)?),
                         [] => {
                             return error!(
                                 span,
@@ -461,8 +467,10 @@ impl Parser {
             },
             "rule" => match tail {
                 [lhs, rhs, rest @ ..] => {
-                    let body = map_fallible(lhs.expect_list("rule query")?, self, Self::fact)?;
-                    let head = map_fallible(rhs.expect_list("rule actions")?, self, Self::actions)?;
+                    let body =
+                        map_fallible(lhs.expect_list("rule query")?, self, Self::parse_fact)?;
+                    let head =
+                        map_fallible(rhs.expect_list("rule actions")?, self, Self::parse_action)?;
                     let head = GenericActions(head.into_iter().flatten().collect());
 
                     let mut ruleset = "".into();
@@ -485,8 +493,8 @@ impl Parser {
             },
             "rewrite" => match tail {
                 [lhs, rhs, rest @ ..] => {
-                    let lhs = self.expr(lhs)?;
-                    let rhs = self.expr(rhs)?;
+                    let lhs = self.parse_expr(lhs)?;
+                    let rhs = self.parse_expr(rhs)?;
 
                     let mut ruleset = "".into();
                     let mut conditions = Vec::new();
@@ -499,7 +507,7 @@ impl Parser {
                                 conditions = map_fallible(
                                     w.expect_list("rewrite conditions")?,
                                     self,
-                                    Self::fact,
+                                    Self::parse_fact,
                                 )?
                             }
                             _ => return error!(span, "could not parse rewrite options"),
@@ -521,8 +529,8 @@ impl Parser {
             },
             "birewrite" => match tail {
                 [lhs, rhs, rest @ ..] => {
-                    let lhs = self.expr(lhs)?;
-                    let rhs = self.expr(rhs)?;
+                    let lhs = self.parse_expr(lhs)?;
+                    let rhs = self.parse_expr(rhs)?;
 
                     let mut ruleset = "".into();
                     let mut conditions = Vec::new();
@@ -533,7 +541,7 @@ impl Parser {
                                 conditions = map_fallible(
                                     w.expect_list("rewrite conditions")?,
                                     self,
-                                    Self::fact,
+                                    Self::parse_fact,
                                 )?
                             }
                             _ => return error!(span, "could not parse birewrite options"),
@@ -575,7 +583,7 @@ impl Parser {
 
                 let until = match options(rest)?.as_slice() {
                     [] => None,
-                    [(":until", facts)] => Some(map_fallible(facts, self, Self::fact)?),
+                    [(":until", facts)] => Some(map_fallible(facts, self, Self::parse_fact)?),
                     _ => return error!(span, "could not parse run options"),
                 };
 
@@ -593,7 +601,7 @@ impl Parser {
                 [s, e] => vec![Command::Simplify {
                     span,
                     schedule: self.schedule(s)?,
-                    expr: self.expr(e)?,
+                    expr: self.parse_expr(e)?,
                 }],
                 _ => return error!(span, "usage: (simplify <schedule> <expr>)"),
             },
@@ -606,13 +614,16 @@ impl Parser {
                     };
                     vec![Command::QueryExtract {
                         span,
-                        expr: self.expr(e)?,
+                        expr: self.parse_expr(e)?,
                         variants,
                     }]
                 }
                 _ => return error!(span, "usage: (query-extract <:variants <uint>>? <expr>)"),
             },
-            "check" => vec![Command::Check(span, map_fallible(tail, self, Self::fact)?)],
+            "check" => vec![Command::Check(
+                span,
+                map_fallible(tail, self, Self::parse_fact)?,
+            )],
             "push" => match tail {
                 [] => vec![Command::Push(1)],
                 [n] => vec![Command::Push(n.expect_uint("number of times to push")?)],
@@ -660,7 +671,7 @@ impl Parser {
                 [file, exprs @ ..] => vec![Command::Output {
                     span,
                     file: file.expect_string("file name")?,
-                    exprs: map_fallible(exprs, self, Self::expr)?,
+                    exprs: map_fallible(exprs, self, Self::parse_expr)?,
                 }],
                 _ => return error!(span, "usage: (output <file name> <expr>+)"),
             },
@@ -670,7 +681,7 @@ impl Parser {
             },
             "fail" => match tail {
                 [subcommand] => {
-                    let mut cs = self.commands(subcommand)?;
+                    let mut cs = self.parse_command(subcommand)?;
                     if cs.len() != 1 {
                         todo!("extend Fail to work with multiple parsed commands")
                     }
@@ -679,7 +690,7 @@ impl Parser {
                 _ => return error!(span, "usage: (fail <command>)"),
             },
             _ => self
-                .actions(sexp)
+                .parse_action(sexp)
                 .into_iter()
                 .flatten()
                 .map(Command::Action)
@@ -687,52 +698,7 @@ impl Parser {
         })
     }
 
-    pub fn rec_datatype(&self, sexp: &Sexp) -> Result<(Span, Symbol, Subdatatypes), ParseError> {
-        let (head, tail, span) = sexp.expect_call("datatype")?;
-
-        Ok(match head.into() {
-            "sort" => match tail {
-                [name, call] => {
-                    let name = name.expect_atom("sort name")?;
-                    let (func, args, _) = call.expect_call("container sort declaration")?;
-                    let args = map_fallible(args, self, Self::expr)?;
-                    (span, name, Subdatatypes::NewSort(func, args))
-                }
-                _ => {
-                    return error!(
-                        span,
-                        "usage: (sort <name> (<container sort> <argument sort>*))"
-                    )
-                }
-            },
-            _ => {
-                let variants = map_fallible(tail, self, Self::variant)?;
-                (span, head, Subdatatypes::Variants(variants))
-            }
-        })
-    }
-
-    pub fn variant(&self, sexp: &Sexp) -> Result<Variant, ParseError> {
-        let (name, tail, span) = sexp.expect_call("datatype variant")?;
-
-        let (types, cost) = match tail {
-            [types @ .., Sexp::Atom(o, _), c] if *o == ":cost".into() => {
-                (types, Some(c.expect_uint("cost")?))
-            }
-            types => (types, None),
-        };
-
-        Ok(Variant {
-            span,
-            name,
-            types: map_fallible(types, self, |_, sexp| {
-                sexp.expect_atom("variant argument type")
-            })?,
-            cost,
-        })
-    }
-
-    pub fn schedule(&self, sexp: &Sexp) -> Result<Schedule, ParseError> {
+    pub fn schedule(&mut self, sexp: &Sexp) -> Result<Schedule, ParseError> {
         if let Sexp::Atom(ruleset, span) = sexp {
             return Ok(Schedule::Run(
                 span.clone(),
@@ -780,7 +746,7 @@ impl Parser {
 
                 let until = match options(rest)?.as_slice() {
                     [] => None,
-                    [(":until", facts)] => Some(map_fallible(facts, self, Self::fact)?),
+                    [(":until", facts)] => Some(map_fallible(facts, self, Self::parse_fact)?),
                     _ => return error!(span, "could not parse run options"),
                 };
 
@@ -790,10 +756,10 @@ impl Parser {
         })
     }
 
-    pub fn actions(&self, sexp: &Sexp) -> Result<Vec<Action>, ParseError> {
+    pub fn parse_action(&mut self, sexp: &Sexp) -> Result<Vec<Action>, ParseError> {
         let (head, tail, span) = sexp.expect_call("action")?;
 
-        if let Some(func) = self.actions.get(&head) {
+        if let Some(func) = self.actions.get(&head).cloned() {
             return func.parse(tail, span, self);
         }
 
@@ -802,15 +768,15 @@ impl Parser {
                 [name, value] => vec![Action::Let(
                     span,
                     name.expect_atom("binding name")?,
-                    self.expr(value)?,
+                    self.parse_expr(value)?,
                 )],
                 _ => return error!(span, "usage: (let <name> <expr>)"),
             },
             "set" => match tail {
                 [call, value] => {
                     let (func, args, _) = call.expect_call("table lookup")?;
-                    let args = map_fallible(args, self, Self::expr)?;
-                    let value = self.expr(value)?;
+                    let args = map_fallible(args, self, Self::parse_expr)?;
+                    let value = self.parse_expr(value)?;
                     vec![Action::Set(span, func, args, value)]
                 }
                 _ => return error!(span, "usage: (set (<table name> <expr>*) <expr>)"),
@@ -818,7 +784,7 @@ impl Parser {
             "delete" => match tail {
                 [call] => {
                     let (func, args, _) = call.expect_call("table lookup")?;
-                    let args = map_fallible(args, self, Self::expr)?;
+                    let args = map_fallible(args, self, Self::parse_expr)?;
                     vec![Action::Change(span, Change::Delete, func, args)]
                 }
                 _ => return error!(span, "usage: (delete (<table name> <expr>*))"),
@@ -826,13 +792,17 @@ impl Parser {
             "subsume" => match tail {
                 [call] => {
                     let (func, args, _) = call.expect_call("table lookup")?;
-                    let args = map_fallible(args, self, Self::expr)?;
+                    let args = map_fallible(args, self, Self::parse_expr)?;
                     vec![Action::Change(span, Change::Subsume, func, args)]
                 }
                 _ => return error!(span, "usage: (subsume (<table name> <expr>*))"),
             },
             "union" => match tail {
-                [e1, e2] => vec![Action::Union(span, self.expr(e1)?, self.expr(e2)?)],
+                [e1, e2] => vec![Action::Union(
+                    span,
+                    self.parse_expr(e1)?,
+                    self.parse_expr(e2)?,
+                )],
                 _ => return error!(span, "usage: (union <expr> <expr>)"),
             },
             "panic" => match tail {
@@ -842,29 +812,33 @@ impl Parser {
             "extract" => match tail {
                 [e] => vec![Action::Extract(
                     span.clone(),
-                    self.expr(e)?,
+                    self.parse_expr(e)?,
                     Expr::Lit(span, Literal::Int(0)),
                 )],
-                [e, v] => vec![Action::Extract(span, self.expr(e)?, self.expr(v)?)],
+                [e, v] => vec![Action::Extract(
+                    span,
+                    self.parse_expr(e)?,
+                    self.parse_expr(v)?,
+                )],
                 _ => return error!(span, "usage: (extract <expr> <number of variants>?)"),
             },
-            _ => vec![Action::Expr(span, self.expr(sexp)?)],
+            _ => vec![Action::Expr(span, self.parse_expr(sexp)?)],
         })
     }
 
-    fn fact(&self, sexp: &Sexp) -> Result<Fact, ParseError> {
+    pub fn parse_fact(&mut self, sexp: &Sexp) -> Result<Fact, ParseError> {
         let (head, tail, span) = sexp.expect_call("fact")?;
 
         Ok(match head.into() {
             "=" => match tail {
-                [e1, e2] => Fact::Eq(span, self.expr(e1)?, self.expr(e2)?),
+                [e1, e2] => Fact::Eq(span, self.parse_expr(e1)?, self.parse_expr(e2)?),
                 _ => return error!(span, "usage: (= <expr> <expr>)"),
             },
-            _ => Fact::Fact(self.expr(sexp)?),
+            _ => Fact::Fact(self.parse_expr(sexp)?),
         })
     }
 
-    fn expr(&self, sexp: &Sexp) -> Result<Expr, ParseError> {
+    pub fn parse_expr(&mut self, sexp: &Sexp) -> Result<Expr, ParseError> {
         Ok(match sexp {
             Sexp::Literal(literal, span) => Expr::Lit(span.clone(), literal.clone()),
             Sexp::Atom(symbol, span) => Expr::Var(span.clone(), *symbol),
@@ -873,13 +847,62 @@ impl Parser {
                 _ => {
                     let (head, tail, span) = sexp.expect_call("call expression")?;
 
-                    if let Some(func) = self.exprs.get(&head) {
+                    if let Some(func) = self.exprs.get(&head).cloned() {
                         return func.parse(tail, span, self);
                     }
 
-                    Expr::Call(span.clone(), head, map_fallible(tail, self, Self::expr)?)
+                    Expr::Call(
+                        span.clone(),
+                        head,
+                        map_fallible(tail, self, Self::parse_expr)?,
+                    )
+                }
+            },
+        })
+    }
+
+    fn rec_datatype(&mut self, sexp: &Sexp) -> Result<(Span, Symbol, Subdatatypes), ParseError> {
+        let (head, tail, span) = sexp.expect_call("datatype")?;
+
+        Ok(match head.into() {
+            "sort" => match tail {
+                [name, call] => {
+                    let name = name.expect_atom("sort name")?;
+                    let (func, args, _) = call.expect_call("container sort declaration")?;
+                    let args = map_fallible(args, self, Self::parse_expr)?;
+                    (span, name, Subdatatypes::NewSort(func, args))
+                }
+                _ => {
+                    return error!(
+                        span,
+                        "usage: (sort <name> (<container sort> <argument sort>*))"
+                    )
                 }
             },
+            _ => {
+                let variants = map_fallible(tail, self, Self::variant)?;
+                (span, head, Subdatatypes::Variants(variants))
+            }
+        })
+    }
+
+    pub fn variant(&mut self, sexp: &Sexp) -> Result<Variant, ParseError> {
+        let (name, tail, span) = sexp.expect_call("datatype variant")?;
+
+        let (types, cost) = match tail {
+            [types @ .., Sexp::Atom(o, _), c] if *o == ":cost".into() => {
+                (types, Some(c.expect_uint("cost")?))
+            }
+            types => (types, None),
+        };
+
+        Ok(Variant {
+            span,
+            name,
+            types: map_fallible(types, self, |_, sexp| {
+                sexp.expect_atom("variant argument type")
+            })?,
+            cost,
         })
     }
 }
@@ -1099,7 +1122,7 @@ mod tests {
             Ok(Expr::Call(
                 span,
                 y.into(),
-                map_fallible(tail, macros, Parser::expr)?,
+                map_fallible(tail, macros, Parser::parse_expr)?,
             ))
         })));
         let s = r#"(f (qqqq a 3) 4.0 (H "hello"))"#;

From c331027fa7c9dd1c0c32ebddec1925ca62d3386e Mon Sep 17 00:00:00 2001
From: Yihong Zhang <yz489@cs.washington.edu>
Date: Sun, 5 Jan 2025 13:54:38 -0800
Subject: [PATCH 3/6] address nits

---
 src/ast/parse.rs | 108 ++++++++++++++++++++++-------------------------
 1 file changed, 51 insertions(+), 57 deletions(-)

diff --git a/src/ast/parse.rs b/src/ast/parse.rs
index 3fd12a02..57bf3eb5 100644
--- a/src/ast/parse.rs
+++ b/src/ast/parse.rs
@@ -131,13 +131,7 @@ impl Display for Span {
 // error messages are defined in the same place that they are created,
 // making it easier to improve errors over time.
 #[derive(Debug, Error)]
-pub struct ParseError(Span, String);
-
-impl ParseError {
-    pub fn new(span: Span, message: &str) -> Self {
-        Self(span, message.to_string())
-    }
-}
+pub struct ParseError(pub Span, pub String);
 
 impl std::fmt::Display for ParseError {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
@@ -223,46 +217,6 @@ fn map_fallible<T>(
         .collect::<Result<_, _>>()
 }
 
-// helper for parsing a list of options
-fn options(sexps: &[Sexp]) -> Result<Vec<(&str, &[Sexp])>, ParseError> {
-    fn option_name(sexp: &Sexp) -> Option<&str> {
-        if let Ok(symbol) = sexp.expect_atom("") {
-            let s: &str = symbol.into();
-            if let Some(':') = s.chars().next() {
-                return Some(s);
-            }
-        }
-        None
-    }
-
-    let mut out = Vec::new();
-    let mut i = 0;
-    while i < sexps.len() {
-        let Some(key) = option_name(&sexps[i]) else {
-            return error!(sexps[i].span(), "option key must start with ':'");
-        };
-        i += 1;
-
-        let start = i;
-        while i < sexps.len() && option_name(&sexps[i]).is_none() {
-            i += 1;
-        }
-        out.push((key, &sexps[start..i]));
-    }
-    Ok(out)
-}
-
-fn schema(input: &Sexp, output: &Sexp) -> Result<Schema, ParseError> {
-    Ok(Schema {
-        input: input
-            .expect_list("input sorts")?
-            .iter()
-            .map(|sexp| sexp.expect_atom("input sort"))
-            .collect::<Result<_, _>>()?,
-        output: output.expect_atom("output sort")?,
-    })
-}
-
 pub trait Macro<T> {
     fn get_head(&self) -> String;
     fn parse(&self, args: &[Sexp], span: Span, parser: &mut Parser) -> Result<T, ParseError>;
@@ -391,8 +345,8 @@ impl Parser {
             "function" => match tail {
                 [name, inputs, output, rest @ ..] => vec![Command::Function {
                     name: name.expect_atom("function name")?,
-                    schema: schema(inputs, output)?,
-                    merge: match options(rest)?.as_slice() {
+                    schema: self.parse_schema(inputs, output)?,
+                    merge: match self.parse_options(rest)?.as_slice() {
                         [(":no-merge", [])] => None,
                         [(":merge", [e])] => Some(self.parse_expr(e)?),
                         [] => {
@@ -415,7 +369,7 @@ impl Parser {
                 [name, inputs, output, rest @ ..] => {
                     let mut cost = None;
                     let mut unextractable = false;
-                    match options(rest)?.as_slice() {
+                    match self.parse_options(rest)?.as_slice() {
                         [] => {}
                         [(":unextractable", [])] => unextractable = true,
                         [(":cost", [c])] => cost = Some(c.expect_uint("cost")?),
@@ -425,7 +379,7 @@ impl Parser {
                     vec![Command::Constructor {
                         span,
                         name: name.expect_atom("constructor name")?,
-                        schema: schema(inputs, output)?,
+                        schema: self.parse_schema(inputs, output)?,
                         cost,
                         unextractable,
                     }]
@@ -475,7 +429,7 @@ impl Parser {
 
                     let mut ruleset = "".into();
                     let mut name = "".into();
-                    for option in options(rest)? {
+                    for option in self.parse_options(rest)? {
                         match option {
                             (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
                             (":name", [s]) => name = s.expect_string("rule name")?.into(),
@@ -499,7 +453,7 @@ impl Parser {
                     let mut ruleset = "".into();
                     let mut conditions = Vec::new();
                     let mut subsume = false;
-                    for option in options(rest)? {
+                    for option in self.parse_options(rest)? {
                         match option {
                             (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
                             (":subsume", []) => subsume = true,
@@ -534,7 +488,7 @@ impl Parser {
 
                     let mut ruleset = "".into();
                     let mut conditions = Vec::new();
-                    for option in options(rest)? {
+                    for option in self.parse_options(rest)? {
                         match option {
                             (":ruleset", [r]) => ruleset = r.expect_atom("ruleset name")?,
                             (":when", [w]) => {
@@ -581,7 +535,7 @@ impl Parser {
                     )
                 };
 
-                let until = match options(rest)?.as_slice() {
+                let until = match self.parse_options(rest)?.as_slice() {
                     [] => None,
                     [(":until", facts)] => Some(map_fallible(facts, self, Self::parse_fact)?),
                     _ => return error!(span, "could not parse run options"),
@@ -607,7 +561,7 @@ impl Parser {
             },
             "query-extract" => match tail {
                 [rest @ .., e] => {
-                    let variants = match options(rest)?.as_slice() {
+                    let variants = match self.parse_options(rest)?.as_slice() {
                         [] => 0,
                         [(":variants", [v])] => v.expect_uint("number of variants")?,
                         _ => return error!(span, "could not parse query-extract options"),
@@ -744,7 +698,7 @@ impl Parser {
                     ("".into(), tail)
                 };
 
-                let until = match options(rest)?.as_slice() {
+                let until = match self.parse_options(rest)?.as_slice() {
                     [] => None,
                     [(":until", facts)] => Some(map_fallible(facts, self, Self::parse_fact)?),
                     _ => return error!(span, "could not parse run options"),
@@ -905,6 +859,46 @@ impl Parser {
             cost,
         })
     }
+
+    // helper for parsing a list of options
+    pub fn parse_options<'a>(&self, sexps: &'a [Sexp]) -> Result<Vec<(&'a str, &'a [Sexp])>, ParseError> {
+        fn option_name(sexp: &Sexp) -> Option<&str> {
+            if let Ok(symbol) = sexp.expect_atom("") {
+                let s: &str = symbol.into();
+                if let Some(':') = s.chars().next() {
+                    return Some(s);
+                }
+            }
+            None
+        }
+
+        let mut out = Vec::new();
+        let mut i = 0;
+        while i < sexps.len() {
+            let Some(key) = option_name(&sexps[i]) else {
+                return error!(sexps[i].span(), "option key must start with ':'");
+            };
+            i += 1;
+
+            let start = i;
+            while i < sexps.len() && option_name(&sexps[i]).is_none() {
+                i += 1;
+            }
+            out.push((key, &sexps[start..i]));
+        }
+        Ok(out)
+    }
+
+    pub fn parse_schema(&self, input: &Sexp, output: &Sexp) -> Result<Schema, ParseError> {
+        Ok(Schema {
+            input: input
+                .expect_list("input sorts")?
+                .iter()
+                .map(|sexp| sexp.expect_atom("input sort"))
+                .collect::<Result<_, _>>()?,
+            output: output.expect_atom("output sort")?,
+        })
+    }
 }
 
 #[derive(Clone, Debug)]

From cb0716a4d8a0cf6b70396047db0cf717060ee64c Mon Sep 17 00:00:00 2001
From: Yihong Zhang <yz489@cs.washington.edu>
Date: Sun, 5 Jan 2025 15:44:32 -0800
Subject: [PATCH 4/6] get_head -> name

---
 src/ast/parse.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/ast/parse.rs b/src/ast/parse.rs
index 57bf3eb5..90aacfc1 100644
--- a/src/ast/parse.rs
+++ b/src/ast/parse.rs
@@ -218,7 +218,7 @@ fn map_fallible<T>(
 }
 
 pub trait Macro<T> {
-    fn get_head(&self) -> String;
+    fn name(&self) -> String;
     fn parse(&self, args: &[Sexp], span: Span, parser: &mut Parser) -> Result<T, ParseError>;
 }
 
@@ -237,7 +237,7 @@ impl<T, F> Macro<T> for SimpleMacro<T, F>
 where
     F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError>,
 {
-    fn get_head(&self) -> String {
+    fn name(&self) -> String {
         self.0.clone()
     }
 
@@ -287,15 +287,15 @@ impl Parser {
     }
 
     pub fn add_command_macro(&mut self, ma: Arc<dyn Macro<Vec<Command>>>) {
-        self.commands.insert(ma.get_head().into(), ma);
+        self.commands.insert(ma.name().into(), ma);
     }
 
     pub fn add_action_macro(&mut self, ma: Arc<dyn Macro<Vec<Action>>>) {
-        self.actions.insert(ma.get_head().into(), ma);
+        self.actions.insert(ma.name().into(), ma);
     }
 
     pub fn add_expr_macro(&mut self, ma: Arc<dyn Macro<Expr>>) {
-        self.exprs.insert(ma.get_head().into(), ma);
+        self.exprs.insert(ma.name().into(), ma);
     }
 
     pub fn parse_command(&mut self, sexp: &Sexp) -> Result<Vec<Command>, ParseError> {

From 0e8a8d1d5971ca7ae6e738a40a54c5af213c1d52 Mon Sep 17 00:00:00 2001
From: Yihong Zhang <yz489@cs.washington.edu>
Date: Sun, 5 Jan 2025 15:46:00 -0800
Subject: [PATCH 5/6] nits

---
 src/ast/parse.rs | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/ast/parse.rs b/src/ast/parse.rs
index 90aacfc1..68191760 100644
--- a/src/ast/parse.rs
+++ b/src/ast/parse.rs
@@ -861,7 +861,10 @@ impl Parser {
     }
 
     // helper for parsing a list of options
-    pub fn parse_options<'a>(&self, sexps: &'a [Sexp]) -> Result<Vec<(&'a str, &'a [Sexp])>, ParseError> {
+    pub fn parse_options<'a>(
+        &self,
+        sexps: &'a [Sexp],
+    ) -> Result<Vec<(&'a str, &'a [Sexp])>, ParseError> {
         fn option_name(sexp: &Sexp) -> Option<&str> {
             if let Ok(symbol) = sexp.expect_atom("") {
                 let s: &str = symbol.into();

From 63520b9488f33adf4033a787892e9226dbc917ca Mon Sep 17 00:00:00 2001
From: Yihong Zhang <yz489@cs.washington.edu>
Date: Sun, 5 Jan 2025 15:47:50 -0800
Subject: [PATCH 6/6] make Macro::name return Symbol than String

---
 src/ast/parse.rs | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/ast/parse.rs b/src/ast/parse.rs
index 68191760..d619897b 100644
--- a/src/ast/parse.rs
+++ b/src/ast/parse.rs
@@ -218,11 +218,11 @@ fn map_fallible<T>(
 }
 
 pub trait Macro<T> {
-    fn name(&self) -> String;
+    fn name(&self) -> Symbol;
     fn parse(&self, args: &[Sexp], span: Span, parser: &mut Parser) -> Result<T, ParseError>;
 }
 
-pub struct SimpleMacro<T, F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError>>(String, F);
+pub struct SimpleMacro<T, F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError>>(Symbol, F);
 
 impl<T, F> SimpleMacro<T, F>
 where
@@ -237,8 +237,8 @@ impl<T, F> Macro<T> for SimpleMacro<T, F>
 where
     F: Fn(&[Sexp], Span, &mut Parser) -> Result<T, ParseError>,
 {
-    fn name(&self) -> String {
-        self.0.clone()
+    fn name(&self) -> Symbol {
+        self.0
     }
 
     fn parse(&self, args: &[Sexp], span: Span, parser: &mut Parser) -> Result<T, ParseError> {
@@ -287,15 +287,15 @@ impl Parser {
     }
 
     pub fn add_command_macro(&mut self, ma: Arc<dyn Macro<Vec<Command>>>) {
-        self.commands.insert(ma.name().into(), ma);
+        self.commands.insert(ma.name(), ma);
     }
 
     pub fn add_action_macro(&mut self, ma: Arc<dyn Macro<Vec<Action>>>) {
-        self.actions.insert(ma.name().into(), ma);
+        self.actions.insert(ma.name(), ma);
     }
 
     pub fn add_expr_macro(&mut self, ma: Arc<dyn Macro<Expr>>) {
-        self.exprs.insert(ma.name().into(), ma);
+        self.exprs.insert(ma.name(), ma);
     }
 
     pub fn parse_command(&mut self, sexp: &Sexp) -> Result<Vec<Command>, ParseError> {