Support user-defined scheduler

src/actions.rs

@@ -227,8 +227,7 @@ impl EGraph {
                     }
                     MergeFn::Expr(merge_prog) => {
                         let values = [old_value, new_value];
-                        let mut stack = vec![];
-                        self.run_actions(&mut stack, &values, &merge_prog, true)?;
+                        let mut stack = self.run_actions(&values, &merge_prog, true)?;
                         stack.pop().unwrap()
                     }
                 };
@@ -243,7 +242,7 @@ impl EGraph {
                     let values = [old_value, new_value];
                     // We need to pass a new stack instead of reusing the old one
                     // because Load(Stack(idx)) use absolute index.
-                    self.run_actions(&mut Vec::new(), &values, &prog, true)?;
+                    self.run_actions(&values, &prog, true)?;
                 }
             }
         } else {
@@ -252,13 +251,15 @@ impl EGraph {
         Ok(())
     }
 
+    /// Runs actions with the given substitution
+    /// Returns the resulting stack if successful
     pub(crate) fn run_actions(
         &mut self,
-        stack: &mut Vec<Value>,
         subst: &[Value],
         program: &Program,
         make_defaults: bool,
-    ) -> Result<(), Error> {
+    ) -> Result<Vec<Value>, Error> {
+        let mut stack = vec![];
         for instr in &program.0 {
             match instr {
                 Instruction::Load(load) => match load {
@@ -337,7 +338,7 @@ impl EGraph {
                     let new_value = stack.pop().unwrap();
                     let new_len = stack.len() - function.schema.input.len();
 
-                    self.perform_set(*f, new_value, stack)?;
+                    self.perform_set(*f, new_value, &mut stack)?;
                     stack.truncate(new_len)
                 }
                 Instruction::Union(arity) => {
@@ -423,6 +424,6 @@ impl EGraph {
                 }
             }
         }
-        Ok(())
+        Ok(stack)
     }
 }

src/ast/desugar.rs

@@ -1,3 +1,5 @@
+use hashbrown::HashMap;
+
 use super::{Rewrite, Rule};
 use crate::*;
 
@@ -76,6 +78,7 @@ fn desugar_rewrite(
         name,
         rule: Rule {
             span: span.clone(),
+            props: rewrite.props.clone(),
             body: [Fact::Eq(
                 span.clone(),
                 vec![Expr::Var(span, var), rewrite.lhs.clone()],
@@ -92,6 +95,7 @@ fn desugar_birewrite(ruleset: Symbol, name: Symbol, rewrite: &Rewrite) -> Vec<NC
     let span = rewrite.span.clone();
     let rw2 = Rewrite {
         span,
+        props: rewrite.props.clone(),
         lhs: rewrite.rhs.clone(),
         rhs: rewrite.lhs.clone(),
         conditions: rewrite.conditions.clone(),
@@ -238,24 +242,28 @@ pub(crate) fn desugar_command(
         }
         Command::Rule {
             ruleset,
-            mut name,
+            name,
             rule,
         } => {
-            if name == "".into() {
-                name = rule.to_string().replace('\"', "'").into();
-            }
+            let name = if name.as_str() == "" {
+                // Adding a unique prefix because there may be duplicate rules (e.g., when
+                // processing a resugared egglog program).
+                desugar.get_fresh().to_string() + &rule.to_string().replace('\"', "'")
+            } else {
+                name.to_string()
+            };
 
             let mut result = vec![NCommand::NormRule {
                 ruleset,
-                name,
+                name: name.clone().into(),
                 rule: rule.clone(),
             }];
 
             if seminaive_transform {
                 if let Some(new_rule) = add_semi_naive_rule(desugar, rule) {
                     result.push(NCommand::NormRule {
                         ruleset,
-                        name,
+                        name: (desugar.get_fresh().to_string() + "__seminaive-" + &name).into(),
                         rule: new_rule,
                     });
                 }
@@ -305,6 +313,7 @@ pub(crate) fn desugar_command(
                 let fresh_rulename = desugar.get_fresh();
                 let rule = Rule {
                     span: span.clone(),
+                    props: HashMap::default(),
                     body: vec![Fact::Eq(
                         span.clone(),
                         vec![Expr::Var(span.clone(), fresh), expr.clone()],

src/ast/mod.rs

@@ -177,7 +177,7 @@ pub(crate) type ResolvedNCommand = GenericNCommand<ResolvedCall, ResolvedVar>;
 /// TODO: The name "NCommand" used to denote normalized command, but this
 /// meaning is obsolete. A future PR should rename this type to something
 /// like "DCommand".
-#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub enum GenericNCommand<Head, Leaf>
 where
     Head: Clone + Display,
@@ -343,6 +343,12 @@ pub(crate) type ResolvedSchedule = GenericSchedule<ResolvedCall, ResolvedVar>;
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum GenericSchedule<Head, Leaf> {
+    WithScheduler(
+        Span,
+        Symbol,
+        Vec<GenericExpr<Head, Leaf>>,
+        Box<GenericSchedule<Head, Leaf>>,
+    ),
     Saturate(Span, Box<GenericSchedule<Head, Leaf>>),
     Repeat(Span, usize, Box<GenericSchedule<Head, Leaf>>),
     Run(Span, GenericRunConfig<Head, Leaf>),
@@ -410,6 +416,11 @@ where
                 span,
                 scheds.into_iter().map(|s| s.visit_exprs(f)).collect(),
             ),
+            GenericSchedule::WithScheduler(span, scheduler, args, sched) => {
+                let sched = Box::new(sched.visit_exprs(f));
+                let args = args.into_iter().map(f).collect();
+                GenericSchedule::WithScheduler(span, scheduler, args, sched)
+            }
         }
     }
 }
@@ -421,6 +432,9 @@ impl<Head: Display, Leaf: Display> ToSexp for GenericSchedule<Head, Leaf> {
             GenericSchedule::Repeat(_ann, size, sched) => list!("repeat", size, sched),
             GenericSchedule::Run(_ann, config) => config.to_sexp(),
             GenericSchedule::Sequence(_ann, scheds) => list!("seq", ++ scheds),
+            GenericSchedule::WithScheduler(_ann, scheduler, args, sched) => {
+                list!("with-scheduler", list!(scheduler, ++ args), sched)
+            }
         }
     }
 }
@@ -1285,6 +1299,16 @@ where
     }
 }
 
+impl<Head, Leaf> Display for Facts<Head, Leaf>
+where
+    Head: Clone + Display,
+    Leaf: Clone + PartialEq + Eq + Display + Hash,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", ListDisplay(&self.0, "\n"))
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct CorrespondingVar<Head, Leaf>
 where
@@ -1580,21 +1604,23 @@ where
 }
 
 #[derive(Clone, Debug)]
-pub(crate) struct CompiledRule {
+pub struct CompiledRule {
+    pub props: HashMap<String, Literal>,
     pub(crate) query: CompiledQuery,
     pub(crate) program: Program,
 }
 
 pub type Rule = GenericRule<Symbol, Symbol>;
 pub(crate) type ResolvedRule = GenericRule<ResolvedCall, ResolvedVar>;
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct GenericRule<Head, Leaf>
 where
     Head: Clone + Display,
     Leaf: Clone + PartialEq + Eq + Display + Hash,
 {
     pub span: Span,
+    pub props: HashMap<String, Literal>,
     pub head: GenericActions<Head, Leaf>,
     pub body: Vec<GenericFact<Head, Leaf>>,
 }
@@ -1610,6 +1636,7 @@ where
     ) -> Self {
         Self {
             span: self.span,
+            props: self.props,
             head: self.head.visit_exprs(f),
             body: self
                 .body
@@ -1708,6 +1735,7 @@ type Rewrite = GenericRewrite<Symbol, Symbol>;
 #[derive(Clone, Debug)]
 pub struct GenericRewrite<Head, Leaf> {
     pub span: Span,
+    pub props: HashMap<String, Literal>,
     pub lhs: GenericExpr<Head, Leaf>,
     pub rhs: GenericExpr<Head, Leaf>,
     pub conditions: Vec<GenericFact<Head, Leaf>>,

src/ast/parse.lalrpop

@@ -61,16 +61,19 @@ Command: Command = {
     <lo:LParen> "rule" <body:List<Fact>> <head:List<Action>> 
         <ruleset:(":ruleset" <Ident>)?> 
         <name:(":name" <String>)?> 
-    <hi:RParen> => Command::Rule{ruleset: ruleset.unwrap_or("".into()), name: name.unwrap_or("".to_string()).into(), rule: Rule { span: Span(srcfile.clone(), lo, hi), head: Actions::new(head), body }},
+        <props:(":prop" <String> <Literal>)*>
+    <hi:RParen> => Command::Rule{ruleset: ruleset.unwrap_or("".into()), name: name.unwrap_or("".to_string()).into(), rule: Rule { span: Span(srcfile.clone(), lo, hi), props: props.into_iter().collect(), head: Actions::new(head), body }},
     <lo:LParen> "rewrite" <lhs:Expr> <rhs:Expr>
         <subsume:(":subsume")?>
         <conditions:(":when" <List<Fact>>)?>
         <ruleset:(":ruleset" <Ident>)?>
-    <hi:RParen> => Command::Rewrite(ruleset.unwrap_or("".into()), Rewrite { span: Span(srcfile.clone(), lo, hi), lhs, rhs, conditions: conditions.unwrap_or_default() }, subsume.is_some()),
+        <props:(":prop" <String> <Literal>)*>
+    <hi:RParen> => Command::Rewrite(ruleset.unwrap_or("".into()), Rewrite { span: Span(srcfile.clone(), lo, hi), props: props.into_iter().collect(), lhs, rhs, conditions: conditions.unwrap_or_default() }, subsume.is_some()),
     <lo:LParen> "birewrite" <lhs:Expr> <rhs:Expr>
         <conditions:(":when" <List<Fact>>)?>
         <ruleset:(":ruleset" <Ident>)?>
-    <hi:RParen> => Command::BiRewrite(ruleset.unwrap_or("".into()), Rewrite { span: Span(srcfile.clone(), lo, hi), lhs, rhs, conditions: conditions.unwrap_or_default() }),
+        <props:(":prop" <String> <Literal>)*>
+    <hi:RParen> => Command::BiRewrite(ruleset.unwrap_or("".into()), Rewrite { span: Span(srcfile.clone(), lo, hi), props: props.into_iter().collect(), lhs, rhs, conditions: conditions.unwrap_or_default() }),
     <lo:LParen> "let" <name:Ident> <expr:Expr> <hi:RParen> => Command::Action(Action::Let(Span(srcfile.clone(), lo, hi), name, expr)),
     <NonLetAction> => Command::Action(<>),
     <lo:LParen> "run" <limit:UNum>  <until:(":until" <(Fact)*>)?> <hi:RParen> => 
@@ -98,6 +101,11 @@ Command: Command = {
 }
 
 Schedule: Schedule = {
+    <lo:LParen> "with-scheduler" 
+        <lo2:LParen> <scheduler:Ident> <args:Expr*> <hi2:RParen>
+        <sched:Schedule>
+    <hi:RParen> =>
+        Schedule::WithScheduler(Span(srcfile.clone(), lo, hi), scheduler, args, Box::new(sched)),
     <lo:LParen> "saturate" <scheds:Schedule*> <hi:RParen> => 
         Schedule::Saturate(Span(srcfile.clone(), lo, hi), Box::new(
             Schedule::Sequence(Span(srcfile.clone(), lo, hi), scheds))),

src/ast/remove_globals.rs

@@ -183,6 +183,7 @@ impl<'a> GlobalRemover<'a> {
 
                 let new_rule = GenericRule {
                     span: rule.span,
+                    props: rule.props,
                     // instrument the old facts and add the new facts to the end
                     body: rule
                         .body

src/core.rs

@@ -802,7 +802,8 @@ where
         Leaf: SymbolLike,
     {
         let GenericRule {
-            span: _,
+            span,
+            props: _,
             head,
             body,
         } = self;
@@ -812,7 +813,7 @@ where
         let (head, _correspondence) =
             head.to_core_actions(typeinfo, &mut binding, &mut fresh_gen)?;
         Ok(GenericCoreRule {
-            span: self.span.clone(),
+            span: span.clone(),
             body,
             head,
         })