1use std::collections::BTreeSet;
4use std::error::Error;
5use std::fmt::{Display, Formatter};
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::{Arc, Mutex};
8use std::time::Instant;
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub enum TerminationReason {
13 Converged,
15 MaxIterations,
17 NoObservations,
19 NumericalFailure,
21 Cancelled,
23 MaxRuntime,
25 MaxEvaluations,
27 Stagnated,
29 Diverged,
31 RepeatedRejections,
33}
34
35impl TerminationReason {
36 #[must_use]
38 pub const fn as_str(self) -> &'static str {
39 match self {
40 Self::Converged => "converged",
41 Self::MaxIterations => "max_iterations",
42 Self::NoObservations => "no_observations",
43 Self::NumericalFailure => "numerical_failure",
44 Self::Cancelled => "cancelled",
45 Self::MaxRuntime => "max_runtime",
46 Self::MaxEvaluations => "max_evaluations",
47 Self::Stagnated => "stagnated",
48 Self::Diverged => "diverged",
49 Self::RepeatedRejections => "repeated_rejections",
50 }
51 }
52}
53
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56pub enum RefinementEventKind {
57 Start,
59 Trial,
61 StepAccepted,
63 StepRejected,
65 Iteration,
67 Checkpoint,
69 Warning,
71 Termination,
73 Failure,
75}
76
77impl RefinementEventKind {
78 #[must_use]
80 pub const fn as_str(self) -> &'static str {
81 match self {
82 Self::Start => "start",
83 Self::Trial => "trial",
84 Self::StepAccepted => "step_accepted",
85 Self::StepRejected => "step_rejected",
86 Self::Iteration => "iteration",
87 Self::Checkpoint => "checkpoint",
88 Self::Warning => "warning",
89 Self::Termination => "termination",
90 Self::Failure => "failure",
91 }
92 }
93}
94
95#[derive(Clone, Debug, PartialEq)]
97pub enum DiagnosticValue {
98 String(String),
100 Bool(bool),
102 Integer(i64),
104 Unsigned(u64),
106 Float(f64),
108 Null,
110}
111
112impl DiagnosticValue {
113 fn validate(&self) -> Result<(), RuntimeError> {
114 if matches!(self, Self::Float(value) if !value.is_finite()) {
115 return Err(RuntimeError::InvalidEvent {
116 message: "floating-point diagnostics must be finite".to_owned(),
117 });
118 }
119 Ok(())
120 }
121}
122
123#[derive(Clone, Debug, PartialEq)]
125pub struct RefinementEvent {
126 kind: RefinementEventKind,
127 stage: String,
128 attempted_iteration: usize,
129 accepted_iterations: usize,
130 evaluations: usize,
131 elapsed_seconds: f64,
132 message: String,
133 diagnostics: Vec<(String, DiagnosticValue)>,
134}
135
136impl RefinementEvent {
137 #[allow(clippy::too_many_arguments)]
144 pub fn new(
145 kind: RefinementEventKind,
146 stage: impl Into<String>,
147 attempted_iteration: usize,
148 accepted_iterations: usize,
149 evaluations: usize,
150 elapsed_seconds: f64,
151 message: impl Into<String>,
152 diagnostics: Vec<(String, DiagnosticValue)>,
153 ) -> Result<Self, RuntimeError> {
154 let stage = stage.into();
155 let message = message.into();
156 if stage.trim().is_empty() {
157 return Err(invalid_event("event stage must be non-empty"));
158 }
159 if accepted_iterations > attempted_iteration {
160 return Err(invalid_event(
161 "accepted iterations cannot exceed attempted iterations",
162 ));
163 }
164 if !elapsed_seconds.is_finite() || elapsed_seconds < 0.0 {
165 return Err(invalid_event(
166 "event elapsed time must be non-negative and finite",
167 ));
168 }
169 if message.is_empty() {
170 return Err(invalid_event("event message must be non-empty"));
171 }
172 let mut keys = BTreeSet::new();
173 for (key, value) in &diagnostics {
174 if key.is_empty() {
175 return Err(invalid_event("diagnostic keys must be non-empty"));
176 }
177 if !keys.insert(key.clone()) {
178 return Err(RuntimeError::InvalidEvent {
179 message: format!("duplicate diagnostic key {key:?}"),
180 });
181 }
182 value.validate()?;
183 }
184 Ok(Self {
185 kind,
186 stage,
187 attempted_iteration,
188 accepted_iterations,
189 evaluations,
190 elapsed_seconds,
191 message,
192 diagnostics,
193 })
194 }
195
196 #[must_use]
198 pub const fn kind(&self) -> RefinementEventKind {
199 self.kind
200 }
201
202 #[must_use]
204 pub fn stage(&self) -> &str {
205 &self.stage
206 }
207
208 #[must_use]
210 pub const fn attempted_iteration(&self) -> usize {
211 self.attempted_iteration
212 }
213
214 #[must_use]
216 pub const fn accepted_iterations(&self) -> usize {
217 self.accepted_iterations
218 }
219
220 #[must_use]
222 pub const fn evaluations(&self) -> usize {
223 self.evaluations
224 }
225
226 #[must_use]
228 pub const fn elapsed_seconds(&self) -> f64 {
229 self.elapsed_seconds
230 }
231
232 #[must_use]
234 pub fn message(&self) -> &str {
235 &self.message
236 }
237
238 #[must_use]
240 pub fn diagnostics(&self) -> &[(String, DiagnosticValue)] {
241 &self.diagnostics
242 }
243}
244
245#[derive(Debug, Default)]
246struct CancellationState {
247 requested: AtomicBool,
248 reason: Mutex<Option<String>>,
249}
250
251#[derive(Clone, Debug, Default)]
253pub struct CancellationToken {
254 state: Arc<CancellationState>,
255}
256
257impl CancellationToken {
258 pub fn request(&self, reason: impl Into<String>) -> Result<bool, CancellationError> {
267 let reason = reason.into();
268 if reason.trim().is_empty() {
269 return Err(CancellationError::InvalidReason);
270 }
271 let mut stored = self
272 .state
273 .reason
274 .lock()
275 .map_err(|_| CancellationError::Poisoned)?;
276 if stored.is_some() {
277 return Ok(false);
278 }
279 *stored = Some(reason);
280 self.state.requested.store(true, Ordering::Release);
281 Ok(true)
282 }
283
284 #[must_use]
286 pub fn is_requested(&self) -> bool {
287 self.state.requested.load(Ordering::Acquire)
288 }
289
290 pub fn reason(&self) -> Result<Option<String>, CancellationError> {
296 self.state
297 .reason
298 .lock()
299 .map(|reason| reason.clone())
300 .map_err(|_| CancellationError::Poisoned)
301 }
302}
303
304#[derive(Clone, Copy, Debug, PartialEq, Eq)]
306pub enum CancellationError {
307 InvalidReason,
309 Poisoned,
311}
312
313impl Display for CancellationError {
314 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
315 match self {
316 Self::InvalidReason => formatter.write_str("cancellation reason must be non-empty"),
317 Self::Poisoned => formatter.write_str("cancellation reason lock is poisoned"),
318 }
319 }
320}
321
322impl Error for CancellationError {}
323
324#[derive(Clone, Copy, Debug, PartialEq)]
326#[allow(clippy::struct_field_names)]
328pub struct RefinementLimits {
329 max_iterations: usize,
330 max_evaluations: usize,
331 max_runtime_seconds: Option<f64>,
332 max_consecutive_rejections: usize,
333}
334
335impl RefinementLimits {
336 pub fn new(
342 max_iterations: usize,
343 max_evaluations: usize,
344 max_runtime_seconds: Option<f64>,
345 max_consecutive_rejections: usize,
346 ) -> Result<Self, RuntimeError> {
347 if max_iterations == 0 || max_evaluations == 0 || max_consecutive_rejections == 0 {
348 return Err(RuntimeError::InvalidLimits);
349 }
350 if max_runtime_seconds.is_some_and(|seconds| !seconds.is_finite() || seconds <= 0.0) {
351 return Err(RuntimeError::InvalidLimits);
352 }
353 Ok(Self {
354 max_iterations,
355 max_evaluations,
356 max_runtime_seconds,
357 max_consecutive_rejections,
358 })
359 }
360
361 #[must_use]
363 pub const fn max_iterations(self) -> usize {
364 self.max_iterations
365 }
366
367 #[must_use]
369 pub const fn max_evaluations(self) -> usize {
370 self.max_evaluations
371 }
372
373 #[must_use]
375 pub const fn max_runtime_seconds(self) -> Option<f64> {
376 self.max_runtime_seconds
377 }
378
379 #[must_use]
381 pub const fn max_consecutive_rejections(self) -> usize {
382 self.max_consecutive_rejections
383 }
384}
385
386impl Default for RefinementLimits {
387 fn default() -> Self {
388 Self {
389 max_iterations: 100,
390 max_evaluations: 1_000,
391 max_runtime_seconds: None,
392 max_consecutive_rejections: 20,
393 }
394 }
395}
396
397pub trait RuntimeClock: Send + Sync {
399 fn now_seconds(&self) -> f64;
401}
402
403#[derive(Debug)]
405pub struct MonotonicClock {
406 origin: Instant,
407}
408
409impl MonotonicClock {
410 #[must_use]
412 pub fn new() -> Self {
413 Self {
414 origin: Instant::now(),
415 }
416 }
417}
418
419impl Default for MonotonicClock {
420 fn default() -> Self {
421 Self::new()
422 }
423}
424
425impl RuntimeClock for MonotonicClock {
426 fn now_seconds(&self) -> f64 {
427 self.origin.elapsed().as_secs_f64()
428 }
429}
430
431pub trait RefinementEventSink: Send {
433 fn emit(&mut self, event: &RefinementEvent) -> Result<(), String>;
440}
441
442impl<F> RefinementEventSink for F
443where
444 F: FnMut(&RefinementEvent) -> Result<(), String> + Send,
445{
446 fn emit(&mut self, event: &RefinementEvent) -> Result<(), String> {
447 self(event)
448 }
449}
450
451pub trait CheckpointSink<C>: Send {
453 fn checkpoint(&mut self, checkpoint: &C) -> Result<(), String>;
459}
460
461impl<C, F> CheckpointSink<C> for F
462where
463 F: FnMut(&C) -> Result<(), String> + Send,
464{
465 fn checkpoint(&mut self, checkpoint: &C) -> Result<(), String> {
466 self(checkpoint)
467 }
468}
469
470#[derive(Clone, Debug, PartialEq, Eq)]
472pub struct RefinementStop {
473 pub reason: TerminationReason,
475 pub message: String,
477}
478
479impl Display for RefinementStop {
480 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
481 formatter.write_str(&self.message)
482 }
483}
484
485impl Error for RefinementStop {}
486
487pub struct RefinementRuntime<C = ()> {
489 limits: RefinementLimits,
490 cancellation: Option<CancellationToken>,
491 event_sink: Option<Box<dyn RefinementEventSink>>,
492 checkpoint_sink: Option<Box<dyn CheckpointSink<C>>>,
493 clock: Arc<dyn RuntimeClock>,
494 started_at: f64,
495 attempted_iteration: usize,
496 accepted_iterations: usize,
497 evaluations: usize,
498 consecutive_rejections: usize,
499 event_sink_error: Option<String>,
500}
501
502impl<C> RefinementRuntime<C> {
503 pub fn new(
509 limits: RefinementLimits,
510 cancellation: Option<CancellationToken>,
511 ) -> Result<Self, RuntimeError> {
512 Self::with_clock(limits, cancellation, Arc::new(MonotonicClock::new()))
513 }
514
515 pub fn with_clock(
521 limits: RefinementLimits,
522 cancellation: Option<CancellationToken>,
523 clock: Arc<dyn RuntimeClock>,
524 ) -> Result<Self, RuntimeError> {
525 let started_at = clock.now_seconds();
526 if !started_at.is_finite() {
527 return Err(RuntimeError::InvalidClock);
528 }
529 Ok(Self {
530 limits,
531 cancellation,
532 event_sink: None,
533 checkpoint_sink: None,
534 clock,
535 started_at,
536 attempted_iteration: 0,
537 accepted_iterations: 0,
538 evaluations: 0,
539 consecutive_rejections: 0,
540 event_sink_error: None,
541 })
542 }
543
544 pub fn set_event_sink(&mut self, sink: impl RefinementEventSink + 'static) {
546 self.event_sink = Some(Box::new(sink));
547 self.event_sink_error = None;
548 }
549
550 pub fn set_checkpoint_sink(&mut self, sink: impl CheckpointSink<C> + 'static) {
552 self.checkpoint_sink = Some(Box::new(sink));
553 }
554
555 pub fn resume_accepted(&mut self, completed_iterations: usize) -> Result<(), RuntimeError> {
564 if self.attempted_iteration != 0
565 || self.accepted_iterations != 0
566 || self.evaluations != 0
567 || completed_iterations > self.limits.max_iterations
568 {
569 return Err(RuntimeError::InvalidResume);
570 }
571 self.attempted_iteration = completed_iterations;
572 self.accepted_iterations = completed_iterations;
573 Ok(())
574 }
575
576 pub fn elapsed_seconds(&self) -> Result<f64, RuntimeError> {
582 let elapsed = self.clock.now_seconds() - self.started_at;
583 if !elapsed.is_finite() || elapsed < 0.0 {
584 return Err(RuntimeError::InvalidClock);
585 }
586 Ok(elapsed)
587 }
588
589 pub fn emit(
596 &mut self,
597 kind: RefinementEventKind,
598 stage: impl Into<String>,
599 message: impl Into<String>,
600 diagnostics: Vec<(String, DiagnosticValue)>,
601 ) -> Result<RefinementEvent, RuntimeError> {
602 let event = RefinementEvent::new(
603 kind,
604 stage,
605 self.attempted_iteration,
606 self.accepted_iterations,
607 self.evaluations,
608 self.elapsed_seconds()?,
609 message,
610 diagnostics,
611 )?;
612 let sink_result = self.event_sink.as_mut().map(|sink| sink.emit(&event));
613 if let Some(Err(message)) = sink_result {
614 self.event_sink_error = Some(message);
615 self.event_sink = None;
616 }
617 Ok(event)
618 }
619
620 pub fn check_boundary(&self) -> Result<(), RuntimeError> {
627 if let Some(token) = &self.cancellation
628 && token.is_requested()
629 {
630 let message = token
631 .reason()
632 .map_err(RuntimeError::Cancellation)?
633 .unwrap_or_else(|| "user requested cancellation".to_owned());
634 return Err(RuntimeError::Stopped(RefinementStop {
635 reason: TerminationReason::Cancelled,
636 message,
637 }));
638 }
639 if let Some(limit) = self.limits.max_runtime_seconds {
640 let elapsed = self.elapsed_seconds()?;
641 if elapsed >= limit {
642 return Err(RuntimeError::Stopped(RefinementStop {
643 reason: TerminationReason::MaxRuntime,
644 message: "runtime limit reached".to_owned(),
645 }));
646 }
647 }
648 if self.evaluations >= self.limits.max_evaluations {
649 return Err(RuntimeError::Stopped(RefinementStop {
650 reason: TerminationReason::MaxEvaluations,
651 message: "model-evaluation limit reached".to_owned(),
652 }));
653 }
654 Ok(())
655 }
656
657 pub fn begin_iteration(&mut self, attempted_iteration: usize) -> Result<(), RuntimeError> {
663 if attempted_iteration <= self.attempted_iteration {
664 return Err(RuntimeError::InvalidIterationOrder);
665 }
666 if attempted_iteration > self.limits.max_iterations {
667 return Err(RuntimeError::Stopped(RefinementStop {
668 reason: TerminationReason::MaxIterations,
669 message: "iteration limit reached".to_owned(),
670 }));
671 }
672 self.attempted_iteration = attempted_iteration;
673 self.check_boundary()
674 }
675
676 pub fn begin_evaluation(&mut self) -> Result<(), RuntimeError> {
682 self.check_boundary()?;
683 self.evaluations = self
684 .evaluations
685 .checked_add(1)
686 .ok_or(RuntimeError::CounterOverflow)?;
687 Ok(())
688 }
689
690 pub fn accept_step(&mut self, checkpoint: Option<&C>) -> Result<(), RuntimeError> {
700 if self.accepted_iterations >= self.attempted_iteration {
701 return Err(RuntimeError::DuplicateAcceptance);
702 }
703 self.accepted_iterations = self
704 .accepted_iterations
705 .checked_add(1)
706 .ok_or(RuntimeError::CounterOverflow)?;
707 self.consecutive_rejections = 0;
708 if let Some(checkpoint) = checkpoint
709 && let Some(sink) = self.checkpoint_sink.as_mut()
710 {
711 sink.checkpoint(checkpoint)
712 .map_err(|message| RuntimeError::CheckpointSink { message })?;
713 self.emit(
714 RefinementEventKind::Checkpoint,
715 "checkpoint",
716 "accepted-state checkpoint completed",
717 Vec::new(),
718 )?;
719 }
720 Ok(())
721 }
722
723 pub fn reject_step(&mut self) -> Result<(), RuntimeError> {
729 self.consecutive_rejections = self
730 .consecutive_rejections
731 .checked_add(1)
732 .ok_or(RuntimeError::CounterOverflow)?;
733 if self.consecutive_rejections >= self.limits.max_consecutive_rejections {
734 return Err(RuntimeError::Stopped(RefinementStop {
735 reason: TerminationReason::RepeatedRejections,
736 message: "consecutive rejected-step limit reached".to_owned(),
737 }));
738 }
739 Ok(())
740 }
741
742 #[must_use]
744 pub const fn attempted_iteration(&self) -> usize {
745 self.attempted_iteration
746 }
747
748 #[must_use]
750 pub const fn accepted_iterations(&self) -> usize {
751 self.accepted_iterations
752 }
753
754 #[must_use]
756 pub const fn evaluations(&self) -> usize {
757 self.evaluations
758 }
759
760 #[must_use]
762 pub const fn consecutive_rejections(&self) -> usize {
763 self.consecutive_rejections
764 }
765
766 #[must_use]
768 pub fn event_sink_error(&self) -> Option<&str> {
769 self.event_sink_error.as_deref()
770 }
771
772 #[must_use]
774 pub const fn has_event_sink(&self) -> bool {
775 self.event_sink.is_some()
776 }
777}
778
779#[derive(Debug)]
781pub enum RuntimeError {
782 InvalidLimits,
784 InvalidClock,
786 InvalidEvent {
788 message: String,
790 },
791 InvalidResume,
793 InvalidIterationOrder,
795 DuplicateAcceptance,
797 CounterOverflow,
799 Cancellation(CancellationError),
801 Stopped(RefinementStop),
803 CheckpointSink {
805 message: String,
807 },
808}
809
810impl Display for RuntimeError {
811 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
812 match self {
813 Self::InvalidLimits => {
814 formatter.write_str("refinement limits must be positive and finite")
815 }
816 Self::InvalidClock => {
817 formatter.write_str("refinement clock must be finite and monotonic")
818 }
819 Self::InvalidEvent { message } | Self::CheckpointSink { message } => {
820 formatter.write_str(message)
821 }
822 Self::InvalidResume => {
823 formatter.write_str("refinement counters cannot be resumed in this state")
824 }
825 Self::InvalidIterationOrder => {
826 formatter.write_str("attempted iterations must increase strictly")
827 }
828 Self::DuplicateAcceptance => {
829 formatter.write_str("at most one step may be accepted per attempted iteration")
830 }
831 Self::CounterOverflow => formatter.write_str("refinement runtime counter overflow"),
832 Self::Cancellation(error) => Display::fmt(error, formatter),
833 Self::Stopped(stop) => Display::fmt(stop, formatter),
834 }
835 }
836}
837
838impl Error for RuntimeError {
839 fn source(&self) -> Option<&(dyn Error + 'static)> {
840 match self {
841 Self::Cancellation(error) => Some(error),
842 Self::Stopped(stop) => Some(stop),
843 _ => None,
844 }
845 }
846}
847
848fn invalid_event(message: &str) -> RuntimeError {
849 RuntimeError::InvalidEvent {
850 message: message.to_owned(),
851 }
852}