use crate::{
Child, ChildPulse, Execution, NodeId, Observer, Outcome, Phase, PhaseSpec, PlanError, Pulse,
Report, Reporter, SharedPulse, Stop, StopReason, Total, sync::OwnerCell,
};
use alloc::{boxed::Box, sync::Arc, vec::Vec};
use core::{
fmt,
sync::atomic::{AtomicBool, AtomicU8, Ordering},
};
#[must_use = "finish the root after its work joins; dropping it unfinished records Abandoned"]
pub struct PulseTree {
node: TreePulse<Arc<dyn Stop>>,
}
impl PulseTree {
pub fn new(phase: Phase, stop: impl Stop + 'static) -> Self {
Self {
node: TreePulse::new(phase, Arc::new(stop) as Arc<dyn Stop>),
}
}
pub fn observer(&self) -> Observer {
self.node.observer()
}
pub fn id(&self) -> NodeId {
self.node.state.observer.id()
}
pub fn finish(self, outcome: Outcome) -> Result<(), PlanError> {
self.node.finish(outcome)
}
}
impl fmt::Debug for PulseTree {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PulseTree")
.field("root", &self.id())
.finish()
}
}
const UNPLANNED: u8 = 0;
const COUNTING: u8 = 1;
const SPLIT: u8 = 2;
const FINISHED: u8 = 3;
pub(crate) trait Policy: Clone + Send + Sync + 'static {
#[track_caller]
fn check(&self, phase: &Observer) -> Result<(), StopReason>;
fn may_stop(&self) -> bool;
}
impl Policy for Arc<dyn Stop> {
#[track_caller]
fn check(&self, _phase: &Observer) -> Result<(), StopReason> {
Stop::check(self)
}
fn may_stop(&self) -> bool {
Stop::may_stop(self)
}
}
struct State<S: Policy> {
owner: OwnerCell<Phase>,
reporter: Reporter,
observer: Observer,
stop: S,
activity: AtomicU8,
closed: AtomicBool,
}
pub(crate) struct TreePulse<S: Policy> {
state: Arc<State<S>>,
}
#[derive(Clone)]
struct View<S: Policy> {
state: Arc<State<S>>,
}
impl<S: Policy> TreePulse<S> {
pub(crate) fn observer(&self) -> Observer {
self.state.observer.clone()
}
pub(crate) fn finish(self, outcome: Outcome) -> Result<(), PlanError> {
self.state.finish(outcome)
}
pub(crate) fn new(phase: Phase, stop: S) -> Self {
Self {
state: Arc::new(State {
reporter: phase.deferred_reporter(),
observer: phase.observer(),
owner: OwnerCell::new(phase),
stop,
activity: AtomicU8::new(UNPLANNED),
closed: AtomicBool::new(false),
}),
}
}
}
impl<S: Policy> State<S> {
fn complete(&self, outcome: Outcome) -> Result<(), PlanError> {
self.with_owner(|phase| phase.record_outcome(outcome))??;
self.activity.store(FINISHED, Ordering::Release);
Ok(())
}
fn start(&self) -> Result<(), PlanError> {
self.with_owner(Phase::start)?
}
fn set_total(&self, total: Total) -> Result<(), PlanError> {
self.with_owner(|p| p.set_total(total))?
}
fn share(self: &Arc<Self>) -> Result<SharedPulse, PlanError> {
Ok(SharedPulse::new(View {
state: Arc::clone(self),
}))
}
fn with_owner<R>(&self, f: impl FnOnce(&mut Phase) -> R) -> Result<R, PlanError> {
if self.closed.load(Ordering::Acquire) {
return Err(PlanError::Finished);
}
self.owner.with(f).ok_or(PlanError::Finished)
}
fn close(&self) {
self.closed.store(true, Ordering::Release);
self.activity.store(FINISHED, Ordering::Release);
drop(self.owner.take());
}
fn finish(&self, outcome: Outcome) -> Result<(), PlanError> {
let result = self.with_owner(|phase| phase.finish_with(outcome))?;
if result.is_ok() {
self.activity.store(FINISHED, Ordering::Release);
}
result
}
fn split(
&self,
execution: Execution,
parts: &[PhaseSpec<'_>],
) -> Result<Vec<Child<'static>>, PlanError> {
PhaseSpec::validate_split(parts)?;
self.activity
.compare_exchange(UNPLANNED, SPLIT, Ordering::AcqRel, Ordering::Acquire)
.map_err(|s| {
if s == FINISHED {
PlanError::Finished
} else {
PlanError::AlreadyInUse
}
})?;
let mut children = Vec::with_capacity(parts.len());
let split = self.with_owner(|phase| {
phase.split_each(execution, parts, &mut |phase| {
children.push(Child::new(TreePulse::new(phase, self.stop.clone())));
})
});
match split {
Ok(Ok(())) => Ok(children),
Ok(Err(error)) | Err(error) => {
let _ = self.activity.compare_exchange(
SPLIT,
UNPLANNED,
Ordering::AcqRel,
Ordering::Acquire,
);
Err(error)
}
}
}
}
impl<S: Policy> Stop for State<S> {
#[track_caller]
fn check(&self) -> Result<(), StopReason> {
self.stop.check(&self.observer)
}
fn may_stop(&self) -> bool {
self.stop.may_stop()
}
}
impl<S: Policy> Report for State<S> {
#[track_caller]
fn advance(&self, n: u64) {
if n == 0 {
return;
}
if self.activity.load(Ordering::Acquire) == COUNTING
|| matches!(
self.activity.compare_exchange(
UNPLANNED,
COUNTING,
Ordering::AcqRel,
Ordering::Acquire
),
Ok(_) | Err(COUNTING) | Err(SPLIT)
)
{
self.reporter.advance(n);
}
}
fn may_report(&self) -> bool {
self.activity.load(Ordering::Relaxed) != FINISHED && self.reporter.may_report()
}
}
impl<S: Policy> Drop for TreePulse<S> {
fn drop(&mut self) {
self.state.close();
}
}
impl<S: Policy> ChildPulse for TreePulse<S> {
fn complete_inferred(self: Box<Self>, outcome: Outcome) {
let _ = self
.state
.with_owner(|phase| phase.complete_inferred(outcome));
}
fn complete_as(self: Box<Self>, outcome: Outcome) {
let _ = self.state.complete(outcome);
}
fn finish(self: Box<Self>, outcome: Outcome) -> Result<(), PlanError> {
self.state.finish(outcome)
}
}
impl<S: Policy> how_far::Complete for TreePulse<S> {
fn complete_as(self, outcome: Outcome) {
let _ = self.state.complete(outcome);
}
}
impl how_far::Complete for PulseTree {
fn complete_as(self, outcome: Outcome) {
let _ = self.node.state.complete(outcome);
}
}
macro_rules! pulse_view {
([$($generic:ident : $bound:path)?] $ty:ty, $field:ident $(, #[$hint:meta])?) => {
impl$(<$generic: $bound>)? Stop for $ty {
$(#[$hint])?
#[track_caller]
fn check(&self) -> Result<(), StopReason> {
self.$field.check()
}
fn may_stop(&self) -> bool {
self.$field.may_stop()
}
}
impl$(<$generic: $bound>)? Report for $ty {
$(#[$hint])?
#[track_caller]
fn advance(&self, n: u64) {
self.$field.advance(n)
}
fn may_report(&self) -> bool {
self.$field.may_report()
}
}
impl$(<$generic: $bound>)? Pulse for $ty {
fn split(
&self,
e: Execution,
parts: &[PhaseSpec<'_>],
) -> Result<Vec<Child<'_>>, PlanError> {
self.$field.split(e, parts)
}
fn start(&self) -> Result<(), PlanError> {
self.$field.start()
}
fn set_total(&self, total: Total) -> Result<(), PlanError> {
self.$field.set_total(total)
}
fn share(&self) -> Result<SharedPulse, PlanError> {
self.$field.share()
}
}
};
}
pulse_view!([S: Policy] TreePulse<S>, state, #[inline]);
pulse_view!([S: Policy] View<S>, state, #[inline]);
pulse_view!([] PulseTree, node);