gammalooprs/processes/
generation_progress.rs

1use std::{
2    cell::RefCell,
3    sync::{
4        Arc, Mutex, OnceLock,
5        atomic::{AtomicU8, Ordering},
6    },
7    time::Duration,
8};
9
10use super::GraphGenerationStats;
11
12thread_local! {
13    static GENERATION_PROGRESS_CONTEXT: RefCell<Vec<String>> = const { RefCell::new(Vec::new()) };
14}
15
16pub struct GenerationProgressContextGuard;
17
18impl Drop for GenerationProgressContextGuard {
19    fn drop(&mut self) {
20        GENERATION_PROGRESS_CONTEXT.with(|context| {
21            context.borrow_mut().pop();
22        });
23    }
24}
25
26pub fn enter_progress_context(context: impl Into<String>) -> GenerationProgressContextGuard {
27    GENERATION_PROGRESS_CONTEXT.with(|stack| stack.borrow_mut().push(context.into()));
28    GenerationProgressContextGuard
29}
30
31pub fn detailed_progress_message(message: &str) -> String {
32    GENERATION_PROGRESS_CONTEXT.with(|context| {
33        let context = context.borrow();
34        if context.is_empty() {
35            message.to_string()
36        } else {
37            format!("{} / {message}", context.join(" / "))
38        }
39    })
40}
41
42pub fn enter_detailed_progress_span(message: &str) -> Option<tracing::span::EnteredSpan> {
43    if !detailed_progress_enabled() {
44        return None;
45    }
46
47    let progress_message = detailed_progress_message(message);
48    Some(
49        tracing::info_span!(
50            "Generation progress",
51            indicatif.pb_show = true,
52            indicatif.pb_msg = %progress_message,
53        )
54        .entered(),
55    )
56}
57
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum GenerationProgressMode {
60    Detailed,
61    Aggregate,
62}
63
64impl GenerationProgressMode {
65    fn as_u8(self) -> u8 {
66        match self {
67            Self::Detailed => 0,
68            Self::Aggregate => 1,
69        }
70    }
71
72    fn from_u8(value: u8) -> Self {
73        match value {
74            1 => Self::Aggregate,
75            _ => Self::Detailed,
76        }
77    }
78}
79
80static GENERATION_PROGRESS_MODE: AtomicU8 = AtomicU8::new(0);
81
82pub struct GenerationProgressModeGuard {
83    previous: u8,
84}
85
86impl GenerationProgressModeGuard {
87    pub fn set(mode: GenerationProgressMode) -> Self {
88        let previous = GENERATION_PROGRESS_MODE.swap(mode.as_u8(), Ordering::Relaxed);
89        Self { previous }
90    }
91}
92
93impl Drop for GenerationProgressModeGuard {
94    fn drop(&mut self) {
95        GENERATION_PROGRESS_MODE.store(self.previous, Ordering::Relaxed);
96    }
97}
98
99pub fn current_generation_progress_mode() -> GenerationProgressMode {
100    GenerationProgressMode::from_u8(GENERATION_PROGRESS_MODE.load(Ordering::Relaxed))
101}
102
103pub fn detailed_progress_enabled() -> bool {
104    current_generation_progress_mode() == GenerationProgressMode::Detailed
105}
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub enum GenerationProcessKind {
109    Amplitude,
110    CrossSection,
111}
112
113#[derive(Debug, Clone, Copy, PartialEq, Eq)]
114pub enum GenerationProgressPhase {
115    GraphPreprocessing,
116    GraphGeneration,
117    Backend,
118}
119
120pub trait GenerationProgressObserver: Send + Sync {
121    fn begin_phase(
122        &self,
123        _phase: GenerationProgressPhase,
124        _kind: GenerationProcessKind,
125        _process: &str,
126        _integrand: &str,
127        _total_graphs: usize,
128        _total_cuts: Option<usize>,
129    ) {
130    }
131
132    fn graph_started(
133        &self,
134        _kind: GenerationProcessKind,
135        _integrand: &str,
136        _graph: &str,
137        _cut_count: Option<usize>,
138    ) {
139    }
140
141    fn graph_finished(
142        &self,
143        _kind: GenerationProcessKind,
144        _integrand: &str,
145        _graph: &str,
146        _stats: &GraphGenerationStats,
147        _completed_cuts: Option<usize>,
148    ) {
149    }
150
151    fn cuts_discovered(
152        &self,
153        _integrand: &str,
154        _graph: &str,
155        _st_cut_count: usize,
156        _valid_cut_count: usize,
157    ) {
158    }
159
160    fn cut_finished(&self, _integrand: &str, _graph: &str, _cut_count: usize) {}
161
162    fn backend_started(&self, _kind: GenerationProcessKind, _integrand: &str, _graph_count: usize) {
163    }
164
165    fn backend_finished(&self, _kind: GenerationProcessKind, _integrand: &str, _elapsed: Duration) {
166    }
167}
168
169static GENERATION_PROGRESS_OBSERVER: OnceLock<Mutex<Option<Arc<dyn GenerationProgressObserver>>>> =
170    OnceLock::new();
171
172fn observer_slot() -> &'static Mutex<Option<Arc<dyn GenerationProgressObserver>>> {
173    GENERATION_PROGRESS_OBSERVER.get_or_init(|| Mutex::new(None))
174}
175
176fn observer() -> Option<Arc<dyn GenerationProgressObserver>> {
177    observer_slot().lock().ok().and_then(|guard| guard.clone())
178}
179
180pub struct GenerationProgressObserverGuard {
181    previous: Option<Arc<dyn GenerationProgressObserver>>,
182}
183
184impl GenerationProgressObserverGuard {
185    pub fn set(observer: Arc<dyn GenerationProgressObserver>) -> Self {
186        let previous = observer_slot()
187            .lock()
188            .expect("generation progress observer mutex is poisoned")
189            .replace(observer);
190        Self { previous }
191    }
192}
193
194impl Drop for GenerationProgressObserverGuard {
195    fn drop(&mut self) {
196        *observer_slot()
197            .lock()
198            .expect("generation progress observer mutex is poisoned") = self.previous.take();
199    }
200}
201
202pub fn begin_phase(
203    phase: GenerationProgressPhase,
204    kind: GenerationProcessKind,
205    process: &str,
206    integrand: &str,
207    total_graphs: usize,
208    total_cuts: Option<usize>,
209) {
210    if let Some(observer) = observer() {
211        observer.begin_phase(phase, kind, process, integrand, total_graphs, total_cuts);
212    }
213}
214
215pub fn graph_started(
216    kind: GenerationProcessKind,
217    integrand: &str,
218    graph: &str,
219    cut_count: Option<usize>,
220) {
221    if let Some(observer) = observer() {
222        observer.graph_started(kind, integrand, graph, cut_count);
223    }
224}
225
226pub fn graph_finished(
227    kind: GenerationProcessKind,
228    integrand: &str,
229    graph: &str,
230    stats: &GraphGenerationStats,
231    completed_cuts: Option<usize>,
232) {
233    if let Some(observer) = observer() {
234        observer.graph_finished(kind, integrand, graph, stats, completed_cuts);
235    }
236}
237
238pub fn cuts_discovered(integrand: &str, graph: &str, st_cut_count: usize, valid_cut_count: usize) {
239    if let Some(observer) = observer() {
240        observer.cuts_discovered(integrand, graph, st_cut_count, valid_cut_count);
241    }
242}
243
244pub fn cut_finished(integrand: &str, graph: &str, cut_count: usize) {
245    if let Some(observer) = observer() {
246        observer.cut_finished(integrand, graph, cut_count);
247    }
248}
249
250pub fn backend_started(kind: GenerationProcessKind, integrand: &str, graph_count: usize) {
251    if let Some(observer) = observer() {
252        observer.backend_started(kind, integrand, graph_count);
253    }
254}
255
256pub fn backend_finished(kind: GenerationProcessKind, integrand: &str, elapsed: Duration) {
257    if let Some(observer) = observer() {
258        observer.backend_finished(kind, integrand, elapsed);
259    }
260}
gammalooprs/processes/generation_progress.rs

gammalooprs/processes/
generation_progress.rs
