From acbb79e0594c5982b53ac69a38cd6038952e6ff8 Mon Sep 17 00:00:00 2001 From: Lukas Wirth Date: Fri, 13 Dec 2024 17:24:35 +0100 Subject: [PATCH] Pool `ActiveQuery`s in the query stack --- src/active_query.rs | 77 ++++++++++++++++++++++++++++++++++++++++++--- src/zalsa_local.rs | 26 ++++++--------- 2 files changed, 82 insertions(+), 21 deletions(-) diff --git a/src/active_query.rs b/src/active_query.rs index 752756c9e..9be5f1d7a 100644 --- a/src/active_query.rs +++ b/src/active_query.rs @@ -1,3 +1,5 @@ +use std::{mem, ops}; + use rustc_hash::FxHashMap; use super::zalsa_local::{EdgeKind, QueryEdges, QueryOrigin, QueryRevisions}; @@ -57,7 +59,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, @@ -71,6 +73,28 @@ 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, + // These two are cleared via `mem::take`` when popped off as revisions. + tracked_struct_ids: _, + accumulated: _, + } = self; + *database_key_index = new_database_key_index; + *durability = Durability::MAX; + *changed_at = Revision::start(); + input_outputs.clear(); + *untracked_read = false; + *cycle = None; + disambiguator_map.clear(); + } + pub(super) fn add_read( &mut self, input: DependencyIndex, @@ -106,11 +130,11 @@ impl ActiveQuery { self.input_outputs.contains(&(EdgeKind::Output, key)) } - pub(crate) fn into_revisions(self) -> QueryRevisions { + fn take_revisions(&mut self) -> QueryRevisions { let input_outputs = if self.input_outputs.is_empty() { EMPTY_DEPENDENCIES.clone() } else { - self.input_outputs.into_iter().collect() + self.input_outputs.iter().copied().collect() }; let edges = QueryEdges::new(input_outputs); @@ -125,8 +149,8 @@ impl ActiveQuery { changed_at: self.changed_at, origin, durability: self.durability, - tracked_struct_ids: self.tracked_struct_ids, - accumulated: self.accumulated, + tracked_struct_ids: mem::take(&mut self.tracked_struct_ids), + accumulated: mem::take(&mut self.accumulated), } } @@ -167,3 +191,46 @@ impl ActiveQuery { result } } + +#[derive(Debug, Default)] +pub(crate) struct QueryStack { + stack: Vec, + 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 { + if self.len == 0 { + return None; + } + self.len -= 1; + Some(self.stack[self.len].take_revisions()) + } +} diff --git a/src/zalsa_local.rs b/src/zalsa_local.rs index 6988d6537..47add12ea 100644 --- a/src/zalsa_local.rs +++ b/src/zalsa_local.rs @@ -2,7 +2,7 @@ use rustc_hash::FxHashMap; use tracing::debug; use crate::accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues}; -use crate::active_query::ActiveQuery; +use crate::active_query::QueryStack; use crate::durability::Durability; use crate::key::DatabaseKeyIndex; use crate::key::DependencyIndex; @@ -37,7 +37,7 @@ pub struct ZalsaLocal { /// /// Unwinding note: pushes onto this vector must be popped -- even /// during unwinding. - query_stack: RefCell>, + query_stack: RefCell, /// 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(&self, c: impl FnOnce(&mut Vec) -> R) -> R { - c(self.query_stack.borrow_mut().as_mut()) + pub(crate) fn with_query_stack(&self, c: impl FnOnce(&mut QueryStack) -> R) -> R { + c(&mut self.query_stack.borrow_mut()) } fn query_in_progress(&self) -> bool { @@ -499,7 +499,7 @@ pub(crate) struct ActiveQueryGuard<'me> { } impl ActiveQueryGuard<'_> { - fn pop_helper(&self) -> ActiveQuery { + fn pop_helper(&self) -> QueryRevisions { self.local_state.with_query_stack(|stack| { // Sanity check: pushes and pops should be balanced. assert_eq!(stack.len(), self.push_len); @@ -507,7 +507,7 @@ impl ActiveQueryGuard<'_> { stack.last().unwrap().database_key_index, self.database_key_index ); - stack.pop().unwrap() + stack.pop_into_revisions().unwrap() }) } @@ -522,7 +522,7 @@ impl ActiveQueryGuard<'_> { } /// Invoked when the query has successfully completed execution. - pub(crate) fn complete(self) -> ActiveQuery { + fn complete(self) -> QueryRevisions { let query = self.pop_helper(); std::mem::forget(self); query @@ -533,13 +533,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