use rexact_syntax::error::ErrorCode;
use std::time::Duration;
#[derive(Debug, Clone, Copy)]
pub struct Clock {
#[cfg(not(target_arch = "wasm32"))]
start: std::time::Instant,
}
impl Clock {
pub fn now() -> Self {
Clock {
#[cfg(not(target_arch = "wasm32"))]
start: std::time::Instant::now(),
}
}
pub fn elapsed(&self) -> Duration {
#[cfg(not(target_arch = "wasm32"))]
{
self.start.elapsed()
}
#[cfg(target_arch = "wasm32")]
{
Duration::ZERO
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Limits {
pub states: u64,
pub time: Duration,
pub memory_bytes: u64,
pub repeat: u64,
pub length: u64,
pub depth: u32,
pub enumerate: u64,
}
impl Default for Limits {
fn default() -> Self {
Self {
states: 200_000,
time: Duration::from_millis(5_000),
memory_bytes: 1024 * 1024 * 1024,
repeat: 10_000,
length: 1_000_000,
depth: 250,
enumerate: 100_000,
}
}
}
#[derive(Debug)]
pub struct Budget {
limits: Limits,
start: Clock,
states: u64,
memory: u64,
hit: Option<ErrorCode>,
}
impl Budget {
pub fn new(limits: Limits) -> Self {
Self {
limits,
start: Clock::now(),
states: 0,
memory: 0,
hit: None,
}
}
pub fn limits(&self) -> &Limits {
&self.limits
}
pub fn limit_hit(&self) -> Option<ErrorCode> {
self.hit
}
fn trip(&mut self, code: ErrorCode) -> Result<(), ErrorCode> {
self.hit.get_or_insert(code);
Err(code)
}
pub fn add_states(&mut self, n: u64) -> Result<(), ErrorCode> {
match self.states.checked_add(n) {
Some(total) if total <= self.limits.states => {
self.states = total;
Ok(())
}
_ => self.trip(ErrorCode::L001),
}
}
pub fn add_memory(&mut self, bytes: u64) -> Result<(), ErrorCode> {
match self.memory.checked_add(bytes) {
Some(total) if total <= self.limits.memory_bytes => {
self.memory = total;
Ok(())
}
_ => self.trip(ErrorCode::L003),
}
}
pub fn check_time(&mut self) -> Result<(), ErrorCode> {
if self.start.elapsed() > self.limits.time {
self.trip(ErrorCode::L002)
} else {
Ok(())
}
}
pub fn check_repeat(&mut self, count: u64) -> Result<(), ErrorCode> {
if count > self.limits.repeat {
self.trip(ErrorCode::L004)
} else {
Ok(())
}
}
pub fn check_length(&mut self, len: u64) -> Result<(), ErrorCode> {
if len > self.limits.length {
self.trip(ErrorCode::L005)
} else {
Ok(())
}
}
pub fn check_depth(&mut self, depth: u32) -> Result<(), ErrorCode> {
if depth > self.limits.depth {
self.trip(ErrorCode::L006)
} else {
Ok(())
}
}
pub fn check_enumerate(&mut self, count: u64) -> Result<(), ErrorCode> {
if count > self.limits.enumerate {
self.trip(ErrorCode::L007)
} else {
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn small() -> Limits {
Limits {
states: 10,
repeat: 5,
length: 8,
depth: 3,
enumerate: 4,
..Limits::default()
}
}
#[test]
fn states_breach_reports_l001_and_remembers_it() {
let mut b = Budget::new(small());
assert!(b.add_states(10).is_ok());
assert_eq!(b.add_states(1), Err(ErrorCode::L001));
assert_eq!(b.limit_hit(), Some(ErrorCode::L001));
}
#[test]
fn first_breach_wins() {
let mut b = Budget::new(small());
let _ = b.check_depth(99);
let _ = b.check_repeat(99);
assert_eq!(b.limit_hit(), Some(ErrorCode::L006));
}
#[test]
fn overflow_is_a_breach_not_a_wrap() {
let mut b = Budget::new(Limits {
states: u64::MAX,
..small()
});
assert!(b.add_states(u64::MAX).is_ok());
assert_eq!(b.add_states(1), Err(ErrorCode::L001));
let mut m = Budget::new(Limits {
memory_bytes: u64::MAX,
..small()
});
assert!(m.add_memory(u64::MAX).is_ok());
assert_eq!(m.add_memory(1), Err(ErrorCode::L003));
}
#[test]
fn boundaries_are_inclusive() {
let mut b = Budget::new(small());
assert!(b.check_repeat(5).is_ok());
assert_eq!(b.check_repeat(6), Err(ErrorCode::L004));
let mut b = Budget::new(small());
assert!(b.check_length(8).is_ok());
assert_eq!(b.check_length(9), Err(ErrorCode::L005));
let mut b = Budget::new(small());
assert!(b.check_depth(3).is_ok());
assert_eq!(b.check_depth(4), Err(ErrorCode::L006));
let mut b = Budget::new(small());
assert!(b.check_enumerate(4).is_ok());
assert_eq!(b.check_enumerate(5), Err(ErrorCode::L007));
}
#[test]
fn time_limit_trips_l002() {
let mut b = Budget::new(Limits {
time: Duration::ZERO,
..small()
});
std::thread::sleep(Duration::from_millis(2));
assert_eq!(b.check_time(), Err(ErrorCode::L002));
}
#[test]
fn memory_breach_reports_l003() {
let mut b = Budget::new(Limits {
memory_bytes: 100,
..small()
});
assert!(b.add_memory(100).is_ok());
assert_eq!(b.add_memory(1), Err(ErrorCode::L003));
}
}