pub struct SendCopyPool {
backing: Vec<u8>,
slot_size: u32,
count: u16,
free_list: Vec<u16>,
slot_offset: Vec<u32>, slot_remaining: Vec<u32>, in_use: Vec<bool>, }
impl SendCopyPool {
pub fn new(count: u16, slot_size: u32) -> Self {
let total = count as usize * slot_size as usize;
let backing = vec![0u8; total];
let free_list: Vec<u16> = (0..count).rev().collect();
let n = count as usize;
SendCopyPool {
backing,
slot_size,
count,
free_list,
slot_offset: vec![0u32; n],
slot_remaining: vec![0u32; n],
in_use: vec![false; n],
}
}
pub fn copy_in(&mut self, data: &[u8]) -> Option<(u16, *const u8, u32)> {
if data.len() > self.slot_size as usize {
return None;
}
let idx = self.free_list.pop()?;
let offset = idx as usize * self.slot_size as usize;
self.backing[offset..offset + data.len()].copy_from_slice(data);
let ptr = self.backing.as_ptr().wrapping_add(offset);
self.slot_offset[idx as usize] = 0;
self.slot_remaining[idx as usize] = data.len() as u32;
self.in_use[idx as usize] = true;
Some((idx, ptr, data.len() as u32))
}
pub unsafe fn copy_in_gather(
&mut self,
parts: &[(*const u8, usize)],
total_len: usize,
) -> Option<(u16, *const u8, u32)> {
if total_len > self.slot_size as usize {
return None;
}
let idx = self.free_list.pop()?;
let base = idx as usize * self.slot_size as usize;
let mut dest_offset = 0;
for &(ptr, len) in parts {
let src = unsafe { std::slice::from_raw_parts(ptr, len) };
self.backing[base + dest_offset..base + dest_offset + len].copy_from_slice(src);
dest_offset += len;
}
let out_ptr = self.backing.as_ptr().wrapping_add(base);
self.slot_offset[idx as usize] = 0;
self.slot_remaining[idx as usize] = total_len as u32;
self.in_use[idx as usize] = true;
Some((idx, out_ptr, total_len as u32))
}
pub fn release(&mut self, idx: u16) {
debug_assert!((idx as usize) < self.count as usize);
if !self.in_use[idx as usize] {
return; }
self.in_use[idx as usize] = false;
self.slot_offset[idx as usize] = 0;
self.slot_remaining[idx as usize] = 0;
self.free_list.push(idx);
}
pub fn try_advance(&mut self, slot: u16, bytes_sent: u32) -> Option<(*const u8, u32)> {
let i = slot as usize;
debug_assert!(self.in_use[i]);
debug_assert!(bytes_sent <= self.slot_remaining[i]);
let new_remaining = self.slot_remaining[i] - bytes_sent;
if new_remaining == 0 {
return None;
}
self.slot_offset[i] += bytes_sent;
self.slot_remaining[i] = new_remaining;
let base = i * self.slot_size as usize;
let ptr = self
.backing
.as_ptr()
.wrapping_add(base + self.slot_offset[i] as usize);
Some((ptr, new_remaining))
}
pub fn original_len(&self, slot: u16) -> u32 {
let i = slot as usize;
debug_assert!(self.in_use[i]);
self.slot_offset[i] + self.slot_remaining[i]
}
pub fn slot_size(&self) -> u32 {
self.slot_size
}
#[allow(dead_code)]
pub fn free_count(&self) -> usize {
self.free_list.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn copy_in_and_release() {
let mut pool = SendCopyPool::new(4, 128);
assert_eq!(pool.free_count(), 4);
let (idx, ptr, len) = pool.copy_in(b"hello").unwrap();
assert_eq!(len, 5);
assert_eq!(pool.free_count(), 3);
let slice = unsafe { std::slice::from_raw_parts(ptr, len as usize) };
assert_eq!(slice, b"hello");
pool.release(idx);
assert_eq!(pool.free_count(), 4);
}
#[test]
fn exhaust_pool() {
let mut pool = SendCopyPool::new(2, 64);
let _ = pool.copy_in(b"a").unwrap();
let _ = pool.copy_in(b"b").unwrap();
assert!(pool.copy_in(b"c").is_none());
}
#[test]
fn data_too_large() {
let mut pool = SendCopyPool::new(4, 4);
assert!(pool.copy_in(b"toolarge").is_none());
}
#[test]
fn try_advance_partial() {
let mut pool = SendCopyPool::new(4, 128);
let (idx, _ptr, len) = pool.copy_in(b"hello world").unwrap();
assert_eq!(len, 11);
assert_eq!(pool.original_len(idx), 11);
let result = pool.try_advance(idx, 5);
assert!(result.is_some());
let (new_ptr, new_remaining) = result.unwrap();
assert_eq!(new_remaining, 6);
assert_eq!(pool.original_len(idx), 11);
let slice = unsafe { std::slice::from_raw_parts(new_ptr, new_remaining as usize) };
assert_eq!(slice, b" world");
let result = pool.try_advance(idx, 4);
assert!(result.is_some());
let (new_ptr2, new_remaining2) = result.unwrap();
assert_eq!(new_remaining2, 2);
assert_eq!(pool.original_len(idx), 11);
let slice2 = unsafe { std::slice::from_raw_parts(new_ptr2, new_remaining2 as usize) };
assert_eq!(slice2, b"ld");
let result = pool.try_advance(idx, 2);
assert!(result.is_none());
pool.release(idx);
assert_eq!(pool.free_count(), 4);
}
#[test]
fn try_advance_full_send() {
let mut pool = SendCopyPool::new(4, 128);
let (idx, _ptr, len) = pool.copy_in(b"hello").unwrap();
assert_eq!(len, 5);
let result = pool.try_advance(idx, 5);
assert!(result.is_none());
assert_eq!(pool.original_len(idx), 5);
pool.release(idx);
}
#[test]
fn release_clears_tracking() {
let mut pool = SendCopyPool::new(4, 128);
let (idx, _ptr, _len) = pool.copy_in(b"test data").unwrap();
pool.try_advance(idx, 4);
assert_eq!(pool.original_len(idx), 9);
pool.release(idx);
let (idx2, _ptr2, len2) = pool.copy_in(b"new").unwrap();
assert_eq!(idx2, idx);
assert_eq!(len2, 3);
assert_eq!(pool.original_len(idx2), 3);
}
}