Make custom schedulers compose with the term/proof encoding

CHANGELOG.md

@@ -9,6 +9,7 @@
 - Add a BigRat to-i64 primitive for integral rationals.
 - Add f64 exp, log, and sqrt primitives.
 - Add `RunReport::can_stop` so scheduler progress can be reported separately from database updates.
+- Make custom schedulers compose with the term/proof encoding. The term-encoding maintenance schedule (rebuilding/congruence/deferred deletes) is now carried in the program via an internal `(internal-rebuild-schedule ...)` command and run to fixpoint after every step inside `step_rules` and `step_rules_with_scheduler`, rather than being woven around each `run` in the desugared schedule. This keeps the desugared program self-contained while letting a scheduler drive only the user's ruleset, with maintenance always running unscheduled.
 - Add `EGraph::typecheck_expr_with_bindings_and_output`, `Core::eval_resolved_expr`, and `Core::apply_primitive` for body-defined primitive support, including normal command-path global rewrites for expressions typechecked through the helper.
 - Allow `unstable-fn` function containers to target primitive overloads.
 - Desugar `relation`s to `constructor`s to simplify the language and implementation. Relations no longer return unit `()` values.

src/ast/check_shadowing.rs

@@ -76,6 +76,7 @@ impl Names {
             }
             ResolvedNCommand::Extract(..) => Ok(()),
             ResolvedNCommand::RunSchedule(..) => Ok(()),
+            ResolvedNCommand::SetRebuildSchedule(..) => Ok(()),
             ResolvedNCommand::PrintOverallStatistics(..) => Ok(()),
             ResolvedNCommand::ProveExists(..) => Ok(()),
             ResolvedNCommand::PrintFunction(..) => Ok(()),

src/ast/desugar.rs

@@ -161,6 +161,9 @@ pub(crate) fn desugar_command(
         Command::RunSchedule(sched) => {
             vec![NCommand::RunSchedule(sched.clone())]
         }
+        Command::SetRebuildSchedule(sched) => {
+            vec![NCommand::SetRebuildSchedule(sched.clone())]
+        }
         Command::PrintOverallStatistics(span, file) => {
             vec![NCommand::PrintOverallStatistics(span, file.clone())]
         }

src/ast/mod.rs

@@ -76,6 +76,9 @@ where
     CoreAction(GenericAction<Head, Leaf>),
     Extract(Span, GenericExpr<Head, Leaf>, GenericExpr<Head, Leaf>),
     RunSchedule(GenericSchedule<Head, Leaf>),
+    /// Internal: records the maintenance schedule the term encoding runs
+    /// between steps. Not for direct use; emitted by the term-encoding pass.
+    SetRebuildSchedule(GenericSchedule<Head, Leaf>),
     PrintOverallStatistics(Span, Option<String>),
     Check(Span, Vec<GenericFact<Head, Leaf>>),
     PrintFunction(
@@ -155,6 +158,9 @@ where
             }
             GenericNCommand::NormRule { rule } => GenericCommand::Rule { rule: rule.clone() },
             GenericNCommand::RunSchedule(schedule) => GenericCommand::RunSchedule(schedule.clone()),
+            GenericNCommand::SetRebuildSchedule(schedule) => {
+                GenericCommand::SetRebuildSchedule(schedule.clone())
+            }
             GenericNCommand::PrintOverallStatistics(span, file) => {
                 GenericCommand::PrintOverallStatistics(span.clone(), file.clone())
             }
@@ -226,6 +232,7 @@ where
             | GenericNCommand::Pop(..)
             | GenericNCommand::Input { .. }
             | GenericNCommand::UserDefined(..)
+            | GenericNCommand::SetRebuildSchedule(..)
             | GenericNCommand::ProveExists(..) => self,
         }
     }
@@ -262,6 +269,9 @@ where
             GenericNCommand::RunSchedule(schedule) => {
                 GenericNCommand::RunSchedule(schedule.visit_exprs(f))
             }
+            GenericNCommand::SetRebuildSchedule(schedule) => {
+                GenericNCommand::SetRebuildSchedule(schedule.visit_exprs(f))
+            }
             GenericNCommand::PrintOverallStatistics(span, file) => {
                 GenericNCommand::PrintOverallStatistics(span, file)
             }
@@ -849,6 +859,10 @@ where
     ///
     /// See [`Schedule`] for more details.
     RunSchedule(GenericSchedule<Head, Leaf>),
+    /// Internal command emitted by the term-encoding pass to record the
+    /// maintenance schedule that `step_rules` runs between steps. Not intended
+    /// for direct use.
+    SetRebuildSchedule(GenericSchedule<Head, Leaf>),
     /// Print runtime statistics about rules
     /// and rulesets so far.
     PrintOverallStatistics(Span, Option<String>),
@@ -1051,6 +1065,9 @@ where
             }
             GenericCommand::Rule { rule } => rule.fmt(f),
             GenericCommand::RunSchedule(sched) => write!(f, "(run-schedule {sched})"),
