Merge pull request #928 from saulshanabrook/codex/proof-selector-merge-main

Avoid redundant proof merge rows for eager current values

Author: saulshanabrook

src/proofs/proof_encoding.rs

@@ -10,6 +10,11 @@ pub(crate) struct EncodingState {
     pub uf_function: HashMap<String, String>,
     /// Maps sort name -> proof function name (set from :internal-proof-func annotation).
     pub proof_func_parent: HashMap<String, String>,
+    /// Function name -> (hidden current-value function, input arity). The
+    /// current function uses the original eager backend merge, so cleanup can
+    /// discard stale proof-view candidates whenever the current value already
+    /// has a proof witness.
+    pub merge_current: HashMap<String, (String, usize)>,
     pub term_header_added: bool,
     // TODO this is very ugly- we should separate out a typechecking struct
     // since we didn't need an entire e-graph
@@ -30,6 +35,7 @@ impl EncodingState {
             uf_parent: HashMap::default(),
             uf_function: HashMap::default(),
             proof_func_parent: HashMap::default(),
+            merge_current: HashMap::default(),
             term_header_added: false,
             original_typechecking: None,
             proofs_enabled: false,
@@ -255,13 +261,25 @@ impl<'a> ProofInstrumentor<'a> {
             .as_ref()
             .unwrap_or_else(|| panic!("Proofs don't support :no-merge"));
 
+        let current_name = self
+            .egraph
+            .parser
+            .symbol_gen
+            .fresh(&format!("{name}Current"));
+        self.egraph
+            .proof_state
+            .merge_current
+            .insert(name.clone(), (current_name.clone(), child_names.len()));
+
         let fresh_name = self.egraph.parser.symbol_gen.fresh("merge_rule");
         let cleanup_name = self.egraph.parser.symbol_gen.fresh("merge_cleanup");
+        let current_cleanup_name = self.egraph.parser.symbol_gen.fresh("merge_current_cleanup");
 
         let p1_fresh = self.egraph.parser.symbol_gen.fresh("p1");
         let p2_fresh = self.egraph.parser.symbol_gen.fresh("p2");
         let view_name = self.view_name(&fdecl.name);
         let rebuilding_cleanup_ruleset = self.proof_names().rebuilding_cleanup_ruleset_name.clone();
+        let input_sorts = ListDisplay(&fdecl.schema.input, " ");
         let proof_query = if self.egraph.proof_state.proofs_enabled {
             // View is a function with proof output; bind proof variables
             format!(
@@ -310,7 +328,11 @@ impl<'a> ProofInstrumentor<'a> {
         // The first runs the merge function adding a new row.
         // The second deletes rows with old values for the old variable, while the third deletes rows with new values for the new variable.
         format!(
-            "(sort {fresh_sort})
+            "(function {current_name} ({input_sorts}) {output_sort}
+                    :merge {merge_fn}
+                    :unextractable
+                    :internal-hidden)
+                 (sort {fresh_sort})
                  (constructor {cleanup_constructor} ({output_sort} {output_sort}) {fresh_sort} :internal-hidden)
                  (rule (({view_name} {child_names_str} old)
                         ({view_name} {child_names_str} new)
@@ -333,6 +355,13 @@ impl<'a> ProofInstrumentor<'a> {
                        ((delete ({view_name} {child_names_str} old)))
                         :ruleset {rebuilding_cleanup_ruleset}
                         :name \"{cleanup_name}\")
+                 (rule ((= selected ({current_name} {child_names_str}))
+                        ({view_name} {child_names_str} selected)
+                        ({view_name} {child_names_str} old)
+                        (!= selected old))
+                       ((delete ({view_name} {child_names_str} old)))
+                        :ruleset {rebuilding_cleanup_ruleset}
+                        :name \"{current_cleanup_name}\")
                 ",
         )
     }
@@ -999,7 +1028,16 @@ impl<'a> ProofInstrumentor<'a> {
     /// View is always a function (returning Proof or Unit).
     fn update_view(&mut self, fname: &str, args: &[String], proof: &str) -> String {
         let view_name = self.view_name(fname);
-        format!("(set ({view_name} {}) {proof})", ListDisplay(args, " "))
+        let view_update = format!("(set ({view_name} {}) {proof})", ListDisplay(args, " "));
+        if let Some((current_name, input_arity)) =
+            self.egraph.proof_state.merge_current.get(fname).cloned()
+            && args.len() == input_arity + 1
+        {
+            let inputs = ListDisplay(&args[..input_arity], " ");
+            let output = &args[input_arity];
+            return format!("{view_update}\n(set ({current_name} {inputs}) {output})");
+        }
+        view_update
     }
 
     /// Return some code adding to the view and term tables.

src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap

@@ -13,13 +13,15 @@ expression: snapshot
 (function __addView (i64 i64 i64) Unit :merge old :unextractable :internal-term-constructor add)
 (constructor __to_delete_add (i64 i64) __view :internal-hidden)
 (constructor __to_subsume_add (i64 i64) __view :internal-hidden)
+(function __addCurrent (i64 i64) i64 :merge old :unextractable :internal-hidden)
 (sort __mergecleanupsort)
 (constructor __mergecleanup (i64 i64) __mergecleanupsort :internal-hidden)
 (rule ((__addView c0_ c1_ old)
        (__addView c0_ c1_ new)
        (!= old new)
        (= (ordering-max old new) new))
       ((set (__addView c0_ c1_ old) ())
+       (set (__addCurrent c0_ c1_) old)
        (__mergecleanup old old)
        (__mergecleanup old new))
         :ruleset __rebuilding :name "__merge_rule")
@@ -29,6 +31,12 @@ expression: snapshot
        (!= merged old))
       ((delete (__addView c0_ c1_ old)))
         :ruleset __rebuilding_cleanup :name "__merge_cleanup")
+(rule ((= selected (__addCurrent c0_ c1_))
+       (__addView c0_ c1_ selected)
+       (__addView c0_ c1_ old)
+       (!= selected old))
+      ((delete (__addView c0_ c1_ old)))
+        :ruleset __rebuilding_cleanup :name "__merge_current_cleanup")
 (rule ((__to_delete_add c0_ c1_)
        (__addView c0_ c1_ out))
       ((delete (__addView c0_ c1_ out))

tests/proof_selector_merge.rs

@@ -0,0 +1,35 @@
+use egglog::EGraph;
+
+#[test]
+fn selector_merge_cleanup_uses_existing_current_proof() {
+    let mut candidates = String::new();
+    for i in 0..200 {
+        candidates.push_str(&format!("(Candidate (K) {i})\n"));
+    }
+
+    let program = format!(
+        r#"
+        (sort Key)
+        (constructor K () Key)
+        (function Best (Key) i64 :merge old)
+        (relation Candidate (Key i64))
+        (rule ((Candidate k v))
+              ((set (Best k) v)))
+        {candidates}
+        (run 1)
+        (check (= (Best (K)) 0))
+        "#
+    );
+
+    let mut egraph = EGraph::new_with_proofs().with_proof_testing();
+    egraph.parse_and_run_program(None, &program).unwrap();
+
+    let merge_rule_matches: usize = egraph
+        .get_overall_run_report()
+        .num_matches_per_rule
+        .iter()
+        .filter(|(rule, _)| rule.as_ref().contains("merge_rule"))
+        .map(|(_, matches)| *matches)
+        .sum();
+    assert_eq!(merge_rule_matches, 0);
+}