Use packed cached hashes for faster table mutations

core-relations/src/action/mod.rs

@@ -3,25 +3,27 @@
 //! This allows us to execute the "right-hand-side" of a rule. The
 //! implementation here is optimized to execute on a batch of rows at a time.
 use std::{
-    ops::Deref,
+    hash::Hasher,
     sync::{
         Arc,
         atomic::{AtomicBool, Ordering},
     },
 };
 
 use crate::{
-    common::HashMap,
     free_join::{invoke_batch, invoke_batch_assign},
     numeric_id::{DenseIdMap, NumericId},
 };
 use egglog_concurrency::NotificationList;
-use smallvec::SmallVec;
+use hashbrown::HashTable;
+use rustc_hash::FxHasher;
 
 use crate::{
     BaseValues, ContainerValues, ExternalFunctionId, WrappedTable,
     common::Value,
-    free_join::{CounterId, Counters, ExternalFunctions, TableId, TableInfo, Variable},
+    free_join::{
+        CounterId, CounterReservation, Counters, ExternalFunctions, TableId, TableInfo, Variable,
+    },
     pool::{Clear, Pooled, with_pool_set},
     table_spec::{ColumnId, MutationBuffer},
 };
@@ -289,36 +291,89 @@ pub(crate) struct ExtractedBinding {
     pub(crate) vals: Pooled<Vec<Value>>,
 }
 
+#[derive(Clone, Copy)]
+struct PredictedEntry {
+    hash: u64,
+    index: usize,
+    table: TableId,
+    key_arity: u32,
+}
+
 #[derive(Default)]
 pub(crate) struct PredictedVals {
-    #[allow(clippy::type_complexity)]
-    data: HashMap<(TableId, SmallVec<[Value; 3]>), Pooled<Vec<Value>>>,
+    index: HashTable<PredictedEntry>,
+    values: Vec<Value>,
 }
 
 impl Clear for PredictedVals {
     fn reuse(&self) -> bool {
-        self.data.capacity() > 0
+        self.index.capacity() > 0 || self.values.capacity() > 0
     }
     fn clear(&mut self) {
-        self.data.clear()
+        self.index.clear();
+        self.values.clear();
     }
     fn bytes(&self) -> usize {
-        self.data.capacity()
-            * (std::mem::size_of::<(TableId, SmallVec<[Value; 3]>)>()
-                + std::mem::size_of::<Pooled<Vec<Value>>>())
+        self.index.capacity() * std::mem::size_of::<PredictedEntry>()
+            + self.values.capacity() * std::mem::size_of::<Value>()
     }
 }
 
 impl PredictedVals {
-    pub(crate) fn get_val(
+    fn hash(table: TableId, key: &[Value]) -> u64 {
+        let mut hasher = FxHasher::default();
+        hasher.write_u32(table.rep());
+        hasher.write_usize(key.len());
+        for value in key {
+            hasher.write_u32(value.rep());
+        }
+        hasher.finish()
+    }
+
+    fn row(&self, entry: PredictedEntry, row_arity: usize) -> &[Value] {
+        &self.values[entry.index..entry.index + row_arity]
+    }
+
+    fn matches(&self, entry: &PredictedEntry, hash: u64, table: TableId, key: &[Value]) -> bool {
+        if entry.hash != hash || entry.table != table || entry.key_arity as usize != key.len() {
+            return false;
+        }
+        self.values[entry.index..entry.index + key.len()] == *key
+    }
+
+    pub(crate) fn get_or_insert_with(
         &mut self,
         table: TableId,
         key: &[Value],
-        default: impl FnOnce() -> Pooled<Vec<Value>>,
-    ) -> impl Deref<Target = Pooled<Vec<Value>>> + '_ {
-        self.data
-            .entry((table, SmallVec::from_slice(key)))
-            .or_insert_with(default)
+        row_arity: usize,
+        default: impl FnOnce(&mut Vec<Value>, usize),
+    ) -> (&[Value], bool) {
+        let hash = Self::hash(table, key);
+        if let Some(entry) = self
+            .index
+            .find(hash, |entry| self.matches(entry, hash, table, key))
+            .copied()
+        {
+            return (self.row(entry, row_arity), false);
+        }
+
+        let entry = PredictedEntry {
+            hash,
+            index: self.values.len(),
+            table,
+            key_arity: u32::try_from(key.len()).expect("predicted key arity must fit in u32"),
+        };
+        self.values.reserve(row_arity);
+        let row_start = self.values.len();
+        self.values.extend_from_slice(key);
+        default(&mut self.values, row_start);
+        assert_eq!(
+            self.values.len() - row_start,
+            row_arity,
+            "predicted row builder produced the wrong arity"
+        );
+        self.index.insert_unique(hash, entry, |entry| entry.hash);
+        (self.row(entry, row_arity), true)
     }
 }
 
