Pool ActiveQuerys in the query stack

Author: Veykril

src/active_query.rs

@@ -1,4 +1,5 @@
 use std::ops::Not;
+use std::{mem, ops};
 
 use super::zalsa_local::{QueryEdges, QueryOrigin, QueryRevisions};
 use crate::accumulator::accumulated_map::AtomicInputAccumulatedValues;
@@ -63,7 +64,7 @@ pub(crate) struct ActiveQuery {
 }
 
 impl ActiveQuery {
-    pub(super) fn new(database_key_index: DatabaseKeyIndex) -> Self {
+    pub(crate) fn new(database_key_index: DatabaseKeyIndex) -> Self {
         ActiveQuery {
             database_key_index,
             durability: Durability::MAX,
@@ -78,6 +79,32 @@ impl ActiveQuery {
         }
     }
 
+    fn reset(&mut self, new_database_key_index: DatabaseKeyIndex) {
+        let Self {
+            database_key_index,
+            durability,
+            changed_at,
+            input_outputs,
+            untracked_read,
+            cycle,
+            disambiguator_map,
+            tracked_struct_ids,
+            accumulated,
+            accumulated_inputs,
+        } = self;
+        *database_key_index = new_database_key_index;
+        *durability = Durability::MAX;
+        *changed_at = Revision::start();
+        input_outputs.clear();
+        *untracked_read = false;
+        *cycle = None;
+        *accumulated_inputs = Default::default();
+        disambiguator_map.clear();
+        // These two are cleared via `mem::take`` when popped off as revisions.
+        debug_assert!(tracked_struct_ids.is_empty());
+        _ = accumulated;
+    }
+
     pub(super) fn add_read(
         &mut self,
         input: InputDependencyIndex,
@@ -113,24 +140,38 @@ impl ActiveQuery {
         self.input_outputs.contains(&QueryEdge::Output(key))
     }
 
-    pub(crate) fn into_revisions(self) -> QueryRevisions {
-        let edges = QueryEdges::new(self.input_outputs);
-        let origin = if self.untracked_read {
+    fn take_revisions(&mut self) -> QueryRevisions {
+        let &mut Self {
+            database_key_index: _,
+            cycle: _,
+            disambiguator_map: _,
+            durability,
+            changed_at,
+            untracked_read,
+            accumulated_inputs,
+            ref mut input_outputs,
+            ref mut tracked_struct_ids,
+            ref mut accumulated,
+        } = self;
+
+        let edges = QueryEdges::new(input_outputs.drain(..));
+        let origin = if untracked_read {
             QueryOrigin::DerivedUntracked(edges)
         } else {
             QueryOrigin::Derived(edges)
         };
-        let accumulated = self
-            .accumulated
+        let accumulated = accumulated
             .is_empty()
             .not()
-            .then(|| Box::new(self.accumulated));
+            .then(|| Box::new(mem::take(accumulated)));
+        let tracked_struct_ids = mem::take(tracked_struct_ids);
+        let accumulated_inputs = AtomicInputAccumulatedValues::new(accumulated_inputs);
         QueryRevisions {
-            changed_at: self.changed_at,
+            changed_at,
+            durability,
             origin,
-            durability: self.durability,
-            tracked_struct_ids: self.tracked_struct_ids,
-            accumulated_inputs: AtomicInputAccumulatedValues::new(self.accumulated_inputs),
+            tracked_struct_ids,
+            accumulated_inputs,
             accumulated,
         }
     }
@@ -166,3 +207,46 @@ impl ActiveQuery {
         self.disambiguator_map.disambiguate(key)
     }
 }
+
+#[derive(Debug, Default)]
+pub(crate) struct QueryStack {
+    stack: Vec<ActiveQuery>,
+    len: usize,
+}
+
+impl ops::Deref for QueryStack {
+    type Target = [ActiveQuery];
+
+    fn deref(&self) -> &Self::Target {
+        &self.stack[..self.len]
+    }
+}
+
+impl ops::DerefMut for QueryStack {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.stack[..self.len]
+    }
+}
+
+impl QueryStack {
+    pub(crate) fn push_new_query(&mut self, database_key_index: DatabaseKeyIndex) {
+        if self.len < self.stack.len() {
+            self.stack[self.len].reset(database_key_index);
+        } else {
+            self.stack.push(ActiveQuery::new(database_key_index));
+        }
+        self.len += 1;
+    }
+
+    pub(crate) fn len(&self) -> usize {
+        self.len
+    }
+
+    pub(crate) fn pop_into_revisions(&mut self) -> Option<QueryRevisions> {
+        if self.len == 0 {
+            return None;
+        }
+        self.len -= 1;
+        Some(self.stack[self.len].take_revisions())
+    }
+}

src/function/execute.rs

@@ -45,7 +45,6 @@ where
 
         // Query was not previously executed, or value is potentially
         // stale, or value is absent. Let's execute!
-        let database_key_index = active_query.database_key_index;
         let id = database_key_index.key_index;
         let value = match Cycle::catch(|| C::execute(db, C::id_to_input(db, id))) {
             Ok(v) => v,

src/tracked_struct.rs

@@ -326,6 +326,10 @@ impl DisambiguatorMap {
         disambiguator.0 += 1;
         result
     }
+
+    pub fn clear(&mut self) {
+        self.map.clear()
+    }
 }
 
 impl<C> IngredientImpl<C>

src/zalsa_local.rs

@@ -4,7 +4,7 @@ use tracing::debug;
 use crate::accumulator::accumulated_map::{
     AccumulatedMap, AtomicInputAccumulatedValues, InputAccumulatedValues,
 };
-use crate::active_query::ActiveQuery;
+use crate::active_query::QueryStack;
 use crate::durability::Durability;
 use crate::key::{DatabaseKeyIndex, InputDependencyIndex, OutputDependencyIndex};
 use crate::runtime::StampedValue;
@@ -37,7 +37,7 @@ pub struct ZalsaLocal {
     ///
     /// Unwinding note: pushes onto this vector must be popped -- even
     /// during unwinding.
-    query_stack: RefCell<Vec<ActiveQuery>>,
+    query_stack: RefCell<QueryStack>,
 
     /// Stores the most recent page for a given ingredient.
     /// This is thread-local to avoid contention.
@@ -47,7 +47,7 @@ pub struct ZalsaLocal {
 impl ZalsaLocal {
     pub(crate) fn new() -> Self {
         ZalsaLocal {
-            query_stack: RefCell::new(vec![]),
+            query_stack: RefCell::new(QueryStack::default()),
             most_recent_pages: RefCell::new(FxHashMap::default()),
         }
     }
@@ -88,7 +88,7 @@ impl ZalsaLocal {
     #[inline]
     pub(crate) fn push_query(&self, database_key_index: DatabaseKeyIndex) -> ActiveQueryGuard<'_> {
         let mut query_stack = self.query_stack.borrow_mut();
-        query_stack.push(ActiveQuery::new(database_key_index));
+        query_stack.push_new_query(database_key_index);
         ActiveQueryGuard {
             local_state: self,
             database_key_index,
@@ -97,8 +97,8 @@ impl ZalsaLocal {
     }
 
     /// Executes a closure within the context of the current active query stacks.
-    pub(crate) fn with_query_stack<R>(&self, c: impl FnOnce(&mut Vec<ActiveQuery>) -> R) -> R {
-        c(self.query_stack.borrow_mut().as_mut())
+    pub(crate) fn with_query_stack<R>(&self, c: impl FnOnce(&mut QueryStack) -> R) -> R {
+        c(&mut self.query_stack.borrow_mut())
     }
 
     /// Returns the index of the active query along with its *current* durability/changed-at
@@ -485,15 +485,15 @@ pub(crate) struct ActiveQueryGuard<'me> {
 }
 
 impl ActiveQueryGuard<'_> {
-    fn pop_helper(&self) -> ActiveQuery {
+    fn pop_impl(&self) -> QueryRevisions {
         self.local_state.with_query_stack(|stack| {
             // Sanity check: pushes and pops should be balanced.
             assert_eq!(stack.len(), self.push_len);
             debug_assert_eq!(
                 stack.last().unwrap().database_key_index,
                 self.database_key_index
             );
-            stack.pop().unwrap()
+            stack.pop_into_revisions().unwrap()
         })
     }
 
@@ -508,8 +508,8 @@ impl ActiveQueryGuard<'_> {
     }
 
     /// Invoked when the query has successfully completed execution.
-    pub(crate) fn complete(self) -> ActiveQuery {
-        let query = self.pop_helper();
+    fn complete(self) -> QueryRevisions {
+        let query = self.pop_impl();
         std::mem::forget(self);
         query
     }
@@ -519,13 +519,7 @@ impl ActiveQueryGuard<'_> {
     /// query's execution.
     #[inline]
     pub(crate) fn pop(self) -> QueryRevisions {
-        // Extract accumulated inputs.
-        let popped_query = self.complete();
-
-        // If this frame were a cycle participant, it would have unwound.
-        assert!(popped_query.cycle.is_none());
-
-        popped_query.into_revisions()
+        self.complete()
     }
 
     /// If the active query is registered as a cycle participant, remove and
@@ -538,6 +532,6 @@ impl ActiveQueryGuard<'_> {
 
 impl Drop for ActiveQueryGuard<'_> {
     fn drop(&mut self) {
-        self.pop_helper();
+        self.pop_impl();
     }
 }