Defer explanation costs

src/approx_bitset.rs

@@ -0,0 +1,57 @@
+use egg::Id;
+use hashbrown::HashSet;
+use std::fmt::Debug;
+
+type Elt = u128;
+pub const BITS: usize = Elt::BITS as usize;
+
+#[derive(Debug)]
+pub(crate) struct ApproxBitSet(Elt);
+
+pub(crate) trait IdSet: Debug {
+    fn empty() -> Self;
+
+    fn insert(&mut self, val: Id);
+
+    fn may_contain(&self, val: Id) -> bool;
+
+    fn might_confuse(id1: Id, id2: Id) -> bool;
+}
+
+impl IdSet for ApproxBitSet {
+    fn empty() -> Self {
+        ApproxBitSet(0)
+    }
+
+    fn insert(&mut self, val: Id) {
+        let shift = usize::from(val) % BITS;
+        self.0 |= 1 << shift
+    }
+
+    fn may_contain(&self, val: Id) -> bool {
+        let shift = usize::from(val) % BITS;
+        self.0 & (1 << shift) != 0
+    }
+
+    fn might_confuse(id1: Id, id2: Id) -> bool {
+        usize::from(id1) % BITS == usize::from(id2) % BITS
+    }
+}
+
+impl IdSet for HashSet<Id> {
+    fn empty() -> Self {
+        HashSet::default()
+    }
+
+    fn insert(&mut self, val: Id) {
+        self.insert(val);
+    }
+
+    fn may_contain(&self, val: Id) -> bool {
+        self.contains(&val)
+    }
+
+    fn might_confuse(id1: Id, id2: Id) -> bool {
+        id1 == id2
+    }
+}

src/egraph.rs

@@ -85,7 +85,10 @@ impl<D> EGraph<D> {
             // so we do not need to distinguish between nodes that are
             // equivalent modulo ground equalities
             |_, id, _| Some(id),
-            |this, id1, id2| this.explain.union(id1, id2, Justification::CONGRUENCE),
+            |this, existing_id, new_id| {
+                this.explain
+                    .union(existing_id, Justification::CONGRUENCE, existing_id, new_id)
+            },
             |this, id, _| {
                 this.explain.add(id);
                 mk_data(id)
@@ -100,10 +103,17 @@ impl<D> EGraph<D> {
         justification: Justification,
         mut merge: impl FnMut(&mut D, D),
     ) {
-        self.inner.raw_union(id1, id2, |data1, _, _, data2, _, _| {
-            merge(data1, data2);
-            self.explain.union(id1, id2, justification)
-        })
+        let id1_root = self.inner.find_mut(id1);
+        self.inner
+            .raw_union(id1, id2, |data1, _, _, data2, old_root, _| {
+                merge(data1, data2);
+                let (old_id, new_id) = if id1_root == old_root {
+                    (id1, id2)
+                } else {
+                    (id2, id1)
+                };
+                self.explain.union(old_root, justification, old_id, new_id)
+            })
     }
 
     pub fn rebuild(&mut self, mut merge: impl FnMut(&mut D, D)) {