+            GenericCommand::SetRebuildSchedule(sched) => {
+                write!(f, "(internal-rebuild-schedule {sched})")
+            }
             GenericCommand::PrintOverallStatistics(_span, file) => match file {
                 Some(file) => write!(f, "(print-stats :file {file})"),
                 None => write!(f, "(print-stats)"),
@@ -1743,6 +1760,9 @@ where
             GenericCommand::RunSchedule(schedule) => {
                 GenericCommand::RunSchedule(schedule.map_string_symbols(fun))
             }
+            GenericCommand::SetRebuildSchedule(schedule) => {
+                GenericCommand::SetRebuildSchedule(schedule.map_string_symbols(fun))
+            }
             GenericCommand::PrintOverallStatistics(span, file) => {
                 GenericCommand::PrintOverallStatistics(span, file)
             }
@@ -1852,6 +1872,9 @@ where
             GenericCommand::RunSchedule(schedule) => {
                 GenericCommand::RunSchedule(schedule.visit_exprs(f))
             }
+            GenericCommand::SetRebuildSchedule(schedule) => {
+                GenericCommand::SetRebuildSchedule(schedule.visit_exprs(f))
+            }
             GenericCommand::Fail(span, cmd) => {
                 GenericCommand::Fail(span, Box::new(cmd.visit_exprs(f)))
             }
@@ -1963,6 +1986,9 @@ where
             GenericCommand::RunSchedule(schedule) => {
                 GenericCommand::RunSchedule(schedule.map_symbols(head, leaf))
             }
+            GenericCommand::SetRebuildSchedule(schedule) => {
+                GenericCommand::SetRebuildSchedule(schedule.map_symbols(head, leaf))
+            }
             GenericCommand::PrintOverallStatistics(span, file) => {
                 GenericCommand::PrintOverallStatistics(span, file)
             }

src/ast/parse.rs

@@ -631,6 +631,10 @@ impl Parser {
                 span,
                 map_fallible(tail, self, Self::parse_schedule)?,
             ))],
