diff --git a/src/active_query.rs b/src/active_query.rs index a3b55d8dc..baca24e69 100644 --- a/src/active_query.rs +++ b/src/active_query.rs @@ -377,7 +377,7 @@ pub struct Backtrace(Box<[CapturedQuery]>); impl Backtrace { pub fn capture() -> Option { crate::with_attached_database(|db| { - db.zalsa_local().with_query_stack(|stack| { + db.zalsa_local().try_with_query_stack(|stack| { Backtrace( stack .iter() @@ -392,7 +392,7 @@ impl Backtrace { .collect(), ) }) - }) + })? } } diff --git a/src/zalsa_local.rs b/src/zalsa_local.rs index a45a5c55b..322c771c6 100644 --- a/src/zalsa_local.rs +++ b/src/zalsa_local.rs @@ -110,15 +110,32 @@ impl ZalsaLocal { } } - /// Executes a closure within the context of the current active query stacks. + /// Executes a closure within the context of the current active query stacks (mutable). #[inline(always)] - pub(crate) fn with_query_stack( + pub(crate) fn with_query_stack_mut( &self, c: impl UnwindSafe + FnOnce(&mut QueryStack) -> R, ) -> R { c(&mut self.query_stack.borrow_mut()) } + #[inline(always)] + pub(crate) fn with_query_stack(&self, c: impl UnwindSafe + FnOnce(&QueryStack) -> R) -> R { + c(&mut self.query_stack.borrow()) + } + + #[inline(always)] + pub(crate) fn try_with_query_stack( + &self, + c: impl UnwindSafe + FnOnce(&QueryStack) -> R, + ) -> Option { + self.query_stack + .try_borrow() + .ok() + .as_ref() + .map(|stack| c(stack)) + } + /// Returns the index of the active query along with its *current* durability/changed-at /// information. As the query continues to execute, naturally, that information may change. pub(crate) fn active_query(&self) -> Option<(DatabaseKeyIndex, Stamp)> { @@ -137,7 +154,7 @@ impl ZalsaLocal { index: IngredientIndex, value: A, ) -> Result<(), ()> { - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { if let Some(top_query) = stack.last_mut() { top_query.accumulate(index, value); Ok(()) @@ -149,7 +166,7 @@ impl ZalsaLocal { /// Add an output to the current query's list of dependencies pub(crate) fn add_output(&self, entity: DatabaseKeyIndex) { - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { if let Some(top_query) = stack.last_mut() { top_query.add_output(entity) } @@ -158,7 +175,7 @@ impl ZalsaLocal { /// Check whether `entity` is an output of the currently active query (if any) pub(crate) fn is_output_of_active_query(&self, entity: DatabaseKeyIndex) -> bool { - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { if let Some(top_query) = stack.last_mut() { top_query.is_output(entity) } else { @@ -182,7 +199,7 @@ impl ZalsaLocal { "report_tracked_read(input={:?}, durability={:?}, changed_at={:?})", input, durability, changed_at ); - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { if let Some(top_query) = stack.last_mut() { top_query.add_read( input, @@ -208,7 +225,7 @@ impl ZalsaLocal { "report_tracked_read(input={:?}, durability={:?}, changed_at={:?})", input, durability, changed_at ); - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { if let Some(top_query) = stack.last_mut() { top_query.add_read_simple(input, durability, changed_at); } @@ -222,7 +239,7 @@ impl ZalsaLocal { /// * `current_revision`, the current revision #[inline(always)] pub(crate) fn report_untracked_read(&self, current_revision: Revision) { - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { if let Some(top_query) = stack.last_mut() { top_query.add_untracked_read(current_revision); } @@ -234,7 +251,7 @@ impl ZalsaLocal { // FIXME: Use or remove this. #[allow(dead_code)] pub(crate) fn report_synthetic_read(&self, durability: Durability, revision: Revision) { - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { if let Some(top_query) = stack.last_mut() { top_query.add_synthetic_read(durability, revision); } @@ -253,7 +270,7 @@ impl ZalsaLocal { /// * the disambiguator index #[track_caller] pub(crate) fn disambiguate(&self, key: IdentityHash) -> (Stamp, Disambiguator) { - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { let top_query = stack.last_mut().expect( "cannot create a tracked struct disambiguator outside of a tracked function", ); @@ -274,7 +291,7 @@ impl ZalsaLocal { #[track_caller] pub(crate) fn store_tracked_struct_id(&self, identity: Identity, id: Id) { - self.with_query_stack(|stack| { + self.with_query_stack_mut(|stack| { let top_query = stack .last_mut() .expect("cannot store a tracked struct ID outside of a tracked function"); @@ -481,7 +498,7 @@ pub(crate) struct ActiveQueryGuard<'me> { impl ActiveQueryGuard<'_> { /// Initialize the tracked struct ids with the values from the prior execution. pub(crate) fn seed_tracked_struct_ids(&self, tracked_struct_ids: &IdentityMap) { - self.local_state.with_query_stack(|stack| { + self.local_state.with_query_stack_mut(|stack| { #[cfg(debug_assertions)] assert_eq!(stack.len(), self.push_len); let frame = stack.last_mut().unwrap(); @@ -495,7 +512,7 @@ impl ActiveQueryGuard<'_> { where I: IntoIterator + UnwindSafe, { - self.local_state.with_query_stack(|stack| { + self.local_state.with_query_stack_mut(|stack| { #[cfg(debug_assertions)] assert_eq!(stack.len(), self.push_len); let frame = stack.last_mut().unwrap(); @@ -508,7 +525,7 @@ impl ActiveQueryGuard<'_> { /// Invoked when the query has successfully completed execution. fn complete(self) -> QueryRevisions { - let query = self.local_state.with_query_stack(|stack| { + let query = self.local_state.with_query_stack_mut(|stack| { stack.pop_into_revisions( self.database_key_index, #[cfg(debug_assertions)] @@ -530,7 +547,7 @@ impl ActiveQueryGuard<'_> { impl Drop for ActiveQueryGuard<'_> { fn drop(&mut self) { - self.local_state.with_query_stack(|stack| { + self.local_state.with_query_stack_mut(|stack| { stack.pop( self.database_key_index, #[cfg(debug_assertions)]