use crate::types::Atom;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
static THUNK_ID_COUNTER: AtomicU64 = AtomicU64::new(1);
#[derive(Debug, Clone)]
pub struct Thunk {
pub id: u64,
pub created_at: u64,
pub probation_deadline: u64,
pub generator_id: String,
pub context: Vec<u8>,
evaluated: Option<Atom>,
}
impl Thunk {
pub fn new(generator_id: &str, context: Vec<u8>, probation_micros: u64) -> Self {
let now = now_micros();
Self {
id: THUNK_ID_COUNTER.fetch_add(1, Ordering::SeqCst),
created_at: now,
probation_deadline: now + probation_micros,
generator_id: generator_id.to_string(),
context,
evaluated: None,
}
}
pub fn in_probation(&self) -> bool {
now_micros() < self.probation_deadline
}
pub fn is_evaluated(&self) -> bool {
self.evaluated.is_some()
}
pub fn evaluate(&mut self, value: Atom) {
self.evaluated = Some(value);
}
pub fn value(&self) -> Option<&Atom> {
self.evaluated.as_ref()
}
pub fn time_remaining(&self) -> u64 {
self.probation_deadline.saturating_sub(now_micros())
}
}
#[derive(Debug)]
pub struct ThunkRegistry {
thunks: HashMap<u64, Thunk>,
default_probation: u64,
}
impl ThunkRegistry {
pub fn new() -> Self {
Self::with_probation(5_000_000) }
pub fn with_probation(probation_micros: u64) -> Self {
Self {
thunks: HashMap::new(),
default_probation: probation_micros,
}
}
pub fn register(&mut self, generator_id: &str, context: Vec<u8>) -> u64 {
let thunk = Thunk::new(generator_id, context, self.default_probation);
let id = thunk.id;
self.thunks.insert(id, thunk);
id
}
pub fn get(&self, id: u64) -> Option<&Thunk> {
self.thunks.get(&id)
}
pub fn get_mut(&mut self, id: u64) -> Option<&mut Thunk> {
self.thunks.get_mut(&id)
}
pub fn evaluate(&mut self, id: u64, value: Atom) -> bool {
if let Some(thunk) = self.thunks.get_mut(&id) {
thunk.evaluate(value);
true
} else {
false
}
}
pub fn gc(&mut self) -> usize {
let before = self.thunks.len();
self.thunks.retain(|_, thunk| {
thunk.is_evaluated() || thunk.in_probation()
});
before - self.thunks.len()
}
pub fn len(&self) -> usize {
self.thunks.len()
}
pub fn is_empty(&self) -> bool {
self.thunks.is_empty()
}
pub fn unevaluated_count(&self) -> usize {
self.thunks.values().filter(|t| !t.is_evaluated()).count()
}
pub fn expired_count(&self) -> usize {
self.thunks
.values()
.filter(|t| !t.is_evaluated() && !t.in_probation())
.count()
}
}
impl Default for ThunkRegistry {
fn default() -> Self {
Self::new()
}
}
fn now_micros() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_micros() as u64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_thunk_probation() {
let thunk = Thunk::new("test_gen", vec![], 1_000_000);
assert!(thunk.in_probation());
assert!(!thunk.is_evaluated());
}
#[test]
fn test_thunk_evaluation() {
let mut thunk = Thunk::new("test_gen", vec![], 1_000_000);
thunk.evaluate(Atom::Float(3.14));
assert!(thunk.is_evaluated());
assert_eq!(thunk.value(), Some(&Atom::Float(3.14)));
}
#[test]
fn test_registry_gc() {
let mut registry = ThunkRegistry::with_probation(0);
registry.register("gen1", vec![]);
registry.register("gen2", vec![]);
std::thread::sleep(std::time::Duration::from_millis(1));
let collected = registry.gc();
assert_eq!(collected, 2);
assert!(registry.is_empty());
}
}