+            "internal-rebuild-schedule" => match tail {
+                [sched] => vec![Command::SetRebuildSchedule(self.parse_schedule(sched)?)],
+                _ => return error!(span, "usage: (internal-rebuild-schedule <schedule>)"),
+            },
             "extract" => match tail {
                 [e] => vec![Command::Extract(
                     span.clone(),

src/lib.rs

@@ -308,6 +308,11 @@ pub struct EGraph {
     proof_state: EncodingState,
     /// In proof mode, this is the program before proof instrumentation and the version we use for proof checking.
     proof_check_program: Vec<ResolvedNCommand>,
+    /// Maintenance schedule run to fixpoint after each step under term
+    /// encoding, set via the internal `(internal-rebuild-schedule ...)` command
+    /// so it travels with the desugared program. `None` when term encoding is
+    /// off.
+    rebuild_schedule: Option<ResolvedSchedule>,
 }
 
 /// A user-defined command allows users to inject custom command that can be called
@@ -417,6 +422,7 @@ impl Default for EGraph {
             command_macros: Default::default(),
             proof_state,
             proof_check_program: vec![],
+            rebuild_schedule: None,
         };
         add_base_sort(&mut eg, UnitSort, span!()).unwrap();
         add_base_sort(&mut eg, StringSort, span!()).unwrap();
@@ -1130,7 +1136,36 @@ impl EGraph {
             .run_rules(&rule_ids)
             .map_err(|e| Error::BackendError(e.to_string()))?;
 
-        Ok(RunReport::singleton(ruleset, iteration_report))
+        let mut report = RunReport::singleton(ruleset, iteration_report);
+        report.union(self.maybe_run_rebuild_schedule(ruleset)?);
+        Ok(report)
+    }
+
+    /// Under term encoding, run the maintenance schedule (set via
+    /// `(internal-rebuild-schedule ...)`) to fixpoint so the next step and any
+    /// query observe a canonical e-graph. A no-op when no schedule is set, or
+    /// when `ruleset` is itself part of the maintenance schedule (which must not
+    /// re-trigger it).
+    fn maybe_run_rebuild_schedule(&mut self, ruleset: &str) -> Result<RunReport, Error> {
+        let Some(sched) = self.rebuild_schedule.clone() else {
+            return Ok(RunReport::default());
+        };
+        if Self::schedule_mentions_ruleset(&sched, ruleset) {
+            return Ok(RunReport::default());
+        }
+        self.run_schedule(&sched)
+    }
+
+    fn schedule_mentions_ruleset(schedule: &ResolvedSchedule, ruleset: &str) -> bool {
+        match schedule {
+            GenericSchedule::Run(_, config) => config.ruleset == ruleset,
+            GenericSchedule::Saturate(_, inner) | GenericSchedule::Repeat(_, _, inner) => {
+                Self::schedule_mentions_ruleset(inner, ruleset)
+            }
+            GenericSchedule::Sequence(_, scheds) => scheds
+                .iter()
+                .any(|s| Self::schedule_mentions_ruleset(s, ruleset)),
+        }
     }
 
     fn add_rule(&mut self, rule: ast::ResolvedRule) -> Result<String, Error> {
@@ -1628,6 +1663,9 @@ impl EGraph {
                 self.overall_run_report.union(report.clone());
                 return Ok(vec![CommandOutput::RunSchedule(report)]);
             }
+            ResolvedNCommand::SetRebuildSchedule(sched) => {
+                self.rebuild_schedule = Some(sched);
+            }
             ResolvedNCommand::PrintOverallStatistics(span, file) => match file {
                 None => {
                     log::info!("Printed overall statistics");

src/proofs/proof_encoding.rs

@@ -1281,7 +1281,11 @@ impl<'a> ProofInstrumentor<'a> {
     fn instrument_schedule(&mut self, schedule: &ResolvedSchedule) -> Schedule {
         match schedule {
             ResolvedSchedule::Run(span, config) => {
-                let new_run = match config.until {
+                // A run needs no maintenance wrapped around it: `step_rules`
+                // runs the schedule recorded by `(internal-rebuild-schedule
+                // ...)` after each step. Only the run's `:until` facts are
+                // rewritten to query the view tables.
+                match config.until {
                     Some(ref facts) => {
                         let (instrumented, _lookups, _proof) = self.instrument_facts(facts);
                         let instrumented_facts = self.parse_facts(&instrumented);
@@ -1300,8 +1304,7 @@ impl<'a> ProofInstrumentor<'a> {
                             until: None,
                         },
                     ),
-                };
-                Schedule::Sequence(span.clone(), vec![new_run, self.rebuild()])
+                }
             }
             ResolvedSchedule::Sequence(span, schedules) => Schedule::Sequence(
                 span.clone(),
@@ -1371,6 +1374,12 @@ impl<'a> ProofInstrumentor<'a> {
             ResolvedNCommand::RunSchedule(schedule) => {
                 res.push(Command::RunSchedule(self.instrument_schedule(schedule)));
             }
+            // Only emitted by this pass, never present in its input; pass through.
+            ResolvedNCommand::SetRebuildSchedule(schedule) => {
+                res.push(Command::SetRebuildSchedule(
+                    self.parse_schedule(schedule.to_string()),
+                ));
+            }
             ResolvedNCommand::Fail(span, cmd) => {
                 self.term_encode_command(cmd, res);
                 let last = res.pop().unwrap();
@@ -1441,16 +1450,26 @@ impl<'a> ProofInstrumentor<'a> {
                 let proof_header = self.proof_header();
                 res.extend(self.parse_program(&proof_header));
             }
+            // Record the maintenance schedule so `step_rules` runs it between
+            // steps. Carried in the program so the desugared output stays
+            // self-contained (e.g. when re-run without term encoding).
+            let rebuild = self.rebuild();
+            res.push(Command::SetRebuildSchedule(rebuild));
             self.egraph.proof_state.term_header_added = true;
         }
 
         for command in program {
             self.term_encode_command(&command, &mut res);
 
-            // run rebuilding after every command except a few
+            // Rebuild after each command's instrumented output, except where it
+            // can't have changed the e-graph (declarations) or already rebuilt
+            // internally (`RunSchedule`, via the fold in `step_rules`). Emitted
+            // per user command so a command's multi-command desugaring (e.g. a
+            // merge `set`) runs before maintenance reconciles it.
             if let ResolvedNCommand::Function(..)
             | ResolvedNCommand::NormRule { .. }
-            | ResolvedNCommand::Sort { .. } = &command
+            | ResolvedNCommand::Sort { .. }
+            | ResolvedNCommand::RunSchedule(..) = &command
             {
             } else {
                 res.push(Command::RunSchedule(self.rebuild()));

src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function1.snap

@@ -8,6 +8,7 @@ expression: snapshot
 (ruleset __rebuilding)
 (ruleset __rebuilding_cleanup)
 (ruleset __delete_subsume_ruleset)
+(internal-rebuild-schedule (seq (saturate (seq (run __rebuilding_cleanup) (saturate (run __single_parent)) (saturate (run __parent)) (saturate (run __uf_function_index)) (run __rebuilding))) (run __delete_subsume_ruleset)))
 (sort Math :internal-uf __UF_Math)
 (function __UF_Math (Math Math) Unit :merge old :unextractable :internal-hidden)
 (function __UF_Mathf (Math) Math :merge new :unextractable :internal-hidden)

src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap

@@ -8,6 +8,7 @@ expression: snapshot
 (ruleset __rebuilding)
 (ruleset __rebuilding_cleanup)
 (ruleset __delete_subsume_ruleset)
+(internal-rebuild-schedule (seq (saturate (seq (run __rebuilding_cleanup) (saturate (run __single_parent)) (saturate (run __parent)) (saturate (run __uf_function_index)) (run __rebuilding))) (run __delete_subsume_ruleset)))
 (sort __view)
 (constructor add (i64 i64 i64) __view :unextractable :internal-hidden)
 (function __addView (i64 i64 i64) Unit :merge old :unextractable :internal-term-constructor add)