use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct SearchBudget {
pub max_expansions: Option<usize>,
pub max_duration: Option<Duration>,
}
impl SearchBudget {
pub const UNLIMITED: Self = Self {
max_expansions: None,
max_duration: None,
};
#[must_use]
pub const fn max_expansions(limit: usize) -> Self {
Self {
max_expansions: Some(limit),
max_duration: None,
}
}
#[must_use]
pub const fn max_duration(limit: Duration) -> Self {
Self {
max_expansions: None,
max_duration: Some(limit),
}
}
#[must_use]
pub const fn with_max_expansions(mut self, limit: usize) -> Self {
self.max_expansions = Some(limit);
self
}
#[must_use]
pub const fn with_max_duration(mut self, limit: Duration) -> Self {
self.max_duration = Some(limit);
self
}
#[must_use]
pub const fn is_unlimited(&self) -> bool {
self.max_expansions.is_none() && self.max_duration.is_none()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum BudgetExhausted {
Expansions {
limit: usize,
expansions: usize,
},
Duration {
limit: Duration,
},
}
impl std::fmt::Display for BudgetExhausted {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Expansions { limit, expansions } => write!(
formatter,
"search budget exhausted after {expansions} expansions (limit {limit})"
),
Self::Duration { limit } => {
write!(formatter, "search budget exhausted after {limit:?}")
}
}
}
}
impl std::error::Error for BudgetExhausted {}
#[derive(Debug, Clone, Copy)]
pub struct BudgetWatch {
budget: SearchBudget,
deadline: Option<Instant>,
}
impl BudgetWatch {
#[must_use]
pub fn start(budget: SearchBudget) -> Self {
let deadline = budget.max_duration.map(|limit| Instant::now() + limit);
Self { budget, deadline }
}
#[must_use]
pub const fn budget(&self) -> SearchBudget {
self.budget
}
#[must_use]
pub const fn is_unlimited(&self) -> bool {
self.budget.is_unlimited()
}
pub fn check(&self, expansions: usize) -> Result<(), BudgetExhausted> {
if let Some(limit) = self.budget.max_expansions
&& expansions >= limit
{
return Err(BudgetExhausted::Expansions { limit, expansions });
}
if let (Some(deadline), Some(limit)) = (self.deadline, self.budget.max_duration)
&& Instant::now() >= deadline
{
return Err(BudgetExhausted::Duration { limit });
}
Ok(())
}
#[must_use]
pub const fn has_limits(&self) -> bool {
!self.budget.is_unlimited()
}
}
#[cfg(test)]
mod tests {
use super::{BudgetExhausted, BudgetWatch, SearchBudget};
use std::time::Duration;
#[test]
fn unlimited_watch_never_exhausts() {
let watch = BudgetWatch::start(SearchBudget::UNLIMITED);
assert!(watch.check(0).is_ok());
assert!(watch.check(usize::MAX).is_ok());
}
#[test]
fn expansion_limit_trips_at_limit() {
let watch = BudgetWatch::start(SearchBudget::max_expansions(3));
assert!(watch.check(2).is_ok());
assert_eq!(
watch.check(3),
Err(BudgetExhausted::Expansions {
limit: 3,
expansions: 3
})
);
}
#[test]
fn duration_limit_trips_after_deadline() {
let watch = BudgetWatch::start(SearchBudget::max_duration(Duration::from_millis(0)));
let result = watch.check(0);
assert!(matches!(result, Err(BudgetExhausted::Duration { .. })));
}
}