use super::{checked_compact_strides, PooledUninitOutput};
use crate::buffer_pool::BufferPool;
fn recycled(pool: &mut BufferPool) -> tenferro_tensor::TypedTensor<f64> {
let mut output = PooledUninitOutput::<f64>::new(pool, vec![4]).unwrap();
for (i, value) in output.as_uninit_slice_mut().iter_mut().enumerate() {
value.write(i as f64);
}
unsafe { output.assume_init_recycled() }.unwrap()
}
#[test]
fn recycled_output_returns_only_after_final_group_owner_drops() {
let mut pool = BufferPool::new();
let tensor = recycled(&mut pool);
let pointer = tensor.as_slice().unwrap().as_ptr();
let (group, slots) = tenferro_tensor::AllocationGroup::from_tensors(vec![
tenferro_tensor::Tensor::from_typed::<f64>(tensor),
])
.unwrap();
assert!(pool.is_empty());
{
let reads = group.read_views(&slots).unwrap();
assert_eq!(reads[0].shape(), &[4]);
assert!(pool.is_empty());
}
drop(group);
assert_eq!(pool.len(), 1);
let tensor = recycled(&mut pool);
assert_eq!(tensor.as_slice().unwrap().as_ptr(), pointer);
drop(tensor);
assert_eq!(pool.len(), 1);
}
#[test]
fn recycled_drop_during_another_checkout_does_not_reenter_the_execution_lock() {
let mut pool = BufferPool::new();
let tensor = recycled(&mut pool);
let output = PooledUninitOutput::<f64>::new(&pool, vec![4]).unwrap();
drop(tensor);
drop(output);
assert_eq!(pool.len(), 2);
}
#[test]
fn output_and_scratch_checkouts_coexist_without_borrowing_the_pool() {
let mut pool = BufferPool::new();
let mut output = PooledUninitOutput::<f64>::new(&pool, vec![2]).unwrap();
let mut scratch = PooledUninitOutput::<f64>::new(&pool, vec![2]).unwrap();
output.as_uninit_slice_mut()[0].write(2.0);
output.as_uninit_slice_mut()[1].write(3.0);
scratch.as_uninit_slice_mut()[0].write(5.0);
scratch.as_uninit_slice_mut()[1].write(7.0);
let zeroed = pool.acquire_zeroed::<f64>(1);
assert_eq!(zeroed, vec![0.0]);
let result = unsafe { output.assume_init_recycled() }.unwrap();
let temporary = unsafe { scratch.assume_init_recycled() }.unwrap();
assert_eq!(result.as_slice().unwrap(), &[2.0, 3.0]);
assert_eq!(temporary.as_slice().unwrap(), &[5.0, 7.0]);
drop(result);
drop(temporary);
assert_eq!(pool.len(), 2);
}
#[test]
fn ranked_recycled_handoff_returns_once_and_rank_error_cancels_checkout() {
use tenferro_tensor::Rank;
let pool = BufferPool::new();
let mut wrong = PooledUninitOutput::<f64>::new(&pool, vec![2]).unwrap();
wrong.as_uninit_slice_mut()[0].write(1.0);
wrong.as_uninit_slice_mut()[1].write(2.0);
assert!(unsafe { wrong.assume_init_as_recycled::<Rank<2>>() }.is_err());
assert_eq!(pool.len(), 1);
let mut output = PooledUninitOutput::<f64>::new(&pool, vec![2]).unwrap();
output.as_uninit_slice_mut()[0].write(3.0);
output.as_uninit_slice_mut()[1].write(4.0);
let owner = unsafe { output.assume_init_as_recycled::<Rank<1>>() }.unwrap();
assert_eq!(owner.as_slice().unwrap(), &[3.0, 4.0]);
assert!(pool.is_empty());
drop(owner);
assert_eq!(pool.len(), 1);
}
#[test]
fn recycled_output_extraction_disarms_return_and_pool_may_drop_first() {
let mut pool = BufferPool::new();
let tensor = recycled(&mut pool);
let pointer = tensor.as_slice().unwrap().as_ptr();
let data = tensor.into_host_vec().unwrap();
assert_eq!(data.as_ptr(), pointer);
assert!(pool.is_empty());
drop(data);
assert!(pool.is_empty());
let tensor = recycled(&mut pool);
drop(pool);
assert_eq!(tensor.as_slice().unwrap(), &[0., 1., 2., 3.]);
drop(tensor);
}
#[test]
fn partial_output_returns_to_the_pool_on_error_and_unwind() {
let mut pool = BufferPool::new();
for unwind in [false, true] {
drop(recycled(&mut pool));
assert_eq!(pool.len(), 1);
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let mut output = PooledUninitOutput::<f64>::new(&pool, vec![4]).unwrap();
output.as_uninit_slice_mut()[0].write(f64::NAN);
if unwind {
panic!("kernel stopped after a partial write");
}
Err::<(), _>("kernel returned an error after a partial write")
}));
if unwind {
assert!(outcome.is_err());
} else {
assert!(outcome.unwrap().is_err());
}
assert_eq!(
pool.len(),
1,
"an aborted checkout returns its buffer to the pool"
);
assert_eq!(
pool.retained_capacity_bytes(),
4 * std::mem::size_of::<f64>()
);
let tensor = recycled(&mut pool);
assert_eq!(tensor.as_slice().unwrap(), &[0., 1., 2., 3.]);
drop(tensor);
pool.clear();
}
}
#[test]
fn recycled_output_respects_current_limit_and_clear() {
let mut pool = BufferPool::new();
let tensor = recycled(&mut pool);
pool.set_max_retained_capacity_bytes(0);
drop(tensor);
assert!(pool.is_empty());
pool.set_max_retained_capacity_bytes(1024);
drop(recycled(&mut pool));
assert_eq!(pool.len(), 1);
pool.clear();
assert!(pool.is_empty());
assert_eq!(pool.retained_capacity_bytes(), 0);
}
#[test]
fn compact_stride_validation_reports_dimension_and_product_overflow() {
let dimension_error = checked_compact_strides(&[isize::MAX as usize + 1]).unwrap_err();
assert!(dimension_error
.to_string()
.contains("dimension exceeds isize"));
let stride_error = checked_compact_strides(&[isize::MAX as usize, 2]).unwrap_err();
assert!(stride_error.to_string().contains("compact stride overflow"));
}
#[test]
fn pooled_uninit_output_public_contract_covers_zero_length_handoff() {
let pool = BufferPool::new();
let mut output = PooledUninitOutput::<i32>::new(&pool, vec![0]).unwrap();
assert!(output.as_uninit_slice_mut().is_empty());
assert!(output.as_uninit_bytes_mut().is_empty());
assert_eq!(output.as_uninit_view_mut().unwrap().dims(), &[0]);
let tensor = unsafe { output.assume_init_recycled() }.unwrap();
assert_eq!(tensor.shape(), &[0]);
assert!(pool.is_empty());
}