Merge pull request #445 from thaliaarchi/update-hashbrown

Update hashbrown
egraphs-good · Oct 24, 2024 · af49ae2 · af49ae2
2 parents ae69e65 + 90e6e69
commit af49ae2
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 70 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -26,7 +26,7 @@ bin = ["dep:clap", "dep:env_logger", "egraph-serialize/serde", "dep:serde_json"]
 wasm-bindgen = ["instant/wasm-bindgen", "dep:getrandom"]
 
 [dependencies]
-hashbrown = { version = "0.14", features = ["raw"] }
+hashbrown = { version = "0.15" }
 indexmap = "2.0"
 instant = "0.1"
 log = "0.4"

diff --git a/src/function/table.rs b/src/function/table.rs
@@ -32,7 +32,7 @@ use std::{
     ops::Range,
 };
 
-use hashbrown::raw::RawTable;
+use hashbrown::HashTable;
 
 use super::binary_search::binary_search_table_by_key;
 use crate::{util::BuildHasher as BH, TupleOutput, Value, ValueVec};
@@ -51,11 +51,11 @@ struct TableOffset {
 pub(crate) struct Table {
     max_ts: u32,
     n_stale: usize,
-    table: RawTable<TableOffset>,
+    table: HashTable<TableOffset>,
     pub(crate) vals: Vec<(Input, TupleOutput)>,
 }
 
-/// Used for the RawTable probe sequence.
+/// Used for the HashTable probe sequence.
 macro_rules! search_for {
     ($slf:expr, $hash:expr, $inp:expr) => {
         |to| {
@@ -97,19 +97,18 @@ impl Table {
 
     /// Rehashes the table, invalidating any offsets stored into the table.
     pub(crate) fn rehash(&mut self) {
-        let mut src = 0usize;
         let mut dst = 0usize;
         self.table.clear();
         self.vals.retain(|(inp, _)| {
             if inp.live() {
                 let hash = hash_values(inp.data());
+                let to = TableOffset { hash, off: dst };
                 self.table
-                    .insert(hash, TableOffset { hash, off: dst }, |to| to.hash);
-                src += 1;
+                    .entry(hash, |to2| to2 == &to, |to2| to2.hash)
+                    .insert(to);
                 dst += 1;
                 true
             } else {
-                src += 1;
                 false
             }
         });
@@ -120,16 +119,16 @@ impl Table {
     /// table.
     pub(crate) fn get(&self, inputs: &[Value]) -> Option<&TupleOutput> {
         let hash = hash_values(inputs);
-        let TableOffset { off, .. } = self.table.get(hash, search_for!(self, hash, inputs))?;
-        debug_assert!(self.vals[*off].0.live());
-        Some(&self.vals[*off].1)
+        let &TableOffset { off, .. } = self.table.find(hash, search_for!(self, hash, inputs))?;
+        debug_assert!(self.vals[off].0.live());
+        Some(&self.vals[off].1)
     }
 
     pub(crate) fn get_mut(&mut self, inputs: &[Value]) -> Option<&mut TupleOutput> {
         let hash: u64 = hash_values(inputs);
-        let TableOffset { off, .. } = self.table.get(hash, search_for!(self, hash, inputs))?;
-        debug_assert!(self.vals[*off].0.live());
-        Some(&mut self.vals[*off].1)
+        let &TableOffset { off, .. } = self.table.find(hash, search_for!(self, hash, inputs))?;
+        debug_assert!(self.vals[off].0.live());
+        Some(&mut self.vals[off].1)
     }
 
     /// Insert the given data into the table at the given timestamp. Return the
@@ -161,7 +160,7 @@ impl Table {
         self.max_ts = ts;
         let hash = hash_values(inputs);
         if let Some(TableOffset { off, .. }) =
-            self.table.get_mut(hash, search_for!(self, hash, inputs))
+            self.table.find_mut(hash, search_for!(self, hash, inputs))
         {
             let (inp, prev) = &mut self.vals[*off];
             let prev_subsumed = prev.subsumed;
@@ -193,14 +192,13 @@ impl Table {
                 subsumed,
             },
         ));
-        self.table.insert(
+        let to = TableOffset {
             hash,
-            TableOffset {
-                hash,
-                off: new_offset,
-            },
-            |off| off.hash,
-        );
+            off: new_offset,
+        };
+        self.table
+            .entry(hash, |to2| to2 == &to, |to2| to2.hash)
+            .insert(to);
     }
 
     /// One more than the maximum (potentially) valid offset into the table.
@@ -237,13 +235,11 @@ impl Table {
     /// removed.
     pub(crate) fn remove(&mut self, inp: &[Value], ts: u32) -> bool {
         let hash = hash_values(inp);
-        let entry = if let Some(entry) = self.table.remove_entry(hash, search_for!(self, hash, inp))
-        {
-            entry
-        } else {
+        let Ok(entry) = self.table.find_entry(hash, search_for!(self, hash, inp)) else {
             return false;
         };
-        self.vals[entry.off].0.stale_at = ts;
+        let (TableOffset { off, .. }, _) = entry.remove();
+        self.vals[off].0.stale_at = ts;
         self.n_stale += 1;
         true
     }

diff --git a/src/termdag.rs b/src/termdag.rs
@@ -1,6 +1,6 @@
 use crate::{
     ast::Literal,
-    util::{HashMap, HashSet},
+    util::{HashMap, HashSet, IndexSet},
     Expr, GenericExpr, Symbol,
 };
 
@@ -21,14 +21,8 @@ pub enum Term {
 /// A hashconsing arena for [`Term`]s.
 #[derive(Clone, PartialEq, Eq, Debug, Default)]
 pub struct TermDag {
-    // think of nodes as a map from indices to Terms.
-    // invariant: the nodes map and the hashcons map are inverses.
-    // note that this implies:
-    // - no duplicates in nodes
-    // - every element of node is a key in hashcons
-    // - every key of hashcons is in nodes
-    pub nodes: Vec<Term>,
-    pub hashcons: HashMap<Term, TermId>,
+    /// A bidirectional map between deduplicated `Term`s and indices.
+    nodes: IndexSet<Term>,
 }
 
 #[macro_export]
@@ -54,14 +48,14 @@ impl TermDag {
     ///
     /// Panics if the term does not already exist in this [TermDag].
     pub fn lookup(&self, node: &Term) -> TermId {
-        *self.hashcons.get(node).unwrap()
+        self.nodes.get_index_of(node).unwrap()
     }
 
     /// Convert the given id to the corresponding term.
     ///
     /// Panics if the id is not valid.
-    pub fn get(&self, id: TermId) -> Term {
-        self.nodes[id].clone()
+    pub fn get(&self, id: TermId) -> &Term {
+        self.nodes.get_index(id).unwrap()
     }
 
     /// Make and return a [`Term::App`] with the given head symbol and children,
@@ -97,10 +91,8 @@ impl TermDag {
     }
 
     fn add_node(&mut self, node: &Term) {
-        if self.hashcons.get(node).is_none() {
-            let idx = self.nodes.len();
-            self.nodes.push(node.clone());
-            self.hashcons.insert(node.clone(), idx);
+        if self.nodes.get(node).is_none() {
+            self.nodes.insert(node.clone());
         }
     }
 
@@ -138,10 +130,7 @@ impl TermDag {
             Term::App(op, args) => {
                 let args: Vec<_> = args
                     .iter()
-                    .map(|a| {
-                        let term = self.get(*a);
-                        self.term_to_expr(&term)
-                    })
+                    .map(|a| self.term_to_expr(self.get(*a)))
                     .collect();
                 Expr::call_no_span(*op, args)
             }
@@ -215,7 +204,7 @@ mod tests {
         //     x, y, (g x y), and the root call to f
         // so we can compute expected answer by hand:
         assert_eq!(
-            td.nodes,
+            td.nodes.as_slice().iter().cloned().collect::<Vec<_>>(),
             vec![
                 Term::Var("x".into()),
                 Term::Var("y".into()),
@@ -236,7 +225,7 @@ mod tests {
         let (td, t) = parse_term(s);
         match_term_app!(t; {
             ("f", [_, x, _, _]) =>
-                assert_eq!(td.term_to_expr(&td.get(*x)), ast::GenericExpr::Var(DUMMY_SPAN.clone(), Symbol::new("x"))),
+                assert_eq!(td.term_to_expr(td.get(*x)), ast::GenericExpr::Var(DUMMY_SPAN.clone(), Symbol::new("x"))),
             (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!())
         })
     }