1use std::sync::Arc;
2use std::sync::atomic::{AtomicI32, AtomicU32, AtomicU64, Ordering};
3
4use parking_lot::Mutex;
5
6use crate::error::{ADError, ADResult};
7use crate::ndarray::{NDArray, NDDataBuffer, NDDataType, NDDimension};
8use crate::ndarray_handle::{NDArrayHandle, pooled_array};
9use crate::timestamp::EpicsTimestamp;
10
11const THRESHOLD_SIZE_RATIO: f64 = 1.5;
14
15static NEXT_POOL_ID: AtomicU64 = AtomicU64::new(1);
19
20trait BinAcc: Copy {
37 const ZERO: Self;
38 fn bin_add(self, rhs: Self) -> Self;
39}
40
41impl BinAcc for i128 {
42 const ZERO: Self = 0;
43 #[inline]
44 fn bin_add(self, rhs: Self) -> Self {
45 self.wrapping_add(rhs)
46 }
47}
48
49impl BinAcc for f32 {
50 const ZERO: Self = 0.0;
51 #[inline]
52 fn bin_add(self, rhs: Self) -> Self {
53 self + rhs
54 }
55}
56
57impl BinAcc for f64 {
58 const ZERO: Self = 0.0;
59 #[inline]
60 fn bin_add(self, rhs: Self) -> Self {
61 self + rhs
62 }
63}
64
65pub struct NDArrayPool {
72 id: u64,
74 max_memory: usize,
75 allocated_bytes: AtomicU64,
76 next_unique_id: AtomicI32,
77 free_list: Mutex<Vec<NDArray>>,
78 num_alloc_buffers: AtomicU32,
79 num_free_buffers: AtomicU32,
80}
81
82impl NDArrayPool {
83 pub fn new(max_memory: usize) -> Self {
84 Self {
85 id: NEXT_POOL_ID.fetch_add(1, Ordering::Relaxed),
86 max_memory,
87 allocated_bytes: AtomicU64::new(0),
88 next_unique_id: AtomicI32::new(1),
89 free_list: Mutex::new(Vec::new()),
90 num_alloc_buffers: AtomicU32::new(0),
91 num_free_buffers: AtomicU32::new(0),
92 }
93 }
94
95 pub fn id(&self) -> u64 {
97 self.id
98 }
99
100 pub fn alloc(&self, dims: Vec<NDDimension>, data_type: NDDataType) -> ADResult<NDArray> {
108 let num_elements: usize = dims.iter().map(|d| d.size).product();
109 let needed_bytes = num_elements * data_type.element_size();
110
111 let reused = {
115 let mut free = self.free_list.lock();
116 let mut best_idx = None;
117 let mut best_cap = usize::MAX;
118 for (i, arr) in free.iter().enumerate() {
119 let cap = arr.data.capacity_bytes();
120 if cap >= needed_bytes && cap < best_cap {
121 best_cap = cap;
122 best_idx = Some(i);
123 }
124 }
125 if let Some(idx) = best_idx {
126 if best_cap as f64 > needed_bytes as f64 * THRESHOLD_SIZE_RATIO {
127 let dropped = free.swap_remove(idx);
129 self.num_free_buffers.fetch_sub(1, Ordering::Relaxed);
130 self.allocated_bytes
131 .fetch_sub(dropped.data_size as u64, Ordering::Relaxed);
132 self.num_alloc_buffers.fetch_sub(1, Ordering::Relaxed);
133 None
134 } else {
135 let arr = free.swap_remove(idx);
136 self.num_free_buffers.fetch_sub(1, Ordering::Relaxed);
137 Some(arr)
138 }
139 } else {
140 None
141 }
142 };
143
144 let mut arr = if let Some(mut reused) = reused {
145 let old_size = reused.data_size;
151 if reused.data.data_type() != data_type {
152 reused.data = NDDataBuffer::zeros(data_type, num_elements);
153 } else {
154 reused.data.resize(num_elements);
155 }
156 let effective_size = if needed_bytes > old_size {
157 let diff = (needed_bytes - old_size) as u64;
158 if self.max_memory > 0 {
163 loop {
164 let current = self.allocated_bytes.load(Ordering::Relaxed);
165 if current + diff > self.max_memory as u64 {
166 let mut free = self.free_list.lock();
169 free.push(reused);
170 self.num_free_buffers.fetch_add(1, Ordering::Relaxed);
171 return Err(ADError::PoolExhausted(needed_bytes, self.max_memory));
172 }
173 if self
174 .allocated_bytes
175 .compare_exchange_weak(
176 current,
177 current + diff,
178 Ordering::Relaxed,
179 Ordering::Relaxed,
180 )
181 .is_ok()
182 {
183 break;
184 }
185 }
186 } else {
187 self.allocated_bytes.fetch_add(diff, Ordering::Relaxed);
188 }
189 needed_bytes
190 } else {
191 old_size
194 };
195 reused.data_size = effective_size;
196 reused.dims = dims;
197 reused.attributes.clear();
198 reused.codec = None;
199 reused
200 } else {
201 if self.max_memory > 0 {
204 loop {
205 let current = self.allocated_bytes.load(Ordering::Relaxed);
206 if current + needed_bytes as u64 > self.max_memory as u64 {
207 let mut freed_enough = false;
208 {
209 let mut free = self.free_list.lock();
210 free.sort_by(|a, b| {
211 b.data.capacity_bytes().cmp(&a.data.capacity_bytes())
212 });
213 let mut reclaimed = 0u64;
214 let over = (current + needed_bytes as u64)
215 .saturating_sub(self.max_memory as u64);
216 while !free.is_empty() && reclaimed < over {
217 let dropped = free.remove(0);
218 let dropped_size = dropped.data_size as u64;
219 self.allocated_bytes
220 .fetch_sub(dropped_size, Ordering::Relaxed);
221 self.num_free_buffers.fetch_sub(1, Ordering::Relaxed);
222 self.num_alloc_buffers.fetch_sub(1, Ordering::Relaxed);
223 reclaimed += dropped_size;
224 }
225 if reclaimed >= over {
226 freed_enough = true;
227 }
228 }
229 if !freed_enough {
230 return Err(ADError::PoolExhausted(needed_bytes, self.max_memory));
231 }
232 continue;
233 }
234 if self
235 .allocated_bytes
236 .compare_exchange_weak(
237 current,
238 current + needed_bytes as u64,
239 Ordering::Relaxed,
240 Ordering::Relaxed,
241 )
242 .is_ok()
243 {
244 break;
245 }
246 }
247 } else {
248 self.allocated_bytes
249 .fetch_add(needed_bytes as u64, Ordering::Relaxed);
250 }
251 self.num_alloc_buffers.fetch_add(1, Ordering::Relaxed);
252 NDArray::new(dims, data_type)
253 };
254
255 arr.unique_id = self.next_unique_id.fetch_add(1, Ordering::Relaxed);
256 arr.timestamp = EpicsTimestamp::now();
257 arr.pool_id = self.id;
258 Ok(arr)
261 }
262
263 pub fn alloc_copy(&self, source: &NDArray) -> ADResult<NDArray> {
266 let dims = source.dims.clone();
267 let data_type = source.data.data_type();
268 let mut copy = self.alloc(dims, data_type)?;
269 copy.data = source.data.clone();
270 copy.time_stamp = source.time_stamp;
271 copy.attributes = source.attributes.clone();
272 copy.codec = source.codec.clone();
273 Ok(copy)
274 }
275
276 pub fn release(&self, array: NDArray) {
283 if array.pool_id != self.id {
284 return;
288 }
289
290 let mut free = self.free_list.lock();
291 free.push(array);
292 self.num_free_buffers.fetch_add(1, Ordering::Relaxed);
293
294 let total = self.allocated_bytes.load(Ordering::Relaxed) as usize;
298 if self.max_memory > 0 && total > self.max_memory && !free.is_empty() {
299 free.sort_by(|a, b| b.data_size.cmp(&a.data_size));
300 let mut excess = total - self.max_memory;
301 while excess > 0 && !free.is_empty() {
302 let dropped = free.remove(0);
303 let dropped_size = dropped.data_size;
304 self.allocated_bytes
305 .fetch_sub(dropped_size as u64, Ordering::Relaxed);
306 self.num_free_buffers.fetch_sub(1, Ordering::Relaxed);
307 self.num_alloc_buffers.fetch_sub(1, Ordering::Relaxed);
308 if dropped_size >= excess {
309 break;
310 }
311 excess -= dropped_size;
312 }
313 }
314 }
315
316 pub fn empty_free_list(&self) {
318 let mut free = self.free_list.lock();
319 let count = free.len() as u32;
320 for arr in free.drain(..) {
321 self.allocated_bytes
322 .fetch_sub(arr.data_size as u64, Ordering::Relaxed);
323 self.num_alloc_buffers.fetch_sub(1, Ordering::Relaxed);
324 }
325 self.num_free_buffers.fetch_sub(count, Ordering::Relaxed);
326 }
327
328 pub fn allocated_bytes(&self) -> u64 {
329 self.allocated_bytes.load(Ordering::Relaxed)
330 }
331
332 pub fn num_alloc_buffers(&self) -> u32 {
333 self.num_alloc_buffers.load(Ordering::Relaxed)
334 }
335
336 pub fn num_free_buffers(&self) -> u32 {
337 self.num_free_buffers.load(Ordering::Relaxed)
338 }
339
340 pub fn max_memory(&self) -> usize {
341 self.max_memory
342 }
343
344 pub fn alloc_handle(
347 pool: &Arc<Self>,
348 dims: Vec<NDDimension>,
349 data_type: NDDataType,
350 ) -> ADResult<NDArrayHandle> {
351 let array = pool.alloc(dims, data_type)?;
352 Ok(pooled_array(array, pool))
353 }
354
355 pub fn copy(
367 &self,
368 src: &NDArray,
369 out: Option<NDArray>,
370 copy_data: bool,
371 copy_dimensions: bool,
372 copy_data_type: bool,
373 ) -> ADResult<NDArray> {
374 let mut out = match out {
375 Some(o) => o,
376 None => self.alloc(src.dims.clone(), src.data.data_type())?,
377 };
378
379 out.unique_id = src.unique_id;
380 out.time_stamp = src.time_stamp;
381 out.timestamp = src.timestamp;
382 if copy_dimensions {
383 out.dims = src.dims.clone();
384 }
385 out.codec = src.codec.clone();
386
387 if copy_data {
388 if copy_data_type && out.data.data_type() != src.data.data_type() {
389 out.data = src.data.clone();
391 } else if out.data.data_type() == src.data.data_type() {
392 out.data = src.data.clone();
393 } else {
394 out.data = crate::color::convert_data_type(src, out.data.data_type())?.data;
396 }
397 } else if copy_data_type && out.data.data_type() != src.data.data_type() {
398 out.data = NDDataBuffer::zeros(src.data.data_type(), out.data.len());
399 }
400
401 out.attributes.clear();
402 out.attributes.copy_from(&src.attributes);
403 Ok(out)
404 }
405
406 pub fn pre_allocate_buffers(&self, template_array: &NDArray, count: usize) -> ADResult<()> {
412 let mut buffers = Vec::with_capacity(count);
413 for _ in 0..count {
414 buffers.push(self.copy(template_array, None, true, true, true)?);
415 }
416 for arr in buffers {
417 self.release(arr);
418 }
419 Ok(())
420 }
421
422 pub fn convert_type(&self, src: &NDArray, target_type: NDDataType) -> ADResult<NDArray> {
425 if src.codec.is_some() {
427 return Err(ADError::UnsupportedConversion(
428 "convert_type: cannot convert compressed (codec) data".into(),
429 ));
430 }
431 if src.data.data_type() == target_type {
432 return self.alloc_copy(src);
433 }
434 let mut out = self.alloc(src.dims.clone(), target_type)?;
437 let converted = crate::color::convert_data_type(src, target_type)?;
438 out.data = converted.data;
439 out.time_stamp = src.time_stamp;
440 out.timestamp = src.timestamp;
441 out.attributes.copy_from(&src.attributes);
442 Ok(out)
443 }
444
445 pub fn convert(
456 &self,
457 src: &NDArray,
458 dims_out: &[NDDimension],
459 target_type: NDDataType,
460 ) -> ADResult<NDArray> {
461 if src.codec.is_some() {
463 return Err(ADError::UnsupportedConversion(
464 "convert: cannot convert compressed (codec) data".into(),
465 ));
466 }
467
468 let ndims = src.dims.len();
469 if dims_out.len() != ndims {
470 return Err(ADError::InvalidDimensions(format!(
471 "convert: dims_out length {} != source ndims {}",
472 dims_out.len(),
473 ndims,
474 )));
475 }
476
477 let mut out_sizes = Vec::with_capacity(ndims);
479 for (i, d) in dims_out.iter().enumerate() {
480 let bin = d.binning.max(1);
481 if d.size == 0 {
482 return Err(ADError::InvalidDimensions(format!(
483 "convert: dims_out[{}].size is 0",
484 i,
485 )));
486 }
487 let out_size = d.size / bin;
488 if out_size == 0 {
489 return Err(ADError::InvalidDimensions(format!(
490 "convert: dims_out[{}] size {} / binning {} = 0",
491 i, d.size, bin,
492 )));
493 }
494 if d.offset + d.size > src.dims[i].size {
496 return Err(ADError::InvalidDimensions(format!(
497 "convert: dims_out[{}] offset {} + size {} > src dim size {}",
498 i, d.offset, d.size, src.dims[i].size,
499 )));
500 }
501 out_sizes.push(out_size);
502 }
503
504 let mut out_dims = Vec::with_capacity(ndims);
509 for i in 0..ndims {
510 let bin = dims_out[i].binning.max(1);
511 out_dims.push(NDDimension {
512 size: out_sizes[i],
513 offset: src.dims[i].offset + dims_out[i].offset,
514 binning: src.dims[i].binning * bin,
515 reverse: dims_out[i].reverse ^ src.dims[i].reverse,
516 });
517 }
518
519 let total_out: usize = out_sizes.iter().product();
520
521 let mut src_strides = vec![1usize; ndims];
523 for i in 1..ndims {
524 src_strides[i] = src_strides[i - 1] * src.dims[i - 1].size;
525 }
526
527 let mut out_strides = vec![1usize; ndims];
529 for i in 1..ndims {
530 out_strides[i] = out_strides[i - 1] * out_sizes[i - 1];
531 }
532
533 macro_rules! bin_loop {
541 ($src_vec:expr, $DstT:ty, $AccT:ty, $variant:ident) => {{
542 let mut out = vec![0 as $DstT; total_out];
543
544 for out_idx in 0..total_out {
546 let mut remaining = out_idx;
548 let mut out_coords = [0usize; 10]; for i in (0..ndims).rev() {
550 out_coords[i] = remaining / out_strides[i];
551 remaining %= out_strides[i];
552 }
553
554 let mut eff_coords = [0usize; 10];
556 for i in 0..ndims {
557 eff_coords[i] = if dims_out[i].reverse {
558 out_sizes[i] - 1 - out_coords[i]
559 } else {
560 out_coords[i]
561 };
562 }
563
564 let mut acc = <$AccT as BinAcc>::ZERO;
566 let bin_total: usize = dims_out.iter().map(|d| d.binning.max(1)).product();
567
568 for bin_flat in 0..bin_total {
570 let mut br = bin_flat;
571 let mut src_flat = 0usize;
572 let mut valid = true;
573
574 for i in (0..ndims).rev() {
575 let bin = dims_out[i].binning.max(1);
576 let bin_off = br % bin;
577 br /= bin;
578
579 let src_coord = dims_out[i].offset + eff_coords[i] * bin + bin_off;
580 if src_coord >= src.dims[i].size {
581 valid = false;
582 break;
583 }
584 src_flat += src_coord * src_strides[i];
585 }
586
587 if valid {
588 acc = acc.bin_add($src_vec[src_flat] as $AccT);
591 }
592 }
593
594 out[out_idx] = acc as $DstT;
595 }
596
597 NDDataBuffer::$variant(out)
598 }};
599 }
600
601 macro_rules! bin_to_target {
605 ($src_vec:expr) => {
606 match target_type {
607 NDDataType::Int8 => bin_loop!($src_vec, i8, i128, I8),
608 NDDataType::UInt8 => bin_loop!($src_vec, u8, i128, U8),
609 NDDataType::Int16 => bin_loop!($src_vec, i16, i128, I16),
610 NDDataType::UInt16 => bin_loop!($src_vec, u16, i128, U16),
611 NDDataType::Int32 => bin_loop!($src_vec, i32, i128, I32),
612 NDDataType::UInt32 => bin_loop!($src_vec, u32, i128, U32),
613 NDDataType::Int64 => bin_loop!($src_vec, i64, i128, I64),
614 NDDataType::UInt64 => bin_loop!($src_vec, u64, i128, U64),
615 NDDataType::Float32 => bin_loop!($src_vec, f32, f32, F32),
616 NDDataType::Float64 => bin_loop!($src_vec, f64, f64, F64),
617 }
618 };
619 }
620
621 let out_data = match &src.data {
622 NDDataBuffer::I8(v) => bin_to_target!(v),
623 NDDataBuffer::U8(v) => bin_to_target!(v),
624 NDDataBuffer::I16(v) => bin_to_target!(v),
625 NDDataBuffer::U16(v) => bin_to_target!(v),
626 NDDataBuffer::I32(v) => bin_to_target!(v),
627 NDDataBuffer::U32(v) => bin_to_target!(v),
628 NDDataBuffer::I64(v) => bin_to_target!(v),
629 NDDataBuffer::U64(v) => bin_to_target!(v),
630 NDDataBuffer::F32(v) => bin_to_target!(v),
631 NDDataBuffer::F64(v) => bin_to_target!(v),
632 };
633
634 let mut arr = self.alloc(out_dims, target_type)?;
638 arr.timestamp = src.timestamp;
639 arr.time_stamp = src.time_stamp;
640 arr.attributes.copy_from(&src.attributes);
641
642 arr.data = out_data;
646
647 Ok(arr)
648 }
649
650 pub fn report(&self, details: i32) -> String {
654 let mut out = String::new();
655 out.push('\n');
656 out.push_str("NDArrayPool:\n");
657 out.push_str(&format!(
658 " numBuffers={}, numFree={}\n",
659 self.num_alloc_buffers(),
660 self.num_free_buffers()
661 ));
662 out.push_str(&format!(
663 " memorySize={}, maxMemory={}\n",
664 self.allocated_bytes(),
665 self.max_memory
666 ));
667 if details > 5 {
668 let free = self.free_list.lock();
669 out.push_str(" freeList: (index, dataSize, capacity)\n");
670 for (i, arr) in free.iter().enumerate() {
671 out.push_str(&format!(
672 " {} {} {}\n",
673 i,
674 arr.data_size,
675 arr.data.capacity_bytes()
676 ));
677 }
678 if details > 10 {
679 for arr in free.iter() {
680 out.push_str(&arr.report(details));
681 }
682 }
683 }
684 out
685 }
686}
687
688const _: fn() = || {
690 fn assert_send_sync<T: Send + Sync>() {}
691 assert_send_sync::<NDArrayPool>();
692};
693
694#[cfg(test)]
695mod tests {
696 use super::*;
697
698 #[test]
699 fn test_alloc_auto_id() {
700 let pool = NDArrayPool::new(1_000_000);
701 let a1 = pool
702 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
703 .unwrap();
704 let a2 = pool
705 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
706 .unwrap();
707 assert_eq!(a1.unique_id, 1);
708 assert_eq!(a2.unique_id, 2);
709 }
710
711 #[test]
712 fn test_alloc_tracks_bytes() {
713 let pool = NDArrayPool::new(1_000_000);
714 let _ = pool
715 .alloc(vec![NDDimension::new(100)], NDDataType::Float64)
716 .unwrap();
717 assert!(pool.allocated_bytes() >= 800);
718 }
719
720 #[test]
721 fn test_alloc_exceeds_max() {
722 let pool = NDArrayPool::new(100);
723 let result = pool.alloc(vec![NDDimension::new(200)], NDDataType::UInt8);
724 assert!(result.is_err());
725 }
726
727 #[test]
728 fn test_alloc_copy_preserves_data() {
729 let pool = NDArrayPool::new(1_000_000);
730 let mut source = pool
731 .alloc(vec![NDDimension::new(4)], NDDataType::UInt8)
732 .unwrap();
733 if let NDDataBuffer::U8(ref mut v) = source.data {
734 v[0] = 1;
735 v[1] = 2;
736 v[2] = 3;
737 v[3] = 4;
738 }
739
740 let copy = pool.alloc_copy(&source).unwrap();
741 assert_ne!(copy.unique_id, source.unique_id);
742 assert_eq!(copy.dims.len(), source.dims.len());
743 if let NDDataBuffer::U8(ref v) = copy.data {
744 assert_eq!(v, &[1, 2, 3, 4]);
745 } else {
746 panic!("wrong type");
747 }
748 }
749
750 #[test]
751 fn test_alloc_copy_tracks_bytes() {
752 let pool = NDArrayPool::new(1_000_000);
753 let source = pool
754 .alloc(vec![NDDimension::new(10)], NDDataType::UInt16)
755 .unwrap();
756 assert_eq!(pool.allocated_bytes(), 20);
757 let _ = pool.alloc_copy(&source).unwrap();
758 assert!(pool.allocated_bytes() >= 40);
759 }
760
761 #[test]
762 fn test_alloc_copy_exceeds_max() {
763 let pool = NDArrayPool::new(60);
764 let source = pool
765 .alloc(vec![NDDimension::new(50)], NDDataType::UInt8)
766 .unwrap();
767 assert!(pool.alloc_copy(&source).is_err());
768 }
769
770 #[test]
773 fn test_release_and_reuse() {
774 let pool = NDArrayPool::new(1_000_000);
775 let arr = pool
776 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
777 .unwrap();
778 let _alloc_bytes_after_first = pool.allocated_bytes();
779 assert_eq!(pool.num_alloc_buffers(), 1);
780
781 pool.release(arr);
783 assert_eq!(pool.num_free_buffers(), 1);
784
785 let arr2 = pool
787 .alloc(vec![NDDimension::new(80)], NDDataType::UInt8)
788 .unwrap();
789 assert_eq!(arr2.data.len(), 80);
790 }
791
792 #[test]
793 fn test_free_list_prefers_smallest_sufficient() {
794 let pool = NDArrayPool::new(10_000_000);
795 let small = pool
796 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
797 .unwrap();
798 let large = pool
799 .alloc(vec![NDDimension::new(10000)], NDDataType::UInt8)
800 .unwrap();
801 let medium = pool
802 .alloc(vec![NDDimension::new(1000)], NDDataType::UInt8)
803 .unwrap();
804
805 pool.release(large);
806 pool.release(medium);
807 pool.release(small);
808 assert_eq!(pool.num_free_buffers(), 3);
809
810 let reused = pool
812 .alloc(vec![NDDimension::new(900)], NDDataType::UInt8)
813 .unwrap();
814 assert!(reused.data.capacity_bytes() >= 900);
815 }
816
817 #[test]
818 fn test_empty_free_list() {
819 let pool = NDArrayPool::new(1_000_000);
820 let a1 = pool
821 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
822 .unwrap();
823 let a2 = pool
824 .alloc(vec![NDDimension::new(200)], NDDataType::UInt8)
825 .unwrap();
826 pool.release(a1);
827 pool.release(a2);
828 assert_eq!(pool.num_free_buffers(), 2);
829
830 pool.empty_free_list();
831 assert_eq!(pool.num_free_buffers(), 0);
832 assert_eq!(pool.num_alloc_buffers(), 0);
833 }
834
835 #[test]
836 fn test_num_free_buffers_tracking() {
837 let pool = NDArrayPool::new(1_000_000);
838 assert_eq!(pool.num_free_buffers(), 0);
839
840 let a = pool
841 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
842 .unwrap();
843 assert_eq!(pool.num_free_buffers(), 0);
844
845 pool.release(a);
846 assert_eq!(pool.num_free_buffers(), 1);
847
848 let _ = pool
849 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
850 .unwrap();
851 assert_eq!(pool.num_free_buffers(), 0);
852 }
853
854 #[test]
855 fn test_concurrent_alloc_release() {
856 use std::sync::Arc;
857 use std::thread;
858
859 let pool = Arc::new(NDArrayPool::new(10_000_000));
860 let mut handles = Vec::new();
861
862 for _ in 0..4 {
863 let pool = pool.clone();
864 handles.push(thread::spawn(move || {
865 for _ in 0..100 {
866 let arr = pool
867 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
868 .unwrap();
869 pool.release(arr);
870 }
871 }));
872 }
873
874 for h in handles {
875 h.join().unwrap();
876 }
877
878 assert!(pool.num_free_buffers() > 0);
880 }
881
882 #[test]
883 fn test_max_memory() {
884 let pool = NDArrayPool::new(42);
885 assert_eq!(pool.max_memory(), 42);
886 }
887
888 #[test]
891 fn test_convert_type_same_type() {
892 let pool = NDArrayPool::new(1_000_000);
893 let mut src = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
894 if let NDDataBuffer::U8(ref mut v) = src.data {
895 v[0] = 10;
896 v[1] = 20;
897 v[2] = 30;
898 v[3] = 40;
899 }
900
901 let out = pool.convert_type(&src, NDDataType::UInt8).unwrap();
902 assert_eq!(out.data.data_type(), NDDataType::UInt8);
903 if let NDDataBuffer::U8(ref v) = out.data {
904 assert_eq!(v, &[10, 20, 30, 40]);
905 } else {
906 panic!("wrong type");
907 }
908 }
909
910 #[test]
911 fn test_convert_type_u8_to_f32() {
912 let pool = NDArrayPool::new(1_000_000);
913 let mut src = NDArray::new(vec![NDDimension::new(3)], NDDataType::UInt8);
914 if let NDDataBuffer::U8(ref mut v) = src.data {
915 v[0] = 0;
916 v[1] = 128;
917 v[2] = 255;
918 }
919
920 let out = pool.convert_type(&src, NDDataType::Float32).unwrap();
921 assert_eq!(out.data.data_type(), NDDataType::Float32);
922 if let NDDataBuffer::F32(ref v) = out.data {
923 assert_eq!(v[0], 0.0);
924 assert_eq!(v[1], 128.0);
925 assert_eq!(v[2], 255.0);
926 } else {
927 panic!("wrong type");
928 }
929 }
930
931 #[test]
932 fn test_convert_type_u16_to_u8() {
933 let pool = NDArrayPool::new(1_000_000);
934 let mut src = NDArray::new(vec![NDDimension::new(2)], NDDataType::UInt16);
935 if let NDDataBuffer::U16(ref mut v) = src.data {
936 v[0] = 100;
937 v[1] = 300; }
939
940 let out = pool.convert_type(&src, NDDataType::UInt8).unwrap();
941 if let NDDataBuffer::U8(ref v) = out.data {
942 assert_eq!(v[0], 100);
943 assert_eq!(v[1], 255); } else {
945 panic!("wrong type");
946 }
947 }
948
949 fn make_4x4_u8() -> NDArray {
953 let mut arr = NDArray::new(
954 vec![NDDimension::new(4), NDDimension::new(4)],
955 NDDataType::UInt8,
956 );
957 if let NDDataBuffer::U8(ref mut v) = arr.data {
958 for i in 0..16 {
959 v[i] = i as u8;
960 }
961 }
962 arr
963 }
964
965 #[test]
966 fn test_convert_identity() {
967 let pool = NDArrayPool::new(1_000_000);
969 let src = make_4x4_u8();
970 let dims_out = vec![
971 NDDimension {
972 size: 4,
973 offset: 0,
974 binning: 1,
975 reverse: false,
976 },
977 NDDimension {
978 size: 4,
979 offset: 0,
980 binning: 1,
981 reverse: false,
982 },
983 ];
984
985 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
986 assert_eq!(out.dims[0].size, 4);
987 assert_eq!(out.dims[1].size, 4);
988 if let NDDataBuffer::U8(ref v) = out.data {
989 for i in 0..16 {
990 assert_eq!(v[i], i as u8);
991 }
992 } else {
993 panic!("wrong type");
994 }
995 }
996
997 #[test]
998 fn test_convert_offset_extraction() {
999 let pool = NDArrayPool::new(1_000_000);
1001 let src = make_4x4_u8();
1002 let dims_out = vec![
1003 NDDimension {
1004 size: 2,
1005 offset: 1,
1006 binning: 1,
1007 reverse: false,
1008 },
1009 NDDimension {
1010 size: 2,
1011 offset: 1,
1012 binning: 1,
1013 reverse: false,
1014 },
1015 ];
1016
1017 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1018 assert_eq!(out.dims[0].size, 2);
1019 assert_eq!(out.dims[1].size, 2);
1020 if let NDDataBuffer::U8(ref v) = out.data {
1024 assert_eq!(v[0], 5);
1025 assert_eq!(v[1], 6);
1026 assert_eq!(v[2], 9);
1027 assert_eq!(v[3], 10);
1028 } else {
1029 panic!("wrong type");
1030 }
1031
1032 assert_eq!(out.dims[0].offset, 1); assert_eq!(out.dims[1].offset, 1);
1035 }
1036
1037 #[test]
1038 fn test_convert_binning_2x2() {
1039 let pool = NDArrayPool::new(1_000_000);
1041 let src = make_4x4_u8();
1042 let dims_out = vec![
1043 NDDimension {
1044 size: 4,
1045 offset: 0,
1046 binning: 2,
1047 reverse: false,
1048 },
1049 NDDimension {
1050 size: 4,
1051 offset: 0,
1052 binning: 2,
1053 reverse: false,
1054 },
1055 ];
1056
1057 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1058 assert_eq!(out.dims[0].size, 2);
1059 assert_eq!(out.dims[1].size, 2);
1060 if let NDDataBuffer::U8(ref v) = out.data {
1065 assert_eq!(v[0], 10);
1066 assert_eq!(v[1], 18);
1067 assert_eq!(v[2], 42);
1068 assert_eq!(v[3], 50);
1069 } else {
1070 panic!("wrong type");
1071 }
1072
1073 assert_eq!(out.dims[0].binning, 2); assert_eq!(out.dims[1].binning, 2);
1076 }
1077
1078 fn make_2x2_filled(data_type: NDDataType, val: i128) -> NDArray {
1081 let mut src = NDArray::new(vec![NDDimension::new(2), NDDimension::new(2)], data_type);
1082 macro_rules! fill {
1083 ($variant:ident, $t:ty) => {
1084 if let NDDataBuffer::$variant(ref mut v) = src.data {
1085 for e in v.iter_mut() {
1086 *e = val as $t;
1087 }
1088 }
1089 };
1090 }
1091 match data_type {
1092 NDDataType::Int8 => fill!(I8, i8),
1093 NDDataType::UInt8 => fill!(U8, u8),
1094 NDDataType::Int16 => fill!(I16, i16),
1095 NDDataType::UInt16 => fill!(U16, u16),
1096 NDDataType::Int32 => fill!(I32, i32),
1097 NDDataType::UInt32 => fill!(U32, u32),
1098 NDDataType::Int64 => fill!(I64, i64),
1099 NDDataType::UInt64 => fill!(U64, u64),
1100 NDDataType::Float32 => fill!(F32, f32),
1101 NDDataType::Float64 => fill!(F64, f64),
1102 }
1103 src
1104 }
1105
1106 fn bin_2x2_to_one() -> Vec<NDDimension> {
1107 vec![
1109 NDDimension {
1110 size: 2,
1111 offset: 0,
1112 binning: 2,
1113 reverse: false,
1114 },
1115 NDDimension {
1116 size: 2,
1117 offset: 0,
1118 binning: 2,
1119 reverse: false,
1120 },
1121 ]
1122 }
1123
1124 #[test]
1129 fn test_convert_binning_widening_target_keeps_full_sum() {
1130 let pool = NDArrayPool::new(1_000_000);
1131 let src = make_2x2_filled(NDDataType::UInt8, 200);
1132 let out = pool
1133 .convert(&src, &bin_2x2_to_one(), NDDataType::UInt16)
1134 .unwrap();
1135 match out.data {
1136 NDDataBuffer::U16(ref v) => assert_eq!(v[0], 800, "C sums in u16: 200*4"),
1137 ref other => panic!("expected U16, got {:?}", other.data_type()),
1138 }
1139 }
1140
1141 #[test]
1145 fn test_convert_binning_integer_overflow_wraps() {
1146 let pool = NDArrayPool::new(1_000_000);
1147 let src = make_2x2_filled(NDDataType::UInt8, 100);
1148 let out = pool
1149 .convert(&src, &bin_2x2_to_one(), NDDataType::UInt8)
1150 .unwrap();
1151 match out.data {
1152 NDDataBuffer::U8(ref v) => assert_eq!(v[0], 144, "400 wraps mod 256 in u8"),
1153 ref other => panic!("expected U8, got {:?}", other.data_type()),
1154 }
1155 }
1156
1157 #[test]
1162 fn test_convert_binning1_int64_above_2pow53_exact() {
1163 let pool = NDArrayPool::new(1_000_000);
1164 let value: i64 = (1i64 << 53) + 1;
1165 let mut src = NDArray::new(
1166 vec![NDDimension::new(1), NDDimension::new(1)],
1167 NDDataType::Int64,
1168 );
1169 if let NDDataBuffer::I64(ref mut v) = src.data {
1170 v[0] = value;
1171 }
1172 let dims_out = vec![
1173 NDDimension {
1174 size: 1,
1175 offset: 0,
1176 binning: 1,
1177 reverse: false,
1178 },
1179 NDDimension {
1180 size: 1,
1181 offset: 0,
1182 binning: 1,
1183 reverse: false,
1184 },
1185 ];
1186 let out = pool.convert(&src, &dims_out, NDDataType::Int64).unwrap();
1187 match out.data {
1188 NDDataBuffer::I64(ref v) => assert_eq!(v[0], value, "2^53+1 kept exactly"),
1189 ref other => panic!("expected I64, got {:?}", other.data_type()),
1190 }
1191 }
1192
1193 #[test]
1194 fn test_convert_reverse_x() {
1195 let pool = NDArrayPool::new(1_000_000);
1197 let mut src = NDArray::new(
1198 vec![NDDimension::new(4), NDDimension::new(1)],
1199 NDDataType::UInt8,
1200 );
1201 if let NDDataBuffer::U8(ref mut v) = src.data {
1202 v[0] = 10;
1203 v[1] = 20;
1204 v[2] = 30;
1205 v[3] = 40;
1206 }
1207
1208 let dims_out = vec![
1209 NDDimension {
1210 size: 4,
1211 offset: 0,
1212 binning: 1,
1213 reverse: true,
1214 },
1215 NDDimension {
1216 size: 1,
1217 offset: 0,
1218 binning: 1,
1219 reverse: false,
1220 },
1221 ];
1222
1223 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1224 if let NDDataBuffer::U8(ref v) = out.data {
1225 assert_eq!(v[0], 40);
1226 assert_eq!(v[1], 30);
1227 assert_eq!(v[2], 20);
1228 assert_eq!(v[3], 10);
1229 } else {
1230 panic!("wrong type");
1231 }
1232 }
1233
1234 #[test]
1235 fn test_convert_reverse_y() {
1236 let pool = NDArrayPool::new(1_000_000);
1238 let mut src = NDArray::new(
1239 vec![NDDimension::new(2), NDDimension::new(2)],
1240 NDDataType::UInt16,
1241 );
1242 if let NDDataBuffer::U16(ref mut v) = src.data {
1243 v[0] = 1;
1245 v[1] = 2;
1246 v[2] = 3;
1247 v[3] = 4;
1248 }
1249
1250 let dims_out = vec![
1251 NDDimension {
1252 size: 2,
1253 offset: 0,
1254 binning: 1,
1255 reverse: false,
1256 },
1257 NDDimension {
1258 size: 2,
1259 offset: 0,
1260 binning: 1,
1261 reverse: true,
1262 },
1263 ];
1264
1265 let out = pool.convert(&src, &dims_out, NDDataType::UInt16).unwrap();
1266 if let NDDataBuffer::U16(ref v) = out.data {
1267 assert_eq!(v[0], 3);
1269 assert_eq!(v[1], 4);
1270 assert_eq!(v[2], 1);
1271 assert_eq!(v[3], 2);
1272 } else {
1273 panic!("wrong type");
1274 }
1275 }
1276
1277 #[test]
1278 fn test_convert_with_type_change() {
1279 let pool = NDArrayPool::new(1_000_000);
1281 let src = make_4x4_u8();
1282 let dims_out = vec![
1283 NDDimension {
1284 size: 4,
1285 offset: 0,
1286 binning: 2,
1287 reverse: false,
1288 },
1289 NDDimension {
1290 size: 4,
1291 offset: 0,
1292 binning: 2,
1293 reverse: false,
1294 },
1295 ];
1296
1297 let out = pool.convert(&src, &dims_out, NDDataType::Float32).unwrap();
1298 assert_eq!(out.data.data_type(), NDDataType::Float32);
1299 assert_eq!(out.dims[0].size, 2);
1300 assert_eq!(out.dims[1].size, 2);
1301 if let NDDataBuffer::F32(ref v) = out.data {
1302 assert_eq!(v[0], 10.0); assert_eq!(v[1], 18.0); } else {
1305 panic!("wrong type");
1306 }
1307 }
1308
1309 #[test]
1310 fn test_convert_cumulative_offset_and_binning() {
1311 let pool = NDArrayPool::new(1_000_000);
1313 let mut src = NDArray::new(
1314 vec![NDDimension::new(4), NDDimension::new(4)],
1315 NDDataType::UInt8,
1316 );
1317 src.dims[0].offset = 10;
1318 src.dims[0].binning = 2;
1319 src.dims[1].offset = 20;
1320 src.dims[1].binning = 3;
1321 if let NDDataBuffer::U8(ref mut v) = src.data {
1322 for i in 0..16 {
1323 v[i] = i as u8;
1324 }
1325 }
1326
1327 let dims_out = vec![
1328 NDDimension {
1329 size: 2,
1330 offset: 1,
1331 binning: 2,
1332 reverse: false,
1333 },
1334 NDDimension {
1335 size: 2,
1336 offset: 1,
1337 binning: 2,
1338 reverse: false,
1339 },
1340 ];
1341
1342 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1343 assert_eq!(out.dims[0].offset, 10 + 1);
1345 assert_eq!(out.dims[1].offset, 20 + 1);
1346 assert_eq!(out.dims[0].binning, 2 * 2);
1348 assert_eq!(out.dims[1].binning, 3 * 2);
1349 }
1350
1351 #[test]
1352 fn test_convert_1d() {
1353 let pool = NDArrayPool::new(1_000_000);
1355 let mut src = NDArray::new(vec![NDDimension::new(8)], NDDataType::UInt16);
1356 if let NDDataBuffer::U16(ref mut v) = src.data {
1357 for i in 0..8 {
1358 v[i] = (i * 10) as u16;
1359 }
1360 }
1362
1363 let dims_out = vec![NDDimension {
1364 size: 4,
1365 offset: 2,
1366 binning: 2,
1367 reverse: false,
1368 }];
1369
1370 let out = pool.convert(&src, &dims_out, NDDataType::UInt16).unwrap();
1371 assert_eq!(out.dims.len(), 1);
1372 assert_eq!(out.dims[0].size, 2);
1373 if let NDDataBuffer::U16(ref v) = out.data {
1374 assert_eq!(v[0], 50);
1377 assert_eq!(v[1], 90);
1378 } else {
1379 panic!("wrong type");
1380 }
1381 }
1382
1383 #[test]
1384 fn test_convert_3d() {
1385 let pool = NDArrayPool::new(1_000_000);
1387 let mut src = NDArray::new(
1388 vec![
1389 NDDimension::new(2),
1390 NDDimension::new(2),
1391 NDDimension::new(2),
1392 ],
1393 NDDataType::UInt8,
1394 );
1395 if let NDDataBuffer::U8(ref mut v) = src.data {
1396 for i in 0..8 {
1397 v[i] = (i + 1) as u8;
1398 }
1399 }
1400
1401 let dims_out = vec![
1402 NDDimension {
1403 size: 2,
1404 offset: 0,
1405 binning: 1,
1406 reverse: false,
1407 },
1408 NDDimension {
1409 size: 2,
1410 offset: 0,
1411 binning: 1,
1412 reverse: false,
1413 },
1414 NDDimension {
1415 size: 2,
1416 offset: 0,
1417 binning: 1,
1418 reverse: false,
1419 },
1420 ];
1421
1422 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1423 if let NDDataBuffer::U8(ref v) = out.data {
1424 for i in 0..8 {
1425 assert_eq!(v[i], (i + 1) as u8);
1426 }
1427 } else {
1428 panic!("wrong type");
1429 }
1430 }
1431
1432 #[test]
1433 fn test_convert_dim_mismatch_error() {
1434 let pool = NDArrayPool::new(1_000_000);
1435 let src = make_4x4_u8();
1436 let dims_out = vec![NDDimension {
1438 size: 4,
1439 offset: 0,
1440 binning: 1,
1441 reverse: false,
1442 }];
1443
1444 let result = pool.convert(&src, &dims_out, NDDataType::UInt8);
1445 assert!(result.is_err());
1446 }
1447
1448 #[test]
1449 fn test_convert_offset_out_of_bounds_error() {
1450 let pool = NDArrayPool::new(1_000_000);
1451 let src = make_4x4_u8();
1452 let dims_out = vec![
1453 NDDimension {
1454 size: 4,
1455 offset: 2,
1456 binning: 1,
1457 reverse: false,
1458 }, NDDimension {
1460 size: 4,
1461 offset: 0,
1462 binning: 1,
1463 reverse: false,
1464 },
1465 ];
1466
1467 let result = pool.convert(&src, &dims_out, NDDataType::UInt8);
1468 assert!(result.is_err());
1469 }
1470
1471 #[test]
1472 fn test_convert_preserves_metadata() {
1473 let pool = NDArrayPool::new(1_000_000);
1474 let mut src = make_4x4_u8();
1475 src.time_stamp = 12345.678;
1476
1477 let dims_out = vec![
1478 NDDimension {
1479 size: 4,
1480 offset: 0,
1481 binning: 1,
1482 reverse: false,
1483 },
1484 NDDimension {
1485 size: 4,
1486 offset: 0,
1487 binning: 1,
1488 reverse: false,
1489 },
1490 ];
1491
1492 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1493 assert_eq!(out.time_stamp, 12345.678);
1494 }
1495
1496 #[test]
1497 fn test_convert_binning_and_reverse_combined() {
1498 let pool = NDArrayPool::new(1_000_000);
1500 let mut src = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt16);
1501 if let NDDataBuffer::U16(ref mut v) = src.data {
1502 v[0] = 1;
1503 v[1] = 2;
1504 v[2] = 3;
1505 v[3] = 4;
1506 }
1507
1508 let dims_out = vec![NDDimension {
1509 size: 4,
1510 offset: 0,
1511 binning: 2,
1512 reverse: true,
1513 }];
1514
1515 let out = pool.convert(&src, &dims_out, NDDataType::UInt16).unwrap();
1516 assert_eq!(out.dims[0].size, 2);
1517 if let NDDataBuffer::U16(ref v) = out.data {
1518 assert_eq!(v[0], 7);
1523 assert_eq!(v[1], 3);
1524 } else {
1525 panic!("wrong type");
1526 }
1527 }
1528
1529 #[test]
1533 fn test_convert_reverse_flag_cumulative() {
1534 let pool = NDArrayPool::new(1_000_000);
1535 let mut src = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
1536 src.dims[0].reverse = true; let dims_out = vec![NDDimension {
1540 size: 4,
1541 offset: 0,
1542 binning: 1,
1543 reverse: true,
1544 }];
1545 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1546 assert!(!out.dims[0].reverse, "true XOR true must be false");
1547
1548 let dims_out2 = vec![NDDimension {
1550 size: 4,
1551 offset: 0,
1552 binning: 1,
1553 reverse: false,
1554 }];
1555 let out2 = pool.convert(&src, &dims_out2, NDDataType::UInt8).unwrap();
1556 assert!(out2.dims[0].reverse, "false XOR true must be true");
1557 }
1558
1559 #[test]
1561 fn test_alloc_tracks_exact_data_size() {
1562 let pool = NDArrayPool::new(0); let a = pool
1564 .alloc(vec![NDDimension::new(333)], NDDataType::UInt16)
1565 .unwrap();
1566 assert_eq!(a.data_size, 666);
1568 assert_eq!(pool.allocated_bytes(), 666);
1569 }
1570
1571 #[test]
1573 fn test_alloc_strict_max_memory_enforcement() {
1574 let pool = NDArrayPool::new(1000);
1575 let _a = pool
1577 .alloc(vec![NDDimension::new(600)], NDDataType::UInt8)
1578 .unwrap();
1579 assert_eq!(pool.allocated_bytes(), 600);
1580 let r = pool.alloc(vec![NDDimension::new(500)], NDDataType::UInt8);
1582 assert!(r.is_err());
1583 assert!(pool.allocated_bytes() <= 1000);
1584 }
1585
1586 #[test]
1588 fn test_convert_output_is_pool_tracked() {
1589 let pool = NDArrayPool::new(1_000_000);
1590 let src = make_4x4_u8();
1591 let before_alloc = pool.num_alloc_buffers();
1592 let dims_out = vec![
1593 NDDimension {
1594 size: 4,
1595 offset: 0,
1596 binning: 1,
1597 reverse: false,
1598 },
1599 NDDimension {
1600 size: 4,
1601 offset: 0,
1602 binning: 1,
1603 reverse: false,
1604 },
1605 ];
1606 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1607 assert_eq!(out.pool_id, pool.id());
1608 assert_eq!(out.data_size, 16);
1609 assert_eq!(pool.num_alloc_buffers(), before_alloc + 1);
1611 let bytes_with_out = pool.allocated_bytes();
1612 assert_eq!(bytes_with_out, 16);
1613 pool.release(out);
1615 assert_eq!(pool.num_free_buffers(), 1);
1616 assert_eq!(pool.allocated_bytes(), 16);
1617 }
1618
1619 #[test]
1621 fn test_convert_rejects_compressed_input() {
1622 let pool = NDArrayPool::new(1_000_000);
1623 let mut src = make_4x4_u8();
1624 src.codec = Some(crate::codec::Codec {
1625 name: crate::codec::CodecName::LZ4,
1626 compressed_size: 4,
1627 level: 0,
1628 shuffle: 0,
1629 compressor: 0,
1630 original_data_type: crate::ndarray::NDDataType::UInt8,
1631 });
1632 let dims_out = vec![
1633 NDDimension {
1634 size: 4,
1635 offset: 0,
1636 binning: 1,
1637 reverse: false,
1638 },
1639 NDDimension {
1640 size: 4,
1641 offset: 0,
1642 binning: 1,
1643 reverse: false,
1644 },
1645 ];
1646 assert!(pool.convert(&src, &dims_out, NDDataType::UInt8).is_err());
1647 assert!(pool.convert_type(&src, NDDataType::UInt16).is_err());
1648 }
1649
1650 #[test]
1652 fn test_release_foreign_array_rejected() {
1653 let pool_a = NDArrayPool::new(1_000_000);
1654 let pool_b = NDArrayPool::new(1_000_000);
1655 let arr = pool_a
1656 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
1657 .unwrap();
1658 let bytes_b_before = pool_b.allocated_bytes();
1659 let free_b_before = pool_b.num_free_buffers();
1660 pool_b.release(arr);
1662 assert_eq!(pool_b.allocated_bytes(), bytes_b_before);
1663 assert_eq!(pool_b.num_free_buffers(), free_b_before);
1664 }
1665
1666 #[test]
1668 fn test_release_non_pool_array_rejected() {
1669 let pool = NDArrayPool::new(1_000_000);
1670 let arr = NDArray::new(vec![NDDimension::new(10)], NDDataType::UInt8);
1671 assert_eq!(arr.pool_id, 0);
1672 pool.release(arr);
1673 assert_eq!(pool.num_free_buffers(), 0);
1674 assert_eq!(pool.allocated_bytes(), 0);
1675 }
1676
1677 #[test]
1679 fn test_copy_allocates_and_copies() {
1680 let pool = NDArrayPool::new(1_000_000);
1681 let mut src = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
1682 if let NDDataBuffer::U8(ref mut v) = src.data {
1683 v.copy_from_slice(&[9, 8, 7, 6]);
1684 }
1685 let out = pool.copy(&src, None, true, true, true).unwrap();
1686 assert_eq!(out.pool_id, pool.id());
1687 assert_eq!(out.dims.len(), 1);
1688 if let NDDataBuffer::U8(ref v) = out.data {
1689 assert_eq!(v, &[9, 8, 7, 6]);
1690 } else {
1691 panic!("wrong type");
1692 }
1693 }
1694
1695 #[test]
1697 fn test_pre_allocate_buffers_warms_free_list() {
1698 let pool = NDArrayPool::new(10_000_000);
1699 let template = pool
1700 .alloc(vec![NDDimension::new(256)], NDDataType::UInt16)
1701 .unwrap();
1702 pool.pre_allocate_buffers(&template, 3).unwrap();
1703 assert_eq!(pool.num_free_buffers(), 3);
1704 }
1705
1706 #[test]
1717 fn test_concurrent_reuse_grow_does_not_overshoot_max_memory() {
1718 use std::sync::Arc;
1719 use std::sync::atomic::AtomicBool;
1720 use std::thread;
1721
1722 const N: usize = 16;
1723 for _ in 0..50 {
1725 let pool = Arc::new(NDArrayPool::new(2000));
1732 let mut warm = Vec::with_capacity(N);
1740 for _ in 0..N {
1741 let mut a = pool
1742 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
1743 .unwrap();
1744 if let NDDataBuffer::U8(ref mut v) = a.data {
1747 let mut big = Vec::with_capacity(200);
1748 big.resize(100, 0u8);
1749 *v = big;
1750 }
1751 assert_eq!(a.data_size, 100);
1752 assert!(a.data.capacity_bytes() >= 200);
1753 warm.push(a);
1754 }
1755 for a in warm {
1756 pool.release(a);
1757 }
1758 assert_eq!(pool.allocated_bytes(), 1600);
1759 assert_eq!(pool.num_free_buffers(), N as u32);
1760
1761 let overshoot = Arc::new(AtomicBool::new(false));
1762 let mut handles = Vec::new();
1763 for _ in 0..N {
1764 let pool = pool.clone();
1765 let overshoot = overshoot.clone();
1766 handles.push(thread::spawn(move || {
1767 let res = pool.alloc(vec![NDDimension::new(200)], NDDataType::UInt8);
1769 if res.is_ok() && pool.allocated_bytes() > pool.max_memory() as u64 {
1770 overshoot.store(true, Ordering::Relaxed);
1771 }
1772 }));
1773 }
1774 for h in handles {
1775 h.join().unwrap();
1776 }
1777
1778 assert!(
1779 !overshoot.load(Ordering::Relaxed),
1780 "allocated_bytes overshot max_memory during concurrent reuse-grow"
1781 );
1782 assert!(
1783 pool.allocated_bytes() <= pool.max_memory() as u64,
1784 "final allocated_bytes {} > max_memory {}",
1785 pool.allocated_bytes(),
1786 pool.max_memory()
1787 );
1788 }
1789 }
1790
1791 #[test]
1793 fn test_pool_report_nonempty() {
1794 let pool = NDArrayPool::new(1_000_000);
1795 let _ = pool
1796 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
1797 .unwrap();
1798 let r = pool.report(10);
1799 assert!(r.contains("NDArrayPool"));
1800 assert!(r.contains("numBuffers"));
1801 }
1802}