refactor(perf): Pool ActiveQuerys in the query stack (#629)

* Do not pass ownership of the `QueryStack` in `DependencyGraph::block_on`

* Pool `ActiveQuery`s in the query stack

* Don't do unnecessary work when discaring the query guard

* Gate debug only field by `debug_assertions`

* Reduce visibility

Author: Veykril

src/accumulator/accumulated_map.rs

@@ -48,6 +48,10 @@ impl AccumulatedMap {
     pub fn is_empty(&self) -> bool {
         self.map.is_empty()
     }
+
+    pub fn clear(&mut self) {
+        self.map.clear()
+    }
 }
 
 /// Tracks whether any input read during a query's execution has any accumulated values.

src/active_query.rs

@@ -1,7 +1,9 @@
 use std::ops::Not;
+use std::{mem, ops};
 
 use super::zalsa_local::{QueryEdges, QueryOrigin, QueryRevisions};
 use crate::accumulator::accumulated_map::AtomicInputAccumulatedValues;
+use crate::runtime::Stamp;
 use crate::tracked_struct::{DisambiguatorMap, IdentityHash, IdentityMap};
 use crate::zalsa_local::QueryEdge;
 use crate::{
@@ -13,18 +15,19 @@ use crate::{
     tracked_struct::Disambiguator,
     Revision,
 };
+use crate::{Accumulator, IngredientIndex};
 
 #[derive(Debug)]
 pub(crate) struct ActiveQuery {
     /// What query is executing
     pub(crate) database_key_index: DatabaseKeyIndex,
 
     /// Minimum durability of inputs observed so far.
-    pub(crate) durability: Durability,
+    durability: Durability,
 
     /// Maximum revision of all inputs observed. If we observe an
     /// untracked read, this will be set to the most recent revision.
-    pub(crate) changed_at: Revision,
+    changed_at: Revision,
 
     /// Inputs: Set of subqueries that were accessed thus far.
     /// Outputs: Tracks values written by this query. Could be...
@@ -52,32 +55,17 @@ pub(crate) struct ActiveQuery {
 
     /// Stores the values accumulated to the given ingredient.
     /// The type of accumulated value is erased but known to the ingredient.
-    pub(crate) accumulated: AccumulatedMap,
+    accumulated: AccumulatedMap,
 
     /// [`InputAccumulatedValues::Empty`] if any input read during the query's execution
     /// has any accumulated values.
-    pub(super) accumulated_inputs: InputAccumulatedValues,
+    accumulated_inputs: InputAccumulatedValues,
 
     /// Provisional cycle results that this query depends on.
-    pub(crate) cycle_heads: CycleHeads,
+    cycle_heads: CycleHeads,
 }
 
 impl ActiveQuery {
-    pub(super) fn new(database_key_index: DatabaseKeyIndex) -> Self {
-        ActiveQuery {
-            database_key_index,
-            durability: Durability::MAX,
-            changed_at: Revision::start(),
-            input_outputs: FxIndexSet::default(),
-            untracked_read: false,
-            disambiguator_map: Default::default(),
-            tracked_struct_ids: Default::default(),
-            accumulated: Default::default(),
-            accumulated_inputs: Default::default(),
-            cycle_heads: Default::default(),
-        }
-    }
-
     pub(super) fn add_read(
         &mut self,
         input: DatabaseKeyIndex,
@@ -116,6 +104,10 @@ impl ActiveQuery {
         self.changed_at = self.changed_at.max(revision);
     }
 
+    pub(super) fn accumulate(&mut self, index: IngredientIndex, value: impl Accumulator) {
+        self.accumulated.accumulate(index, value);
+    }
+
     /// Adds a key to our list of outputs.
     pub(super) fn add_output(&mut self, key: DatabaseKeyIndex) {
         self.input_outputs.insert(QueryEdge::Output(key));
@@ -126,30 +118,195 @@ 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 {
+    pub(super) fn disambiguate(&mut self, key: IdentityHash) -> Disambiguator {
+        self.disambiguator_map.disambiguate(key)
+    }
+
+    pub(super) fn stamp(&self) -> Stamp {
+        Stamp {
+            value: (),
+            durability: self.durability,
+            changed_at: self.changed_at,
+        }
+    }
+}
+
+impl ActiveQuery {
+    fn new(database_key_index: DatabaseKeyIndex) -> Self {
+        ActiveQuery {
+            database_key_index,
+            durability: Durability::MAX,
+            changed_at: Revision::start(),
+            input_outputs: FxIndexSet::default(),
+            untracked_read: false,
+            disambiguator_map: Default::default(),
+            tracked_struct_ids: Default::default(),
+            accumulated: Default::default(),
+            accumulated_inputs: Default::default(),
+            cycle_heads: Default::default(),
+        }
+    }
+
+    fn top_into_revisions(&mut self) -> QueryRevisions {
+        let &mut Self {
+            database_key_index: _,
+            durability,
+            changed_at,
+            ref mut input_outputs,
+            untracked_read,
+            ref mut disambiguator_map,
+            ref mut tracked_struct_ids,
+            ref mut accumulated,
+            accumulated_inputs,
+            ref mut cycle_heads,
+        } = self;
+
+        let edges = QueryEdges::new(input_outputs.drain(..));
+        let origin = if untracked_read {
             QueryOrigin::DerivedUntracked(edges)
         } else {
             QueryOrigin::Derived(edges)
         };
-        let accumulated = self
-            .accumulated
+        disambiguator_map.clear();
+        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);
+        let cycle_heads = mem::take(cycle_heads);
         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,
-            cycle_heads: self.cycle_heads,
+            cycle_heads,
         }
     }
 
-    pub(super) fn disambiguate(&mut self, key: IdentityHash) -> Disambiguator {
-        self.disambiguator_map.disambiguate(key)
+    fn clear(&mut self) {
+        let Self {
+            database_key_index: _,
+            durability: _,
+            changed_at: _,
+            input_outputs,
+            untracked_read: _,
+            disambiguator_map,
+            tracked_struct_ids,
+            accumulated,
+            accumulated_inputs: _,
+            cycle_heads,
+        } = self;
+        input_outputs.clear();
+        disambiguator_map.clear();
+        tracked_struct_ids.clear();
+        accumulated.clear();
+        *cycle_heads = Default::default();
+    }
+
+    fn reset_for(&mut self, new_database_key_index: DatabaseKeyIndex) {
+        let Self {
+            database_key_index,
+            durability,
+            changed_at,
+            input_outputs,
+            untracked_read,
+            disambiguator_map,
+            tracked_struct_ids,
+            accumulated,
+            accumulated_inputs,
+            cycle_heads,
+        } = self;
+        *database_key_index = new_database_key_index;
+        *durability = Durability::MAX;
+        *changed_at = Revision::start();
+        *untracked_read = false;
+        *accumulated_inputs = Default::default();
+        debug_assert!(
+            input_outputs.is_empty(),
+            "`ActiveQuery::clear` or `ActiveQuery::into_revisions` should've been called"
+        );
+        debug_assert!(
+            disambiguator_map.is_empty(),
+            "`ActiveQuery::clear` or `ActiveQuery::into_revisions` should've been called"
+        );
+        debug_assert!(
+            tracked_struct_ids.is_empty(),
+            "`ActiveQuery::clear` or `ActiveQuery::into_revisions` should've been called"
+        );
+        debug_assert!(
+            cycle_heads.is_empty(),
+            "`ActiveQuery::clear` or `ActiveQuery::into_revisions` should've been called"
+        );
+        debug_assert!(
+            accumulated.is_empty(),
+            "`ActiveQuery::clear` or `ActiveQuery::into_revisions` should've been called"
+        );
+    }
+}
+
+#[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_for(database_key_index);
+        } else {
+            self.stack.push(ActiveQuery::new(database_key_index));
+        }
+        self.len += 1;
+    }
+
+    #[cfg(debug_assertions)]
+    pub(crate) fn len(&self) -> usize {
+        self.len
+    }
+
+    pub(crate) fn pop_into_revisions(
+        &mut self,
+        key: DatabaseKeyIndex,
+        #[cfg(debug_assertions)] push_len: usize,
+    ) -> QueryRevisions {
+        #[cfg(debug_assertions)]
+        assert_eq!(push_len, self.len(), "unbalanced push/pop");
+        debug_assert_ne!(self.len, 0, "too many pops");
+        self.len -= 1;
+        debug_assert_eq!(
+            self.stack[self.len].database_key_index, key,
+            "unbalanced push/pop"
+        );
+        self.stack[self.len].top_into_revisions()
+    }
+
+    pub(crate) fn pop(&mut self, key: DatabaseKeyIndex, #[cfg(debug_assertions)] push_len: usize) {
+        #[cfg(debug_assertions)]
+        assert_eq!(push_len, self.len(), "unbalanced push/pop");
+        debug_assert_ne!(self.len, 0, "too many pops");
+        self.len -= 1;
+        debug_assert_eq!(
+            self.stack[self.len].database_key_index, key,
+            "unbalanced push/pop"
+        );
+        self.stack[self.len].clear()
     }
 }

src/runtime.rs

@@ -1,5 +1,4 @@
 use std::{
-    mem,
     panic::AssertUnwindSafe,
     sync::atomic::{AtomicBool, Ordering},
     thread::ThreadId,
@@ -196,16 +195,14 @@ impl Runtime {
         // we cannot enter an inconsistent state for this parameter.
         let query_mutex_guard = AssertUnwindSafe(query_mutex_guard);
         let result = local_state.with_query_stack(|stack| {
-            let (new_stack, result) = DependencyGraph::block_on(
+            DependencyGraph::block_on(
                 { dg }.0,
                 thread_id,
                 database_key,
                 other_id,
-                mem::take(stack),
+                stack,
                 { query_mutex_guard }.0,
-            );
-            *stack = new_stack;
-            result
+            )
         });
 
         match result {