use crate::memory_stats::{get_or_register_slot, update_arena_stats};
use bumpalo::Bump;
use std::cell::RefCell;
const ARENA_RESET_THRESHOLD: usize = 10 * 1024 * 1024;
thread_local! {
static ARENA: RefCell<Bump> = {
get_or_register_slot();
RefCell::new(Bump::new())
};
static ARENA_BYTES_ALLOCATED: RefCell<usize> = const { RefCell::new(0) };
}
pub fn with_arena<F, R>(f: F) -> R
where
F: FnOnce(&Bump) -> R,
{
ARENA.with(|arena| {
let bump = arena.borrow();
let result = f(&bump);
let allocated = bump.allocated_bytes();
drop(bump);
ARENA_BYTES_ALLOCATED.with(|bytes| {
*bytes.borrow_mut() = allocated;
});
update_arena_stats(allocated);
if should_reset() {
arena_reset();
}
result
})
}
pub fn arena_reset() {
ARENA.with(|arena| {
arena.borrow_mut().reset();
});
ARENA_BYTES_ALLOCATED.with(|bytes| {
*bytes.borrow_mut() = 0;
});
update_arena_stats(0);
}
fn should_reset() -> bool {
ARENA_BYTES_ALLOCATED.with(|bytes| *bytes.borrow() > ARENA_RESET_THRESHOLD)
}
pub fn arena_stats() -> ArenaStats {
let allocated = ARENA_BYTES_ALLOCATED.with(|bytes| *bytes.borrow());
ArenaStats {
allocated_bytes: allocated,
}
}
#[derive(Debug, Clone, Copy)]
pub struct ArenaStats {
pub allocated_bytes: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_arena_reset() {
arena_reset();
with_arena(|arena| {
let _s1 = arena.alloc_str("Hello");
let _s2 = arena.alloc_str("World");
});
let stats_before = arena_stats();
assert!(stats_before.allocated_bytes > 0);
arena_reset();
let stats_after = arena_stats();
assert!(
stats_after.allocated_bytes < stats_before.allocated_bytes,
"Arena should have less memory after reset (before: {}, after: {})",
stats_before.allocated_bytes,
stats_after.allocated_bytes
);
}
#[test]
fn test_with_arena() {
arena_reset();
let len = with_arena(|arena| {
let s = arena.alloc_str("Test string");
assert_eq!(s, "Test string");
s.len()
});
assert_eq!(len, 11);
let stats = arena_stats();
assert!(stats.allocated_bytes > 0);
}
#[test]
fn test_auto_reset_threshold() {
arena_reset();
let big_str = "x".repeat(ARENA_RESET_THRESHOLD / 2);
with_arena(|arena| {
let _s = arena.alloc_str(&big_str);
});
let stats1 = arena_stats();
let initial_bytes = stats1.allocated_bytes;
assert!(initial_bytes > 0);
let big_str2 = "y".repeat(ARENA_RESET_THRESHOLD / 2 + 1000);
with_arena(|arena| {
let _s = arena.alloc_str(&big_str2);
});
let stats2 = arena_stats();
assert!(
stats2.allocated_bytes < initial_bytes + (ARENA_RESET_THRESHOLD / 2 + 2000),
"Arena should have reset: stats2={}, initial={}, threshold={}",
stats2.allocated_bytes,
initial_bytes,
ARENA_RESET_THRESHOLD
);
}
}