@@ -360,6 +415,7 @@ pub(crate) struct DbView<'a> {
 pub struct ExecutionState<'a> {
     pub(crate) predicted: PredictedVals,
     pub(crate) db: DbView<'a>,
+    counter_reservations: CounterReservations,
     buffers: MutationBuffers<'a>,
     /// Whether any mutations have been staged via this ExecutionState.
     pub(crate) changed: bool,
@@ -368,6 +424,19 @@ pub struct ExecutionState<'a> {
     stop_match: Arc<AtomicBool>,
 }
 
+#[derive(Default)]
+struct CounterReservations {
+    ranges: DenseIdMap<CounterId, CounterReservation>,
+}
+
+impl CounterReservations {
+    fn next(&mut self, counters: &Counters, ctr: CounterId) -> usize {
+        self.ranges
+            .get_or_insert(ctr, || counters.take_reservation(ctr))
+            .next()
+    }
+}
+
 /// Copyable query-side view of an [`ExecutionState`]. Join tasks only need
 /// immutable database access and the shared early-stop flag; an owned state is
 /// materialized lazily when an action batch actually executes.
@@ -382,6 +451,7 @@ impl<'db> ExecutionStateSeed<'db, '_> {
         ExecutionState {
             predicted: Default::default(),
             db: self.db,
+            counter_reservations: Default::default(),
             buffers: MutationBuffers::new(self.db.notification_list, Default::default()),
             changed: false,
             stop_match: Arc::clone(self.stop_match),
@@ -439,6 +509,7 @@ impl Clone for ExecutionState<'_> {
         ExecutionState {
             predicted: Default::default(),
             db: self.db,
+            counter_reservations: Default::default(),
             buffers: self.buffers.clone(),
             changed: false,
             stop_match: Arc::clone(&self.stop_match),
@@ -454,6 +525,7 @@ impl<'a> ExecutionState<'a> {
         ExecutionState {
             predicted: Default::default(),
             db,
+            counter_reservations: Default::default(),
             buffers: MutationBuffers::new(db.notification_list, buffers),
             changed: false,
             stop_match: Arc::new(AtomicBool::new(false)),
@@ -500,8 +572,8 @@ impl<'a> ExecutionState<'a> {
         self.db.external_funcs[func].invoke(self, args)
     }
 
-    pub fn inc_counter(&self, ctr: CounterId) -> usize {
-        self.db.counters.inc(ctr)
+    pub fn inc_counter(&mut self, ctr: CounterId) -> usize {
+        self.counter_reservations.next(self.db.counters, ctr)
     }
 
     pub fn read_counter(&self, ctr: CounterId) -> usize {
@@ -550,44 +622,31 @@ impl<'a> ExecutionState<'a> {
         if let Some(row) = self.db.table_info[table].table.get_row(key) {
             return row.vals;
         }
-        Pooled::cloned(
+        let row_arity = key.len() + vals.len();
+        let counters = self.db.counters;
+        let counter_reservations = &mut self.counter_reservations;
+        let (row, inserted) =
             self.predicted
-                .get_val(table, key, || {
-                    Self::construct_new_row(
-                        &self.db,
-                        &mut self.buffers,
-                        &mut self.changed,
-                        table,
-                        key,
-                        vals,
-                    )
-                })
-                .deref(),
-        )
-    }
-
-    fn construct_new_row(
-        db: &DbView,
-        buffers: &mut MutationBuffers,
-        changed: &mut bool,
-        table: TableId,
-        key: &[Value],
-        vals: impl ExactSizeIterator<Item = MergeVal>,
-    ) -> Pooled<Vec<Value>> {
+                .get_or_insert_with(table, key, row_arity, |values, _| {
+                    for val in vals {
+                        values.push(match val {
+                            MergeVal::Counter(ctr) => {
+                                Value::from_usize(counter_reservations.next(counters, ctr))
+                            }
+                            MergeVal::Constant(c) => c,
+                        });
+                    }
+                });
+        if inserted {
+            self.buffers
+                .lazy_init(table, || self.db.table_info[table].table.new_buffer());
+            self.buffers.stage_insert(table, row);
+            self.changed = true;
+        }
         with_pool_set(|ps| {
-            let mut new = ps.get::<Vec<Value>>();
-            new.reserve(key.len() + vals.len());
-            new.extend_from_slice(key);
-            for val in vals {
-                new.push(match val {
-                    MergeVal::Counter(ctr) => Value::from_usize(db.counters.inc(ctr)),
-                    MergeVal::Constant(c) => c,
-                })
-            }
-            buffers.lazy_init(table, || db.table_info[table].table.new_buffer());
-            buffers.stage_insert(table, &new);
-            *changed = true;
-            new
+            let mut result = ps.get::<Vec<Value>>();
+            result.extend_from_slice(row);
+            result
         })
     }
 
@@ -603,16 +662,28 @@ impl<'a> ExecutionState<'a> {
         if let Some(val) = self.db.table_info[table].table.get_row_column(key, col) {
             return val;
         }
-        self.predicted.get_val(table, key, || {
-            Self::construct_new_row(
-                &self.db,
-                &mut self.buffers,
-                &mut self.changed,
-                table,
-                key,
-                vals,
-            )
-        })[col.index()]
+        let row_arity = key.len() + vals.len();
+        let counters = self.db.counters;
+        let counter_reservations = &mut self.counter_reservations;
+        let (row, inserted) =
+            self.predicted
+                .get_or_insert_with(table, key, row_arity, |values, _| {
+                    for val in vals {
+                        values.push(match val {
+                            MergeVal::Counter(ctr) => {
+                                Value::from_usize(counter_reservations.next(counters, ctr))
+                            }
+                            MergeVal::Constant(c) => c,
+                        });
+                    }
+                });
+        if inserted {
+            self.buffers
+                .lazy_init(table, || self.db.table_info[table].table.new_buffer());
+            self.buffers.stage_insert(table, row);
+            self.changed = true;
+        }
+        row[col.index()]
     }
 
     /// Trigger early stopping by setting the stop_match flag.
@@ -683,7 +754,6 @@ impl ExecutionState<'_> {
                 dst_col,
                 dst_var,
             } => {
-                let pool = with_pool_set(|ps| ps.get_pool::<Vec<Value>>().clone());
                 self.buffers.lazy_init(*table_id, || {
                     self.db.table_info[*table_id].table.new_buffer()
                 });
@@ -707,42 +777,37 @@ impl ExecutionState<'_> {
                         //
                         // We avoid doing this more than once by using the
                         // `predicted` map.
-                        let prediction_key = (
-                            *table_id,
-                            SmallVec::<[Value; 3]>::from_slice(key.as_slice()),
-                        );
                         let buffers = &mut self.buffers;
                         // Bind some mutable references because the closure passed
                         // to or_insert_with is `move`.
-                        let ctrs = &self.db.counters;
+                        let ctrs = self.db.counters;
+                        let counter_reservations = &mut self.counter_reservations;
                         let bindings = &bindings;
-                        let pool = pool.clone();
-                        let row =
-                            self.predicted
-                                .data
-                                .entry(prediction_key)
-                                .or_insert_with(move || {
-                                    let mut row = pool.get();
-                                    row.extend_from_slice(key.as_slice());
-                                    // Extend the key with the default values.
-                                    row.reserve(default.len());
-                                    for val in default {
-                                        let val = match val {
-                                            WriteVal::QueryEntry(QueryEntry::Const(c)) => *c,
-                                            WriteVal::QueryEntry(QueryEntry::Var(v)) => {
-                                                bindings[*v][offset]
-                                            }
-                                            WriteVal::IncCounter(ctr) => {
-                                                Value::from_usize(ctrs.inc(*ctr))
-                                            }
-                                            WriteVal::CurrentVal(ix) => row[*ix],
-                                        };
-                                        row.push(val)
-                                    }
-                                    // Insert it into the table.
-                                    buffers.stage_insert(*table_id, &row);
-                                    row
-                                });
+                        let row_arity = key.as_slice().len() + default.len();
+                        let (row, inserted) = self.predicted.get_or_insert_with(
+                            *table_id,
+                            key.as_slice(),
+                            row_arity,
+                            |values, row_start| {
+                                // Extend the key with the default values.
+                                for val in default {
+                                    let val = match val {
+                                        WriteVal::QueryEntry(QueryEntry::Const(c)) => *c,
+                                        WriteVal::QueryEntry(QueryEntry::Var(v)) => {
+                                            bindings[*v][offset]
+                                        }
+                                        WriteVal::IncCounter(ctr) => {
+                                            Value::from_usize(counter_reservations.next(ctrs, *ctr))
+                                        }
+                                        WriteVal::CurrentVal(ix) => values[row_start + *ix],
+                                    };
+                                    values.push(val)
+                                }
+                            },
+                        );
+                        if inserted {
+                            buffers.stage_insert(*table_id, row);
+                        }
                         row[dst_col.index()]
                     });
                 });