From f3aeb35b53d7cd68c73e2bcf06561c29779ed6f1 Mon Sep 17 00:00:00 2001 From: oliver Date: Tue, 7 Jul 2026 17:51:39 +0000 Subject: [PATCH] Make custom schedulers compose with the term/proof encoding Custom schedulers such as egglog-experimental's backoff scheduler drive a ruleset through `step_rules_with_scheduler`, but that path never ran the term-encoding maintenance rulesets (rebuilding, congruence, path compression, deferred deletion). Scheduled steps therefore queried stale, un-canonicalized view tables under term/proof encoding. Record the maintenance schedule in the program via a new internal `(internal-rebuild-schedule ...)` command and run it to fixpoint inside `step_rules` and `step_rules_with_scheduler` after each step. Because the schedule travels with the desugared program, the desugared output stays self-contained and runnable without term encoding. Only the user's ruleset is scheduled; maintenance always runs unscheduled. Add tests for term encoding + scheduler, proof mode, the real backoff algorithm (saturation-equivalence on a benchmark), and a fractional scheduler (residual-match soundness). Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 1 + src/ast/check_shadowing.rs | 1 + src/ast/desugar.rs | 3 + src/ast/mod.rs | 26 ++ src/ast/parse.rs | 4 + src/lib.rs | 40 +- src/proofs/proof_encoding.rs | 29 +- ...sts__tests__doc_example_add_function1.snap | 1 + ...sts__tests__doc_example_add_function2.snap | 1 + src/scheduler.rs | 419 ++++++++++++++++++ src/typechecking.rs | 3 + 11 files changed, 522 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bf6fe547..1a8584448 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/ast/check_shadowing.rs b/src/ast/check_shadowing.rs index 7e808dff7..7e61f52ec 100644 --- a/src/ast/check_shadowing.rs +++ b/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(()), diff --git a/src/ast/desugar.rs b/src/ast/desugar.rs index 11614aa73..17f691410 100644 --- a/src/ast/desugar.rs +++ b/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())] } diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 4d1126dce..8a414bc19 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -76,6 +76,9 @@ where CoreAction(GenericAction), Extract(Span, GenericExpr, GenericExpr), RunSchedule(GenericSchedule), + /// Internal: records the maintenance schedule the term encoding runs + /// between steps. Not for direct use; emitted by the term-encoding pass. + SetRebuildSchedule(GenericSchedule), PrintOverallStatistics(Span, Option), Check(Span, Vec>), 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), + /// 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), /// Print runtime statistics about rules /// and rulesets so far. PrintOverallStatistics(Span, Option), @@ -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) } diff --git a/src/ast/parse.rs b/src/ast/parse.rs index 4d0aebb5e..3090f250a 100644 --- a/src/ast/parse.rs +++ b/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 )"), + }, "extract" => match tail { [e] => vec![Command::Extract( span.clone(), diff --git a/src/lib.rs b/src/lib.rs index 1fe33855e..4fd7dc2e9 100644 --- a/src/lib.rs +++ b/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, + /// 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, } /// 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 { + 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 { @@ -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"); diff --git a/src/proofs/proof_encoding.rs b/src/proofs/proof_encoding.rs index a51701e55..729c7f7ce 100644 --- a/src/proofs/proof_encoding.rs +++ b/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())); diff --git a/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function1.snap b/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function1.snap index de728f721..bcd697be0 100644 --- a/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function1.snap +++ b/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) diff --git a/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap b/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap index e664c5f3b..a0d0c2318 100644 --- a/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap +++ b/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) diff --git a/src/scheduler.rs b/src/scheduler.rs index eb5436fc1..1f2eb0435 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -271,6 +271,11 @@ impl EGraph { self.rulesets = rulesets; self.schedulers = schedulers; + // Under term encoding, run maintenance to fixpoint so the next + // scheduled step queries a canonical e-graph. The maintenance rulesets + // are run unscheduled; only the user's ruleset is scheduled. + query_report.union(self.maybe_run_rebuild_schedule(ruleset)?); + Ok(query_report) } } @@ -478,6 +483,420 @@ mod test { assert_eq!(iter, 12); } + #[derive(Clone)] + struct ChooseAllScheduler; + + impl Scheduler for ChooseAllScheduler { + fn filter_matches(&mut self, _rule: &str, _ruleset: &str, matches: &mut Matches) -> bool { + matches.choose_all(); + // Re-seek next step so we keep firing on rows produced by rebuilding. + true + } + } + + // A scheduled step under term encoding must run maintenance (rebuilding / + // congruence / UF indexing) before returning, otherwise the next step and + // any query would observe stale, un-canonicalized view tables. + #[test] + fn test_scheduler_runs_term_encoding_maintenance() { + let mut egraph = EGraph::new_with_term_encoding(); + let scheduler_id = egraph.add_scheduler(Box::new(ChooseAllScheduler)); + egraph + .parse_and_run_program( + None, + r#" + (sort Math) + (constructor Add (i64 i64) Math) + (Add 1 2) + (ruleset commutative) + (rule ((Add a b)) ((union (Add a b) (Add b a))) + :ruleset commutative :name "commutativity") + "#, + ) + .unwrap(); + + // One scheduled step fires the rule, creating `(Add 2 1)` and unioning + // it with `(Add 1 2)`. The equality is only visible in the view tables + // once maintenance has rebuilt them. + egraph + .step_rules_with_scheduler(scheduler_id, "commutative") + .unwrap(); + + // Passes only if maintenance ran during the scheduled step. + egraph + .parse_and_run_program(None, "(check (= (Add 1 2) (Add 2 1)))") + .unwrap(); + } + + // Driving a ruleset through a scheduler to a fixpoint must reach the same + // saturated e-graph as running it unscheduled, even under term encoding + // (where maintenance runs between scheduled steps). Uses `choose_all` so + // termination is unambiguous; the delaying-scheduler case is covered against + // the real backoff scheduler in egglog-experimental. + #[test] + fn test_scheduler_saturation_matches_unscheduled_term_encoding() { + let program = r#" + (sort Math) + (constructor Num (i64) Math) + (constructor Add (Math Math) Math) + (Add (Num 1) (Add (Num 2) (Num 3))) + (ruleset rw) + (rule ((Add x y)) ((union (Add x y) (Add y x))) + :name "comm" :ruleset rw) + (rule ((Add (Add x y) z)) ((union (Add (Add x y) z) (Add x (Add y z)))) + :name "assoc" :ruleset rw) + (rule ((Add x (Add y z))) ((union (Add x (Add y z)) (Add (Add x y) z))) + :name "assoc2" :ruleset rw) + "#; + + // (1) unscheduled saturation + let mut base = EGraph::new_with_term_encoding(); + base.parse_and_run_program(None, program).unwrap(); + base.parse_and_run_program(None, "(run-schedule (saturate rw))") + .unwrap(); + + // (2) driven through a delaying scheduler to a fixpoint + let mut sched = EGraph::new_with_term_encoding(); + let id = sched.add_scheduler(Box::new(ChooseAllScheduler)); + sched.parse_and_run_program(None, program).unwrap(); + let mut iters = 0; + loop { + let report = sched.step_rules_with_scheduler(id, "rw").unwrap(); + iters += 1; + assert!(iters < 10_000, "scheduler did not converge"); + if report.can_stop { + break; + } + } + + // Same number of e-nodes for every user constructor. + for ctor in ["Num", "Add"] { + assert_eq!( + base.get_size(ctor), + sched.get_size(ctor), + "e-node count differs for {ctor}" + ); + } + + // The same equalities and inequalities hold in both e-graphs. + let checks = [ + "(check (= (Add (Num 1) (Add (Num 2) (Num 3))) (Add (Num 3) (Add (Num 2) (Num 1)))))", + "(check (= (Add (Add (Num 1) (Num 2)) (Num 3)) (Add (Num 1) (Add (Num 2) (Num 3)))))", + "(check (!= (Num 1) (Num 2)))", + "(check (!= (Add (Num 1) (Num 2)) (Add (Num 1) (Num 3))))", + ]; + for check in checks { + base.parse_and_run_program(None, check) + .unwrap_or_else(|e| panic!("unscheduled check failed: {check}: {e}")); + sched + .parse_and_run_program(None, check) + .unwrap_or_else(|e| panic!("scheduled check failed: {check}: {e}")); + } + } + + // Fractional scheduler: fires one match at a time for a while (holding + // residual matches across many rebuilds — the "residual staleness" path), + // then chooses everything to converge. Used to check that residual matches + // held across term-encoding maintenance don't change the final result. + #[derive(Clone)] + struct DelayThenAll { + n: usize, + calls: usize, + budget: usize, + } + + impl Scheduler for DelayThenAll { + fn can_stop(&mut self, _rules: &[&str], _ruleset: &str) -> bool { + // Only stoppable once past the delaying phase, so pending residual + // matches during the delay never terminate the loop early. + self.calls >= self.budget + } + + fn filter_matches(&mut self, _rule: &str, _ruleset: &str, matches: &mut Matches) -> bool { + self.calls += 1; + if self.calls < self.budget { + let size = matches.match_size(); + for i in 0..size.min(self.n) { + matches.choose(i); + } + } else { + matches.choose_all(); + } + true + } + } + + #[test] + fn test_fractional_scheduler_matches_unscheduled_term_encoding() { + let program = r#" + (sort Math) + (constructor Num (i64) Math) + (constructor Add (Math Math) Math) + (Add (Num 1) (Add (Num 2) (Num 3))) + (ruleset rw) + (rule ((Add x y)) ((union (Add x y) (Add y x))) + :name "comm" :ruleset rw) + (rule ((Add (Add x y) z)) ((union (Add (Add x y) z) (Add x (Add y z)))) + :name "assoc" :ruleset rw) + (rule ((Add x (Add y z))) ((union (Add x (Add y z)) (Add (Add x y) z))) + :name "assoc2" :ruleset rw) + "#; + + let mut base = EGraph::new_with_term_encoding(); + base.parse_and_run_program(None, program).unwrap(); + base.parse_and_run_program(None, "(run-schedule (saturate rw))") + .unwrap(); + + let mut sched = EGraph::new_with_term_encoding(); + let id = sched.add_scheduler(Box::new(DelayThenAll { + n: 1, + calls: 0, + budget: 40, + })); + sched.parse_and_run_program(None, program).unwrap(); + let mut iters = 0; + loop { + let report = sched.step_rules_with_scheduler(id, "rw").unwrap(); + iters += 1; + assert!(iters < 10_000, "scheduler did not converge"); + if report.can_stop { + break; + } + } + + for ctor in ["Num", "Add"] { + assert_eq!( + base.get_size(ctor), + sched.get_size(ctor), + "e-node count differs for {ctor} (residual staleness affected the result)" + ); + } + for check in [ + "(check (= (Add (Num 1) (Add (Num 2) (Num 3))) (Add (Num 3) (Add (Num 2) (Num 1)))))", + "(check (!= (Num 1) (Num 2)))", + ] { + base.parse_and_run_program(None, check).unwrap(); + sched.parse_and_run_program(None, check).unwrap(); + } + } + + // Scheduling must keep proof tracking consistent: a term derived by a + // scheduled rule firing must still have an extractable, checkable proof. + #[test] + fn test_scheduler_preserves_proofs() { + // Proof-testing mode turns `check` into a proof extraction + check, so + // the equality check below validates the proof produced for the term + // derived by the scheduled rule firing. + let mut egraph = EGraph::new_with_proofs().with_proof_testing(); + let id = egraph.add_scheduler(Box::new(ChooseAllScheduler)); + egraph + .parse_and_run_program( + None, + r#" + (sort Math) + (constructor Add (i64 i64) Math) + (Add 1 2) + (ruleset commutative) + (rule ((Add a b)) ((union (Add a b) (Add b a))) + :ruleset commutative :name "commutativity") + "#, + ) + .unwrap(); + + // Drive the rule under the scheduler to a fixpoint. + let mut iters = 0; + loop { + let report = egraph.step_rules_with_scheduler(id, "commutative").unwrap(); + iters += 1; + assert!(iters < 1000, "scheduler did not converge"); + if report.can_stop { + break; + } + } + + // Under proof-testing this extracts and checks a proof of the equality + // derived by the scheduled rule. + egraph + .parse_and_run_program(None, "(check (= (Add 1 2) (Add 2 1)))") + .unwrap(); + } + + // The backoff scheduler, ported verbatim (minus logging) from + // egglog-experimental `src/scheduling.rs`, so we can check that a real + // banning scheduler reaches the same fixpoint under term encoding. Backoff + // is all-or-nothing (`choose_all` or ban), so it never leaves residual + // matches. + #[derive(Clone)] + struct BackOffScheduler { + default_match_limit: usize, + default_ban_length: usize, + stats: crate::util::HashMap, + } + + #[derive(Clone)] + struct RuleStats { + iteration: usize, + times_applied: usize, + banned_until: usize, + times_banned: usize, + match_limit: usize, + ban_length: usize, + } + + impl BackOffScheduler { + fn new(match_limit: usize, ban_length: usize) -> Self { + Self { + default_match_limit: match_limit, + default_ban_length: ban_length, + stats: crate::util::HashMap::default(), + } + } + + fn get_stats(&mut self, rule: String) -> &mut RuleStats { + self.stats.entry(rule).or_insert_with(|| RuleStats { + iteration: 0, + times_applied: 0, + banned_until: 0, + times_banned: 0, + match_limit: self.default_match_limit, + ban_length: self.default_ban_length, + }) + } + } + + impl Scheduler for BackOffScheduler { + fn can_stop(&mut self, rules: &[&str], _ruleset: &str) -> bool { + let stats = &mut self.stats; + let mut banned: Vec<(&str, RuleStats)> = rules + .iter() + .filter_map(|rule| { + let s = stats.remove(*rule).unwrap(); + if s.banned_until > s.iteration { + Some((*rule, s)) + } else { + None + } + }) + .collect(); + + let result = if banned.is_empty() { + true + } else { + let min_delta = banned + .iter() + .map(|(_, s)| s.banned_until - s.iteration) + .min() + .unwrap(); + for (_, s) in &mut banned { + s.banned_until -= min_delta; + } + false + }; + + for (rule, s) in banned { + stats.insert(rule.to_owned(), s); + } + result + } + + fn filter_matches(&mut self, rule: &str, _ruleset: &str, matches: &mut Matches) -> bool { + let stats = self.get_stats(rule.to_owned()); + stats.iteration += 1; + + if stats.iteration < stats.banned_until { + return false; + } + + let threshold = stats + .match_limit + .checked_shl(stats.times_banned as u32) + .unwrap(); + if matches.match_size() > threshold { + let ban_length = stats.ban_length << stats.times_banned; + stats.times_banned += 1; + stats.banned_until = stats.iteration + ban_length; + false + } else { + stats.times_applied += 1; + matches.choose_all(); + true + } + } + } + + // The real backoff scheduler must reach the same saturated e-graph as an + // unscheduled run on a benchmark that triggers repeated banning, under term + // encoding (maintenance runs between scheduled steps). + #[test] + fn test_backoff_scheduler_matches_unscheduled_term_encoding() { + // An AC-rewrite closure over four numbers: enough matches that a small + // match limit forces `comm`/`assoc` to be banned and later unbanned. + let program = r#" + (sort Math) + (constructor Num (i64) Math) + (constructor Add (Math Math) Math) + (Add (Num 1) (Add (Num 2) (Add (Num 3) (Num 4)))) + (ruleset rw) + (rule ((Add x y)) ((union (Add x y) (Add y x))) + :name "comm" :ruleset rw) + (rule ((Add (Add x y) z)) ((union (Add (Add x y) z) (Add x (Add y z)))) + :name "assoc" :ruleset rw) + (rule ((Add x (Add y z))) ((union (Add x (Add y z)) (Add (Add x y) z))) + :name "assoc2" :ruleset rw) + "#; + + let mut base = EGraph::new_with_term_encoding(); + base.parse_and_run_program(None, program).unwrap(); + base.parse_and_run_program(None, "(run-schedule (saturate rw))") + .unwrap(); + + let mut sched = EGraph::new_with_term_encoding(); + let id = sched.add_scheduler(Box::new(BackOffScheduler::new(4, 2))); + sched.parse_and_run_program(None, program).unwrap(); + let mut banned_at_least_once = false; + let mut iters = 0; + loop { + let report = sched.step_rules_with_scheduler(id, "rw").unwrap(); + iters += 1; + assert!(iters < 10_000, "backoff scheduler did not converge"); + // A step that finds matches but makes no database change means a rule + // was banned this iteration. + if !report.updated && !report.can_stop { + banned_at_least_once = true; + } + if report.can_stop { + break; + } + } + assert!( + banned_at_least_once, + "benchmark never triggered a ban; test would not exercise backoff" + ); + + for ctor in ["Num", "Add"] { + assert_eq!( + base.get_size(ctor), + sched.get_size(ctor), + "e-node count differs for {ctor}" + ); + } + + let checks = [ + "(check (= (Add (Num 1) (Add (Num 2) (Add (Num 3) (Num 4)))) \ + (Add (Num 4) (Add (Num 3) (Add (Num 2) (Num 1))))))", + "(check (= (Add (Add (Num 1) (Num 2)) (Add (Num 3) (Num 4))) \ + (Add (Num 1) (Add (Num 2) (Add (Num 3) (Num 4))))))", + "(check (!= (Num 1) (Num 2)))", + ]; + for check in checks { + base.parse_and_run_program(None, check) + .unwrap_or_else(|e| panic!("unscheduled check failed: {check}: {e}")); + sched + .parse_and_run_program(None, check) + .unwrap_or_else(|e| panic!("backoff check failed: {check}: {e}")); + } + } + #[derive(Clone, Default)] struct DelayStopScheduler { can_stop_calls: usize, diff --git a/src/typechecking.rs b/src/typechecking.rs index 53cf2b405..cf9d97a8f 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -491,6 +491,9 @@ impl EGraph { NCommand::RunSchedule(schedule) => ResolvedNCommand::RunSchedule( self.type_info.typecheck_schedule(symbol_gen, schedule)?, ), + NCommand::SetRebuildSchedule(schedule) => ResolvedNCommand::SetRebuildSchedule( + self.type_info.typecheck_schedule(symbol_gen, schedule)?, + ), NCommand::Pop(span, n) => ResolvedNCommand::Pop(span.clone(), *n), NCommand::Push(n) => ResolvedNCommand::Push(*n), NCommand::AddRuleset(span, ruleset) => {