use tracing::debug;
use crate::durability::Durability;
use crate::runtime::ActiveQuery;
use crate::runtime::Revision;
use crate::Cycle;
use crate::DatabaseKeyIndex;
use std::cell::RefCell;
use triomphe::Arc;
pub(super) struct LocalState {
query_stack: RefCell<Option<Vec<ActiveQuery>>>,
}
#[derive(Debug, Clone)]
pub(crate) struct QueryRevisions {
pub(crate) changed_at: Revision,
pub(crate) durability: Durability,
pub(crate) inputs: QueryInputs,
}
#[derive(Debug, Clone)]
pub(crate) enum QueryInputs {
Tracked { inputs: Arc<[DatabaseKeyIndex]> },
NoInputs,
Untracked,
}
impl Default for LocalState {
fn default() -> Self {
LocalState {
query_stack: RefCell::new(Some(Vec::new())),
}
}
}
impl LocalState {
#[inline]
pub(super) fn push_query(&self, database_key_index: DatabaseKeyIndex) -> ActiveQueryGuard<'_> {
let mut query_stack = self.query_stack.borrow_mut();
let query_stack = query_stack.as_mut().expect("local stack taken");
query_stack.push(ActiveQuery::new(database_key_index));
ActiveQueryGuard {
local_state: self,
database_key_index,
push_len: query_stack.len(),
}
}
fn with_query_stack<R>(&self, c: impl FnOnce(&mut Vec<ActiveQuery>) -> R) -> R {
c(self
.query_stack
.borrow_mut()
.as_mut()
.expect("query stack taken"))
}
pub(super) fn query_in_progress(&self) -> bool {
self.with_query_stack(|stack| !stack.is_empty())
}
pub(super) fn active_query(&self) -> Option<DatabaseKeyIndex> {
self.with_query_stack(|stack| {
stack
.last()
.map(|active_query| active_query.database_key_index)
})
}
pub(super) fn report_query_read_and_unwind_if_cycle_resulted(
&self,
input: DatabaseKeyIndex,
durability: Durability,
changed_at: Revision,
) {
debug!(
"report_query_read_and_unwind_if_cycle_resulted(input={:?}, durability={:?}, changed_at={:?})",
input, durability, changed_at
);
self.with_query_stack(|stack| {
if let Some(top_query) = stack.last_mut() {
top_query.add_read(input, durability, changed_at);
if let Some(cycle) = &top_query.cycle {
cycle.clone().throw()
}
}
})
}
pub(super) fn report_untracked_read(&self, current_revision: Revision) {
self.with_query_stack(|stack| {
if let Some(top_query) = stack.last_mut() {
top_query.add_untracked_read(current_revision);
}
})
}
pub(super) fn report_synthetic_read(&self, durability: Durability, revision: Revision) {
self.with_query_stack(|stack| {
if let Some(top_query) = stack.last_mut() {
top_query.add_synthetic_read(durability, revision);
}
})
}
pub(super) fn take_query_stack(&self) -> Vec<ActiveQuery> {
assert!(
self.query_stack.borrow().is_some(),
"query stack already taken"
);
self.query_stack.take().unwrap()
}
pub(super) fn restore_query_stack(&self, stack: Vec<ActiveQuery>) {
assert!(self.query_stack.borrow().is_none(), "query stack not taken");
self.query_stack.replace(Some(stack));
}
}
impl std::panic::RefUnwindSafe for LocalState {}
pub(crate) struct ActiveQueryGuard<'me> {
local_state: &'me LocalState,
push_len: usize,
database_key_index: DatabaseKeyIndex,
}
impl ActiveQueryGuard<'_> {
fn pop_helper(&self) -> ActiveQuery {
self.local_state.with_query_stack(|stack| {
assert_eq!(stack.len(), self.push_len);
debug_assert_eq!(
stack.last().unwrap().database_key_index,
self.database_key_index
);
stack.pop().unwrap()
})
}
pub(super) fn complete(self) -> ActiveQuery {
let query = self.pop_helper();
std::mem::forget(self);
query
}
#[inline]
pub(crate) fn pop(self) -> QueryRevisions {
let popped_query = self.complete();
assert!(popped_query.cycle.is_none());
popped_query.revisions()
}
pub(crate) fn take_cycle(&self) -> Option<Cycle> {
self.local_state
.with_query_stack(|stack| stack.last_mut()?.cycle.take())
}
}
impl Drop for ActiveQueryGuard<'_> {
fn drop(&mut self) {
self.pop_helper();
}
}