refactor util hashmap

egraphs-good · Oct 28, 2024 · 9f34210 · 9f34210
1 parent 0a47d4e
commit 9f34210
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 26 deletions.
diff --git a/src/extract.rs b/src/extract.rs
@@ -1,9 +1,7 @@
-use hashbrown::hash_map::Entry;
-
 use crate::ast::Symbol;
 use crate::termdag::{Term, TermDag};
 use crate::util::HashMap;
-use crate::{ArcSort, EGraph, Function, Id, Value};
+use crate::{ArcSort, EGraph, Function, HEntry, Id, Value};
 
 pub type Cost = usize;
 
@@ -194,11 +192,11 @@ impl<'a> Extractor<'a> {
 
                             let id = self.egraph.find(&func.schema.output, output.value).bits;
                             match self.costs.entry(id) {
-                                Entry::Vacant(e) => {
+                                HEntry::Vacant(e) => {
                                     did_something = true;
                                     e.insert(make_new_pair());
                                 }
-                                Entry::Occupied(mut e) => {
+                                HEntry::Occupied(mut e) => {
                                     if new_cost < e.get().0 {
                                         did_something = true;
                                         e.insert(make_new_pair());

diff --git a/src/gj.rs b/src/gj.rs
@@ -1,6 +1,7 @@
 use indexmap::map::Entry;
 use log::log_enabled;
 use smallvec::SmallVec;
+use util::HashMap;
 
 use crate::{core::*, function::index::Offset, *};
 use std::{
@@ -791,24 +792,16 @@ impl Debug for LazyTrie {
     }
 }
 
-#[cfg(feature = "nondeterministic")]
-type SparseMap = HashMap<Value, LazyTrie>;
-#[cfg(feature = "nondeterministic")]
-type SEntry<'a, A, B, D> = hashbrown::hash_map::Entry<'a, A, B, D>;
-#[cfg(not(feature = "nondeterministic"))]
-type SparseMap = IndexMap<Value, LazyTrie>;
-#[cfg(not(feature = "nondeterministic"))]
-type SEntry<'a, A, B> = Entry<'a, A, B>;
 type RowIdx = u32;
 
 #[derive(Debug)]
 enum LazyTrieInner {
     Borrowed {
         index: Rc<ColumnIndex>,
-        map: SparseMap,
+        map: HashMap<Value, LazyTrie>,
     },
     Delayed(SmallVec<[RowIdx; 4]>),
-    Sparse(SparseMap),
+    Sparse(HashMap<Value, LazyTrie>),
 }
 
 impl Default for LazyTrie {
@@ -856,7 +849,7 @@ impl LazyTrie {
         let this = unsafe { &mut *self.0.get() };
         match this {
             LazyTrieInner::Borrowed { index, .. } => {
-                let mut map = SparseMap::with_capacity_and_hasher(index.len(), Default::default());
+                let mut map = HashMap::with_capacity_and_hasher(index.len(), Default::default());
                 map.extend(index.iter().filter_map(|(v, ixs)| {
                     LazyTrie::from_indexes(access.filter_live(ixs)).map(|trie| (v, trie))
                 }));
@@ -894,8 +887,8 @@ impl LazyTrie {
             LazyTrieInner::Borrowed { index, map } => {
                 let ixs = index.get(&value)?;
                 match map.entry(value) {
-                    SEntry::Occupied(o) => Some(o.into_mut()),
-                    SEntry::Vacant(v) => {
+                    HEntry::Occupied(o) => Some(o.into_mut()),
+                    HEntry::Vacant(v) => {
                         Some(v.insert(LazyTrie::from_indexes(access.filter_live(ixs))?))
                     }
                 }
@@ -941,20 +934,20 @@ impl<'a> TrieAccess<'a> {
     #[cold]
     fn make_trie_inner(&self, idxs: &[RowIdx]) -> LazyTrieInner {
         let arity = self.function.schema.input.len();
-        let mut map = SparseMap::default();
+        let mut map: HashMap<Value, LazyTrie> = HashMap::default();
         let mut insert = |i: usize, tup: &[Value], out: &TupleOutput, val: Value| {
             if self.timestamp_range.contains(&out.timestamp)
                 && self.constraints.iter().all(|c| c.check(tup, out))
             {
                 match map.entry(val) {
-                    SEntry::Occupied(mut e) => {
+                    HEntry::Occupied(mut e) => {
                         if let LazyTrieInner::Delayed(ref mut v) = e.get_mut().0.get_mut() {
                             v.push(i as RowIdx)
                         } else {
                             unreachable!()
                         }
                     }
-                    SEntry::Vacant(e) => {
+                    HEntry::Vacant(e) => {
                         e.insert(LazyTrie(UnsafeCell::new(LazyTrieInner::Delayed(
                             smallvec::smallvec![i as RowIdx,],
                         ))));

diff --git a/src/typechecking.rs b/src/typechecking.rs
@@ -1,6 +1,5 @@
 use crate::{core::CoreRule, *};
 use ast::Rule;
-use hashbrown::hash_map::Entry;
 
 #[derive(Clone, Debug)]
 pub struct FuncType {
@@ -64,8 +63,8 @@ impl TypeInfo {
     pub fn add_presort<S: Presort>(&mut self, span: Span) -> Result<(), TypeError> {
         let name = S::presort_name();
         match self.presorts.entry(name) {
-            Entry::Occupied(_) => Err(TypeError::SortAlreadyBound(name, span)),
-            Entry::Vacant(e) => {
+            HEntry::Occupied(_) => Err(TypeError::SortAlreadyBound(name, span)),
+            HEntry::Vacant(e) => {
                 e.insert(S::make_sort);
                 self.reserved_primitives.extend(S::reserved_primitives());
                 Ok(())
@@ -77,8 +76,8 @@ impl TypeInfo {
         let name = sort.name();
 
         match self.sorts.entry(name) {
-            Entry::Occupied(_) => Err(TypeError::SortAlreadyBound(name, span)),
-            Entry::Vacant(e) => {
+            HEntry::Occupied(_) => Err(TypeError::SortAlreadyBound(name, span)),
+            HEntry::Vacant(e) => {
                 e.insert(sort.clone());
                 sort.register_primitives(self);
                 Ok(())

diff --git a/src/util.rs b/src/util.rs
@@ -8,9 +8,24 @@ use crate::*;
 
 pub(crate) type BuildHasher = std::hash::BuildHasherDefault<rustc_hash::FxHasher>;
 
+/// Use an index map by default everywhere.
+/// We could fix the seed, but symbol generation is not determinisic so
+/// this doesn't fix the problem.
+#[cfg(not(feature = "nondeterministic"))]
+pub(crate) type HashMap<K, V> = indexmap::IndexMap<K, V, BuildHasher>;
+#[cfg(feature = "nondeterministic")]
 pub(crate) type HashMap<K, V> = hashbrown::HashMap<K, V, BuildHasher>;
+
+#[cfg(not(feature = "nondeterministic"))]
+pub(crate) type HashSet<K> = indexmap::IndexSet<K, BuildHasher>;
+#[cfg(feature = "nondeterministic")]
 pub(crate) type HashSet<K> = hashbrown::HashSet<K, BuildHasher>;
 
+#[cfg(feature = "nondeterministic")]
+pub(crate) type HEntry<'a, A, B, D> = hashbrown::hash_map::Entry<'a, A, B, D>;
+#[cfg(not(feature = "nondeterministic"))]
+pub(crate) type HEntry<'a, A, B> = Entry<'a, A, B>;
+
 pub type IndexMap<K, V> = indexmap::IndexMap<K, V, BuildHasher>;
 pub type IndexSet<K> = indexmap::IndexSet<K, BuildHasher>;