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,
37}
38
39impl TerminationReason {
40 #[must_use]
42 pub const fn as_str(self) -> &'static str {
43 match self {
44 Self::Converged => "converged",
45 Self::MaxIterations => "max_iterations",
46 Self::NoObservations => "no_observations",
47 Self::NumericalFailure => "numerical_failure",
48 Self::Cancelled => "cancelled",
49 Self::MaxRuntime => "max_runtime",
50 Self::MaxEvaluations => "max_evaluations",
51 Self::Stagnated => "stagnated",
52 Self::Diverged => "diverged",
53 Self::RepeatedRejections => "repeated_rejections",
54 }
55 }
56}
57
58#[derive(Clone, Copy, Debug, PartialEq, Eq)]
60pub enum RefinementEventKind {
61 Start,
63 Trial,
65 StepAccepted,
67 StepRejected,
69 Iteration,
71 Checkpoint,
73 Warning,
75 Termination,
77 Failure,
79}
80
81impl RefinementEventKind {
82 #[must_use]
84 pub const fn as_str(self) -> &'static str {
85 match self {
86 Self::Start => "start",
87 Self::Trial => "trial",
88 Self::StepAccepted => "step_accepted",
89 Self::StepRejected => "step_rejected",
90 Self::Iteration => "iteration",
91 Self::Checkpoint => "checkpoint",
92 Self::Warning => "warning",
93 Self::Termination => "termination",
94 Self::Failure => "failure",
95 }
96 }
97}
98
99#[derive(Clone, Debug, PartialEq)]
101pub enum DiagnosticValue {
102 String(String),
104 Bool(bool),
106 Integer(i64),
108 Unsigned(u64),
110 Float(f64),
112 Null,
114}
115
116impl DiagnosticValue {
117 fn validate(&self) -> Result<(), RuntimeError> {
118 if matches!(self, Self::Float(value) if !value.is_finite()) {
119 return Err(RuntimeError::InvalidEvent {
120 message: "floating-point diagnostics must be finite".to_owned(),
121 });
122 }
123 Ok(())
124 }
125}
126
127#[derive(Clone, Debug, PartialEq)]
129pub struct RefinementEvent {
130 kind: RefinementEventKind,
131 stage: String,
132 attempted_iteration: usize,
133 accepted_iterations: usize,
134 evaluations: usize,
135 elapsed_seconds: f64,
136 message: String,
137 diagnostics: Vec<(String, DiagnosticValue)>,
138}
139
140impl RefinementEvent {
141 #[allow(clippy::too_many_arguments)]
148 pub fn new(
149 kind: RefinementEventKind,
150 stage: impl Into<String>,
151 attempted_iteration: usize,
152 accepted_iterations: usize,
153 evaluations: usize,
154 elapsed_seconds: f64,
155 message: impl Into<String>,
156 diagnostics: Vec<(String, DiagnosticValue)>,
157 ) -> Result<Self, RuntimeError> {
158 let stage = stage.into();
159 let message = message.into();
160 if stage.trim().is_empty() {
161 return Err(invalid_event("event stage must be non-empty"));
162 }
163 if accepted_iterations > attempted_iteration {
164 return Err(invalid_event(
165 "accepted iterations cannot exceed attempted iterations",
166 ));
167 }
168 if !elapsed_seconds.is_finite() || elapsed_seconds < 0.0 {
169 return Err(invalid_event(
170 "event elapsed time must be non-negative and finite",
171 ));
172 }
173 if message.is_empty() {
174 return Err(invalid_event("event message must be non-empty"));
175 }
176 let mut keys = BTreeSet::new();
177 for (key, value) in &diagnostics {
178 if key.is_empty() {
179 return Err(invalid_event("diagnostic keys must be non-empty"));
180 }
181 if !keys.insert(key.clone()) {
182 return Err(RuntimeError::InvalidEvent {
183 message: format!("duplicate diagnostic key {key:?}"),
184 });
185 }
186 value.validate()?;
187 }
188 Ok(Self {
189 kind,
190 stage,
191 attempted_iteration,
192 accepted_iterations,
193 evaluations,
194 elapsed_seconds,
195 message,
196 diagnostics,
197 })
198 }
199
200 #[must_use]
202 pub const fn kind(&self) -> RefinementEventKind {
203 self.kind
204 }
205
206 #[must_use]
208 pub fn stage(&self) -> &str {
209 &self.stage
210 }
211
212 #[must_use]
214 pub const fn attempted_iteration(&self) -> usize {
215 self.attempted_iteration
216 }
217
218 #[must_use]
220 pub const fn accepted_iterations(&self) -> usize {
221 self.accepted_iterations
222 }
223
224 #[must_use]
226 pub const fn evaluations(&self) -> usize {
227 self.evaluations
228 }
229
230 #[must_use]
232 pub const fn elapsed_seconds(&self) -> f64 {
233 self.elapsed_seconds
234 }
235
236 #[must_use]
238 pub fn message(&self) -> &str {
239 &self.message
240 }
241
242 #[must_use]
244 pub fn diagnostics(&self) -> &[(String, DiagnosticValue)] {
245 &self.diagnostics
246 }
247}
248
249#[derive(Debug, Default)]
250struct CancellationState {
251 requested: AtomicBool,
252 reason: Mutex<Option<String>>,
253}
254
255#[derive(Clone, Debug, Default)]
257pub struct CancellationToken {
258 state: Arc<CancellationState>,
259}
260
261impl CancellationToken {
262 pub fn request(&self, reason: impl Into<String>) -> Result<bool, CancellationError> {
271 let reason = reason.into();
272 if reason.trim().is_empty() {
273 return Err(CancellationError::InvalidReason);
274 }
275 let mut stored = self
276 .state
277 .reason
278 .lock()
279 .map_err(|_| CancellationError::Poisoned)?;
280 if stored.is_some() {
281 return Ok(false);
282 }
283 *stored = Some(reason);
284 self.state.requested.store(true, Ordering::Release);
285 Ok(true)
286 }
287
288 #[must_use]
290 pub fn is_requested(&self) -> bool {
291 self.state.requested.load(Ordering::Acquire)
292 }
293
294 pub fn reason(&self) -> Result<Option<String>, CancellationError> {
300 self.state
301 .reason
302 .lock()
303 .map(|reason| reason.clone())
304 .map_err(|_| CancellationError::Poisoned)
305 }
306}
307
308#[derive(Clone, Copy, Debug, PartialEq, Eq)]
310pub enum CancellationError {
311 InvalidReason,
313 Poisoned,
315}
316
317impl Display for CancellationError {
318 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
319 match self {
320 Self::InvalidReason => formatter.write_str("cancellation reason must be non-empty"),
321 Self::Poisoned => formatter.write_str("cancellation reason lock is poisoned"),
322 }
323 }
324}
325
326impl Error for CancellationError {}
327
328#[derive(Clone, Copy, Debug, PartialEq)]
330#[allow(clippy::struct_field_names)]
332pub struct RefinementLimits {
333 max_iterations: usize,
334 max_evaluations: usize,
335 max_runtime_seconds: Option<f64>,
336 max_consecutive_rejections: usize,
337}
338
339impl RefinementLimits {
340 pub fn new(
346 max_iterations: usize,
347 max_evaluations: usize,
348 max_runtime_seconds: Option<f64>,
349 max_consecutive_rejections: usize,
350 ) -> Result<Self, RuntimeError> {
351 if max_iterations == 0 || max_evaluations == 0 || max_consecutive_rejections == 0 {
352 return Err(RuntimeError::InvalidLimits);
353 }
354 if max_runtime_seconds.is_some_and(|seconds| !seconds.is_finite() || seconds <= 0.0) {
355 return Err(RuntimeError::InvalidLimits);
356 }
357 Ok(Self {
358 max_iterations,
359 max_evaluations,
360 max_runtime_seconds,
361 max_consecutive_rejections,
362 })
363 }
364
365 #[must_use]
367 pub const fn max_iterations(self) -> usize {
368 self.max_iterations
369 }
370
371 #[must_use]
373 pub const fn max_evaluations(self) -> usize {
374 self.max_evaluations
375 }
376
377 #[must_use]
379 pub const fn max_runtime_seconds(self) -> Option<f64> {
380 self.max_runtime_seconds
381 }
382
383 #[must_use]
385 pub const fn max_consecutive_rejections(self) -> usize {
386 self.max_consecutive_rejections
387 }
388}
389
390impl Default for RefinementLimits {
391 fn default() -> Self {
392 Self {
393 max_iterations: 100,
394 max_evaluations: 1_000,
395 max_runtime_seconds: None,
396 max_consecutive_rejections: 20,
397 }
398 }
399}
400
401pub trait RuntimeClock: Send + Sync {
403 fn now_seconds(&self) -> f64;
405}
406
407#[derive(Debug)]
409pub struct MonotonicClock {
410 origin: Instant,
411}
412
413impl MonotonicClock {
414 #[must_use]
416 pub fn new() -> Self {
417 Self {
418 origin: Instant::now(),
419 }
420 }
421}
422
423impl Default for MonotonicClock {
424 fn default() -> Self {
425 Self::new()
426 }
427}
428
429impl RuntimeClock for MonotonicClock {
430 fn now_seconds(&self) -> f64 {
431 self.origin.elapsed().as_secs_f64()
432 }
433}
434
435pub trait RefinementEventSink: Send {
437 fn emit(&mut self, event: &RefinementEvent) -> Result<(), String>;
444}
445
446impl<F> RefinementEventSink for F
447where
448 F: FnMut(&RefinementEvent) -> Result<(), String> + Send,
449{
450 fn emit(&mut self, event: &RefinementEvent) -> Result<(), String> {
451 self(event)
452 }
453}
454
455pub trait CheckpointSink<C>: Send {
457 fn checkpoint(&mut self, checkpoint: &C) -> Result<(), String>;
463}
464
465impl<C, F> CheckpointSink<C> for F
466where
467 F: FnMut(&C) -> Result<(), String> + Send,
468{
469 fn checkpoint(&mut self, checkpoint: &C) -> Result<(), String> {
470 self(checkpoint)
471 }
472}
473
474#[derive(Clone, Debug, PartialEq, Eq)]
476pub struct RefinementStop {
477 pub reason: TerminationReason,
479 pub message: String,
481}
482
483impl Display for RefinementStop {
484 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
485 formatter.write_str(&self.message)
486 }
487}
488
489impl Error for RefinementStop {}
490
491pub struct RefinementRuntime<C = ()> {
493 limits: RefinementLimits,
494 cancellation: Option<CancellationToken>,
495 event_sink: Option<Box<dyn RefinementEventSink>>,
496 checkpoint_sink: Option<Box<dyn CheckpointSink<C>>>,
497 clock: Arc<dyn RuntimeClock>,
498 started_at: f64,
499 attempted_iteration: usize,
500 accepted_iterations: usize,
501 evaluations: usize,
502 consecutive_rejections: usize,
503 event_sink_error: Option<String>,
504}
505
506impl<C> RefinementRuntime<C> {
507 pub fn new(
513 limits: RefinementLimits,
514 cancellation: Option<CancellationToken>,
515 ) -> Result<Self, RuntimeError> {
516 Self::with_clock(limits, cancellation, Arc::new(MonotonicClock::new()))
517 }
518
519 pub fn with_clock(
525 limits: RefinementLimits,
526 cancellation: Option<CancellationToken>,
527 clock: Arc<dyn RuntimeClock>,
528 ) -> Result<Self, RuntimeError> {
529 let started_at = clock.now_seconds();
530 if !started_at.is_finite() {
531 return Err(RuntimeError::InvalidClock);
532 }
533 Ok(Self {
534 limits,
535 cancellation,
536 event_sink: None,
537 checkpoint_sink: None,
538 clock,
539 started_at,
540 attempted_iteration: 0,
541 accepted_iterations: 0,
542 evaluations: 0,
543 consecutive_rejections: 0,
544 event_sink_error: None,
545 })
546 }
547
548 pub fn set_event_sink(&mut self, sink: impl RefinementEventSink + 'static) {
550 self.event_sink = Some(Box::new(sink));
551 self.event_sink_error = None;
552 }
553
554 pub fn set_checkpoint_sink(&mut self, sink: impl CheckpointSink<C> + 'static) {
556 self.checkpoint_sink = Some(Box::new(sink));
557 }
558
559 pub fn resume_accepted(&mut self, completed_iterations: usize) -> Result<(), RuntimeError> {
568 if self.attempted_iteration != 0
569 || self.accepted_iterations != 0
570 || self.evaluations != 0
571 || completed_iterations > self.limits.max_iterations
572 {
573 return Err(RuntimeError::InvalidResume);
574 }
575 self.attempted_iteration = completed_iterations;
576 self.accepted_iterations = completed_iterations;
577 Ok(())
578 }
579
580 pub fn elapsed_seconds(&self) -> Result<f64, RuntimeError> {
586 let elapsed = self.clock.now_seconds() - self.started_at;
587 if !elapsed.is_finite() || elapsed < 0.0 {
588 return Err(RuntimeError::InvalidClock);
589 }
590 Ok(elapsed)
591 }
592
593 pub fn emit(
600 &mut self,
601 kind: RefinementEventKind,
602 stage: impl Into<String>,
603 message: impl Into<String>,
604 diagnostics: Vec<(String, DiagnosticValue)>,
605 ) -> Result<RefinementEvent, RuntimeError> {
606 let event = RefinementEvent::new(
607 kind,
608 stage,
609 self.attempted_iteration,
610 self.accepted_iterations,
611 self.evaluations,
612 self.elapsed_seconds()?,
613 message,
614 diagnostics,
615 )?;
616 let sink_result = self.event_sink.as_mut().map(|sink| sink.emit(&event));
617 if let Some(Err(message)) = sink_result {
618 self.event_sink_error = Some(message);
619 self.event_sink = None;
620 }
621 Ok(event)
622 }
623
624 pub fn check_boundary(&self) -> Result<(), RuntimeError> {
631 if let Some(token) = &self.cancellation
632 && token.is_requested()
633 {
634 let message = token
635 .reason()
636 .map_err(RuntimeError::Cancellation)?
637 .unwrap_or_else(|| "user requested cancellation".to_owned());
638 return Err(RuntimeError::Stopped(RefinementStop {
639 reason: TerminationReason::Cancelled,
640 message,
641 }));
642 }
643 if let Some(limit) = self.limits.max_runtime_seconds {
644 let elapsed = self.elapsed_seconds()?;
645 if elapsed >= limit {
646 return Err(RuntimeError::Stopped(RefinementStop {
647 reason: TerminationReason::MaxRuntime,
648 message: "runtime limit reached".to_owned(),
649 }));
650 }
651 }
652 if self.evaluations >= self.limits.max_evaluations {
653 return Err(RuntimeError::Stopped(RefinementStop {
654 reason: TerminationReason::MaxEvaluations,
655 message: "model-evaluation limit reached".to_owned(),
656 }));
657 }
658 Ok(())
659 }
660
661 pub fn begin_iteration(&mut self, attempted_iteration: usize) -> Result<(), RuntimeError> {
667 if attempted_iteration <= self.attempted_iteration {
668 return Err(RuntimeError::InvalidIterationOrder);
669 }
670 if attempted_iteration > self.limits.max_iterations {
671 return Err(RuntimeError::Stopped(RefinementStop {
672 reason: TerminationReason::MaxIterations,
673 message: "iteration limit reached".to_owned(),
674 }));
675 }
676 self.attempted_iteration = attempted_iteration;
677 self.check_boundary()
678 }
679
680 pub fn begin_evaluation(&mut self) -> Result<(), RuntimeError> {
686 self.check_boundary()?;
687 self.evaluations = self
688 .evaluations
689 .checked_add(1)
690 .ok_or(RuntimeError::CounterOverflow)?;
691 Ok(())
692 }
693
694 pub fn accept_step(&mut self, checkpoint: Option<&C>) -> Result<(), RuntimeError> {
704 if self.accepted_iterations >= self.attempted_iteration {
705 return Err(RuntimeError::DuplicateAcceptance);
706 }
707 self.accepted_iterations = self
708 .accepted_iterations
709 .checked_add(1)
710 .ok_or(RuntimeError::CounterOverflow)?;
711 self.consecutive_rejections = 0;
712 if let Some(checkpoint) = checkpoint
713 && let Some(sink) = self.checkpoint_sink.as_mut()
714 {
715 sink.checkpoint(checkpoint)
716 .map_err(|message| RuntimeError::CheckpointSink { message })?;
717 self.emit(
718 RefinementEventKind::Checkpoint,
719 "checkpoint",
720 "accepted-state checkpoint completed",
721 Vec::new(),
722 )?;
723 }
724 Ok(())
725 }
726
727 pub fn reject_step(&mut self) -> Result<(), RuntimeError> {
733 self.consecutive_rejections = self
734 .consecutive_rejections
735 .checked_add(1)
736 .ok_or(RuntimeError::CounterOverflow)?;
737 if self.consecutive_rejections >= self.limits.max_consecutive_rejections {
738 return Err(RuntimeError::Stopped(RefinementStop {
739 reason: TerminationReason::RepeatedRejections,
740 message: "consecutive rejected-step limit reached".to_owned(),
741 }));
742 }
743 Ok(())
744 }
745
746 #[must_use]
748 pub const fn attempted_iteration(&self) -> usize {
749 self.attempted_iteration
750 }
751
752 #[must_use]
754 pub const fn accepted_iterations(&self) -> usize {
755 self.accepted_iterations
756 }
757
758 #[must_use]
760 pub const fn evaluations(&self) -> usize {
761 self.evaluations
762 }
763
764 #[must_use]
766 pub const fn consecutive_rejections(&self) -> usize {
767 self.consecutive_rejections
768 }
769
770 #[must_use]
772 pub fn event_sink_error(&self) -> Option<&str> {
773 self.event_sink_error.as_deref()
774 }
775
776 #[must_use]
778 pub const fn has_event_sink(&self) -> bool {
779 self.event_sink.is_some()
780 }
781}
782
783#[derive(Debug)]
785pub enum RuntimeError {
786 InvalidLimits,
788 InvalidClock,
790 InvalidEvent {
792 message: String,
794 },
795 InvalidResume,
797 InvalidIterationOrder,
799 DuplicateAcceptance,
801 CounterOverflow,
803 Cancellation(CancellationError),
805 Stopped(RefinementStop),
807 CheckpointSink {
809 message: String,
811 },
812}
813
814impl Display for RuntimeError {
815 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
816 match self {
817 Self::InvalidLimits => {
818 formatter.write_str("refinement limits must be positive and finite")
819 }
820 Self::InvalidClock => {
821 formatter.write_str("refinement clock must be finite and monotonic")
822 }
823 Self::InvalidEvent { message } | Self::CheckpointSink { message } => {
824 formatter.write_str(message)
825 }
826 Self::InvalidResume => {
827 formatter.write_str("refinement counters cannot be resumed in this state")
828 }
829 Self::InvalidIterationOrder => {
830 formatter.write_str("attempted iterations must increase strictly")
831 }
832 Self::DuplicateAcceptance => {
833 formatter.write_str("at most one step may be accepted per attempted iteration")
834 }
835 Self::CounterOverflow => formatter.write_str("refinement runtime counter overflow"),
836 Self::Cancellation(error) => Display::fmt(error, formatter),
837 Self::Stopped(stop) => Display::fmt(stop, formatter),
838 }
839 }
840}
841
842impl Error for RuntimeError {
843 fn source(&self) -> Option<&(dyn Error + 'static)> {
844 match self {
845 Self::Cancellation(error) => Some(error),
846 Self::Stopped(stop) => Some(stop),
847 _ => None,
848 }
849 }
850}
851
852fn invalid_event(message: &str) -> RuntimeError {
853 RuntimeError::InvalidEvent {
854 message: message.to_owned(),
855 }
856}