use crate::error::ArenaError;
#[derive(Clone)]
pub struct Segment {
data: Vec<f32>,
cursor: usize,
}
impl Segment {
pub fn new(capacity: u32) -> Self {
Self {
data: vec![0.0; capacity as usize],
cursor: 0,
}
}
pub fn alloc(&mut self, len: u32) -> Option<(u32, &mut [f32])> {
let len = len as usize;
let new_cursor = self.cursor.checked_add(len)?;
if new_cursor > self.data.len() {
return None;
}
let offset = self.cursor as u32;
let slice = &mut self.data[self.cursor..new_cursor];
self.cursor = new_cursor;
slice.fill(0.0);
Some((offset, slice))
}
pub fn slice(&self, offset: u32, len: u32) -> Option<&[f32]> {
let start = offset as usize;
let end = start.checked_add(len as usize)?;
if end > self.cursor {
return None;
}
Some(&self.data[start..end])
}
pub fn slice_mut(&mut self, offset: u32, len: u32) -> Option<&mut [f32]> {
let start = offset as usize;
let end = start.checked_add(len as usize)?;
if end > self.cursor {
return None;
}
Some(&mut self.data[start..end])
}
pub fn reset(&mut self) {
self.cursor = 0;
}
pub fn used(&self) -> usize {
self.cursor
}
pub fn capacity(&self) -> usize {
self.data.len()
}
pub fn remaining(&self) -> usize {
self.data.len() - self.cursor
}
pub fn memory_bytes(&self) -> usize {
self.data.len() * std::mem::size_of::<f32>()
}
}
pub struct SegmentList {
segments: Vec<Segment>,
segment_size: u32,
max_segments: u16,
current: usize,
}
impl SegmentList {
pub fn new(segment_size: u32, max_segments: u16) -> Self {
let mut segments = Vec::with_capacity(max_segments as usize);
segments.push(Segment::new(segment_size));
Self {
segments,
segment_size,
max_segments,
current: 0,
}
}
pub fn alloc(&mut self, len: u32) -> Result<(u16, u32), ArenaError> {
if len > self.segment_size {
return Err(ArenaError::CapacityExceeded {
requested: len as usize * std::mem::size_of::<f32>(),
capacity: self.segment_size as usize * std::mem::size_of::<f32>(),
});
}
if let Some((offset, _slice)) = self.segments[self.current].alloc(len) {
return Ok((self.current as u16, offset));
}
let next = self.current + 1;
if next < self.segments.len() {
if let Some((offset, _slice)) = self.segments[next].alloc(len) {
self.current = next;
return Ok((next as u16, offset));
}
}
if self.segments.len() >= self.max_segments as usize {
return Err(ArenaError::CapacityExceeded {
requested: len as usize * std::mem::size_of::<f32>(),
capacity: self.total_capacity_bytes(),
});
}
let mut seg = Segment::new(self.segment_size);
let (offset, _slice) = seg
.alloc(len)
.expect("len <= segment_size, so fresh segment always fits");
self.segments.push(seg);
self.current = self.segments.len() - 1;
Ok((self.current as u16, offset))
}
pub fn slice(&self, segment_index: u16, offset: u32, len: u32) -> Option<&[f32]> {
self.segments
.get(segment_index as usize)?
.slice(offset, len)
}
pub fn slice_mut(&mut self, segment_index: u16, offset: u32, len: u32) -> Option<&mut [f32]> {
self.segments
.get_mut(segment_index as usize)?
.slice_mut(offset, len)
}
pub fn reset(&mut self) {
for seg in &mut self.segments {
seg.reset();
}
self.current = 0;
}
pub fn segment_count(&self) -> usize {
self.segments.len()
}
pub fn memory_bytes(&self) -> usize {
self.segments.iter().map(|s| s.memory_bytes()).sum()
}
pub fn total_used(&self) -> usize {
self.segments.iter().map(|s| s.used()).sum()
}
fn total_capacity_bytes(&self) -> usize {
self.segments.len() * self.segment_size as usize * std::mem::size_of::<f32>()
}
}
impl Clone for SegmentList {
fn clone(&self) -> Self {
Self {
segments: self.segments[..=self.current].to_vec(),
segment_size: self.segment_size,
max_segments: self.max_segments,
current: self.current,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn segment_alloc_returns_zeroed_data() {
let mut seg = Segment::new(1024);
let (offset, data) = seg.alloc(10).unwrap();
assert_eq!(offset, 0);
assert_eq!(data.len(), 10);
assert!(data.iter().all(|&v| v == 0.0));
}
#[test]
fn segment_sequential_alloc() {
let mut seg = Segment::new(1024);
let (off1, _) = seg.alloc(100).unwrap();
let (off2, _) = seg.alloc(200).unwrap();
assert_eq!(off1, 0);
assert_eq!(off2, 100);
assert_eq!(seg.used(), 300);
}
#[test]
fn segment_alloc_fails_when_full() {
let mut seg = Segment::new(100);
assert!(seg.alloc(100).is_some());
assert!(seg.alloc(1).is_none());
}
#[test]
fn segment_reset_allows_realloc() {
let mut seg = Segment::new(100);
seg.alloc(100).unwrap();
seg.reset();
assert_eq!(seg.used(), 0);
assert!(seg.alloc(50).is_some());
}
#[test]
fn segment_slice_reads_written_data() {
let mut seg = Segment::new(1024);
let (offset, data) = seg.alloc(5).unwrap();
data[0] = 1.0;
data[4] = 5.0;
let read = seg.slice(offset, 5).unwrap();
assert_eq!(read[0], 1.0);
assert_eq!(read[4], 5.0);
}
#[test]
fn segment_list_alloc_within_first_segment() {
let mut list = SegmentList::new(1024, 4);
let (seg_idx, offset) = list.alloc(10).unwrap();
assert_eq!(seg_idx, 0);
assert_eq!(offset, 0);
}
#[test]
fn segment_list_grows_on_overflow() {
let mut list = SegmentList::new(100, 4);
list.alloc(100).unwrap(); let (seg_idx, _) = list.alloc(50).unwrap(); assert_eq!(seg_idx, 1);
assert_eq!(list.segment_count(), 2);
}
#[test]
fn segment_list_capacity_exceeded() {
let mut list = SegmentList::new(100, 2);
list.alloc(100).unwrap(); list.alloc(100).unwrap(); let result = list.alloc(1);
assert!(matches!(result, Err(ArenaError::CapacityExceeded { .. })));
}
#[test]
fn segment_list_reset() {
let mut list = SegmentList::new(100, 4);
list.alloc(80).unwrap();
list.alloc(80).unwrap(); assert_eq!(list.segment_count(), 2);
list.reset();
assert_eq!(list.total_used(), 0);
let (seg_idx, offset) = list.alloc(10).unwrap();
assert_eq!(seg_idx, 0);
assert_eq!(offset, 0);
}
#[test]
fn segment_list_slice_roundtrip() {
let mut list = SegmentList::new(1024, 4);
let (seg, off) = list.alloc(5).unwrap();
{
let s = list.slice_mut(seg, off, 5).unwrap();
s[0] = 42.0;
}
let read = list.slice(seg, off, 5).unwrap();
assert_eq!(read[0], 42.0);
}
#[test]
fn oversized_alloc_returns_error_not_panic() {
let mut list = SegmentList::new(100, 4);
let result = list.alloc(101);
assert!(matches!(result, Err(ArenaError::CapacityExceeded { .. })));
}
#[test]
fn exactly_segment_size_alloc_succeeds() {
let mut list = SegmentList::new(100, 4);
let result = list.alloc(100);
assert!(result.is_ok());
}
#[test]
fn test_segment_clone_preserves_data() {
let mut seg = Segment::new(1024);
let (offset, data) = seg.alloc(5).unwrap();
data[0] = 1.0;
data[4] = 5.0;
let cloned = seg.clone();
let read = cloned.slice(offset, 5).unwrap();
assert_eq!(read[0], 1.0);
assert_eq!(read[4], 5.0);
assert_eq!(cloned.used(), 5);
assert_eq!(cloned.capacity(), 1024);
}
#[test]
fn test_segment_clone_independent() {
let mut seg = Segment::new(1024);
let (_offset, data) = seg.alloc(5).unwrap();
data[0] = 1.0;
let mut cloned = seg.clone();
let data = seg.slice_mut(0, 5).unwrap();
data[0] = 99.0;
assert_eq!(cloned.slice(0, 5).unwrap()[0], 1.0);
let cdata = cloned.slice_mut(0, 5).unwrap();
cdata[0] = 77.0;
assert_eq!(seg.slice(0, 5).unwrap()[0], 99.0);
}
#[test]
fn test_segment_list_clone_preserves_slices() {
let mut list = SegmentList::new(100, 4);
let (seg0, off0) = list.alloc(80).unwrap();
list.slice_mut(seg0, off0, 80).unwrap()[0] = 42.0;
let (seg1, off1) = list.alloc(80).unwrap(); list.slice_mut(seg1, off1, 80).unwrap()[0] = 99.0;
let cloned = list.clone();
assert_eq!(cloned.segment_count(), 2);
assert_eq!(cloned.slice(seg0, off0, 80).unwrap()[0], 42.0);
assert_eq!(cloned.slice(seg1, off1, 80).unwrap()[0], 99.0);
}
#[test]
fn test_segment_list_clone_independent() {
let mut list = SegmentList::new(100, 4);
let (seg, off) = list.alloc(10).unwrap();
list.slice_mut(seg, off, 10).unwrap()[0] = 42.0;
let cloned = list.clone();
list.slice_mut(seg, off, 10).unwrap()[0] = 0.0;
assert_eq!(cloned.slice(seg, off, 10).unwrap()[0], 42.0);
}
#[test]
fn segment_slice_returns_none_after_reset() {
let mut seg = Segment::new(1024);
let (_, data) = seg.alloc(100).unwrap();
data[0] = 42.0;
seg.reset();
assert!(seg.slice(0, 100).is_none());
}
#[test]
fn segment_slice_mut_returns_none_after_reset() {
let mut seg = Segment::new(1024);
seg.alloc(100).unwrap();
seg.reset();
assert!(seg.slice_mut(0, 100).is_none());
}
#[test]
fn segment_slice_returns_none_beyond_cursor() {
let mut seg = Segment::new(1024);
seg.alloc(50).unwrap();
assert!(seg.slice(0, 100).is_none());
}
#[test]
fn segment_slice_within_cursor_succeeds() {
let mut seg = Segment::new(1024);
seg.alloc(100).unwrap();
let data = seg.slice(0, 100).unwrap();
assert_eq!(data.len(), 100);
}
#[test]
fn segment_list_slice_returns_none_after_reset() {
let mut list = SegmentList::new(1024, 4);
let (seg, off) = list.alloc(100).unwrap();
list.slice_mut(seg, off, 100).unwrap()[0] = 42.0;
list.reset();
assert!(list.slice(seg, off, 100).is_none());
}
}