use crate::{AsStopReason, Child, Complete, Execution, Outcome, PhaseSpec, PlanError, Pulse};
use alloc::{collections::VecDeque, sync::Arc};
#[must_use = "run stages and hand over the final result; Drop records abandonment"]
pub struct Stages<'a> {
parent: &'a dyn Pulse,
pending: VecDeque<(Child<'a>, Option<StageKey>)>,
state: State,
}
enum State {
Ready,
Running,
Stopped,
}
impl<'a> Stages<'a> {
#[track_caller]
pub fn new(parent: &'a dyn Pulse, parts: &[PhaseSpec<'_>]) -> Self {
let pending = parent
.plan(Execution::Sequence, parts)
.into_iter()
.map(|child| (child, None))
.collect();
Self {
parent,
pending,
state: State::Ready,
}
}
#[track_caller]
pub fn define<const N: usize>(
parent: &'a dyn Pulse,
parts: [PhaseSpec<'_>; N],
) -> (Self, [StageKey; N]) {
let mut stages = Self::new(parent, &parts);
let identity = (N != 0).then(|| Arc::new(()));
let keys = core::array::from_fn(|index| {
let key = StageKey {
identity: Arc::clone(identity.as_ref().expect("nonempty definition")),
index,
};
stages.pending[index].1 = Some(key.clone());
key
});
(stages, keys)
}
#[track_caller]
pub fn begin(&mut self, key: StageKey) -> Stage<'_, 'a> {
self.begin_stage(Some(&key))
}
#[track_caller]
pub fn begin_next(&mut self) -> Stage<'_, 'a> {
self.begin_stage(None)
}
#[track_caller]
fn begin_stage(&mut self, key: Option<&StageKey>) -> Stage<'_, 'a> {
let child = self.begin_child(key);
Stage {
stages: Some(self),
child,
}
}
#[track_caller]
pub fn run<T, E>(&mut self, work: impl FnOnce(&dyn Pulse) -> Result<T, E>) -> Result<T, E>
where
E: AsStopReason,
{
let child = self.begin_child(None);
let result = work(&child);
self.complete_result(child, result)
}
#[track_caller]
pub fn run_classified<T, E>(
&mut self,
is_stop: impl FnOnce(&E) -> bool,
work: impl FnOnce(&dyn Pulse) -> Result<T, E>,
) -> Result<T, E> {
let child = self.begin_child(None);
let result = work(&child);
self.end(child, Outcome::from_result(&result, is_stop));
result
}
fn complete_result<T, E>(&mut self, child: Child<'_>, result: Result<T, E>) -> Result<T, E>
where
E: AsStopReason,
{
let outcome = Outcome::from_result(&result, |error| error.as_stop_reason().is_some());
self.end(child, outcome);
result
}
#[track_caller]
fn begin_child(&mut self, key: Option<&StageKey>) -> Child<'a> {
let child = self.next(key);
if !matches!(self.state, State::Stopped) {
self.state = State::Running;
}
if let Err(error) = child.start() {
self.parent.record_issue(error);
}
child
}
fn end(&mut self, child: Child<'_>, outcome: Outcome) {
child.complete_as(outcome);
if matches!(
outcome,
Outcome::Failed | Outcome::Cancelled | Outcome::Abandoned | Outcome::NotRun
) {
self.state = State::Stopped;
self.resolve_remaining(Outcome::NotRun);
} else if matches!(self.state, State::Running) {
self.state = State::Ready;
}
}
#[track_caller]
pub fn skip(&mut self) {
self.next(None).complete_as(Outcome::Skipped);
}
#[track_caller]
fn next(&mut self, key: Option<&StageKey>) -> Child<'a> {
if matches!(self.state, State::Running) {
self.state = State::Stopped;
self.resolve_remaining(Outcome::NotRun);
}
if matches!(self.state, State::Stopped) {
self.parent.record_issue(PlanError::Finished);
} else if let Some((_, expected)) = self.pending.front() {
if key.is_some_and(|key| {
!expected
.as_ref()
.is_some_and(|expected| expected.matches(key))
}) {
self.parent.record_issue(PlanError::StageMismatch);
} else {
return self.pending.pop_front().expect("checked pending stage").0;
}
} else {
self.parent.record_issue(PlanError::NoMoreStages);
}
Child::untracked(self.parent)
}
fn resolve_remaining(&mut self, outcome: Outcome) {
while let Some((child, _)) = self.pending.pop_front() {
child.complete_inferred(outcome);
}
}
}
impl Complete for Stages<'_> {
fn complete_as(mut self, outcome: Outcome) {
self.resolve_remaining(match outcome {
Outcome::Succeeded | Outcome::Skipped => Outcome::Skipped,
Outcome::Abandoned => Outcome::Abandoned,
_ => Outcome::NotRun,
});
}
}
#[derive(Clone, Debug)]
pub struct StageKey {
identity: Arc<()>,
index: usize,
}
impl StageKey {
fn matches(&self, other: &Self) -> bool {
self.index == other.index && Arc::ptr_eq(&self.identity, &other.identity)
}
}
#[must_use = "hand over the stage result before ?; dropping it records Abandoned"]
pub struct Stage<'s, 'p> {
stages: Option<&'s mut Stages<'p>>,
child: Child<'p>,
}
impl Stage<'_, '_> {
#[track_caller]
pub fn fork<const N: usize>(&self, parts: [PhaseSpec<'_>; N]) -> [Child<'_>; N] {
let mut children = self.fork_children(&parts).into_iter();
core::array::from_fn(|_| children.next().expect("one child per declaration"))
}
#[track_caller]
fn fork_children(&self, parts: &[PhaseSpec<'_>]) -> alloc::vec::Vec<Child<'_>> {
let children = self.child.plan(Execution::ForkJoin, parts);
for child in &children {
if let Err(error) = child.start() {
self.child.record_issue(error);
}
}
children
}
fn close(&mut self, outcome: Outcome) {
if let Some(stages) = self.stages.take() {
let child = core::mem::replace(&mut self.child, Child::inert());
stages.end(child, outcome);
}
}
}
impl core::fmt::Debug for Stage<'_, '_> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Stage")
.field("child", &self.child)
.finish_non_exhaustive()
}
}
impl Complete for Stage<'_, '_> {
fn complete_as(mut self, outcome: Outcome) {
self.close(outcome);
}
}
impl Drop for Stage<'_, '_> {
fn drop(&mut self) {
self.close(Outcome::Abandoned);
}
}
impl crate::Stop for Stage<'_, '_> {
#[inline]
#[track_caller]
fn check(&self) -> Result<(), crate::StopReason> {
self.child.check()
}
#[inline]
fn may_stop(&self) -> bool {
self.child.may_stop()
}
}
impl crate::Report for Stage<'_, '_> {
#[inline]
#[track_caller]
fn advance(&self, completed: u64) {
self.child.advance(completed);
}
#[inline]
fn may_report(&self) -> bool {
self.child.may_report()
}
}
impl Pulse for Stage<'_, '_> {
#[track_caller]
fn record_issue(&self, error: PlanError) {
self.child.record_issue(error);
}
fn split(
&self,
execution: Execution,
parts: &[PhaseSpec<'_>],
) -> Result<alloc::vec::Vec<Child<'_>>, PlanError> {
self.child.split(execution, parts)
}
fn start(&self) -> Result<(), PlanError> {
self.child.start()
}
fn set_total(&self, total: crate::Total) -> Result<(), PlanError> {
self.child.set_total(total)
}
fn share(&self) -> Result<crate::SharedPulse, PlanError> {
self.child.share()
}
}