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};
9
10const THRESHOLD_SIZE_RATIO: f64 = 1.5;
13
14static NEXT_POOL_ID: AtomicU64 = AtomicU64::new(1);
18
19pub struct NDArrayPool {
26 id: u64,
28 max_memory: usize,
29 allocated_bytes: AtomicU64,
30 next_unique_id: AtomicI32,
31 free_list: Mutex<Vec<NDArray>>,
32 num_alloc_buffers: AtomicU32,
33 num_free_buffers: AtomicU32,
34}
35
36impl NDArrayPool {
37 pub fn new(max_memory: usize) -> Self {
38 Self {
39 id: NEXT_POOL_ID.fetch_add(1, Ordering::Relaxed),
40 max_memory,
41 allocated_bytes: AtomicU64::new(0),
42 next_unique_id: AtomicI32::new(1),
43 free_list: Mutex::new(Vec::new()),
44 num_alloc_buffers: AtomicU32::new(0),
45 num_free_buffers: AtomicU32::new(0),
46 }
47 }
48
49 pub fn id(&self) -> u64 {
51 self.id
52 }
53
54 pub fn alloc(&self, dims: Vec<NDDimension>, data_type: NDDataType) -> ADResult<NDArray> {
62 let num_elements: usize = dims.iter().map(|d| d.size).product();
63 let needed_bytes = num_elements * data_type.element_size();
64
65 let reused = {
69 let mut free = self.free_list.lock();
70 let mut best_idx = None;
71 let mut best_cap = usize::MAX;
72 for (i, arr) in free.iter().enumerate() {
73 let cap = arr.data.capacity_bytes();
74 if cap >= needed_bytes && cap < best_cap {
75 best_cap = cap;
76 best_idx = Some(i);
77 }
78 }
79 if let Some(idx) = best_idx {
80 if best_cap as f64 > needed_bytes as f64 * THRESHOLD_SIZE_RATIO {
81 let dropped = free.swap_remove(idx);
83 self.num_free_buffers.fetch_sub(1, Ordering::Relaxed);
84 self.allocated_bytes
85 .fetch_sub(dropped.data_size as u64, Ordering::Relaxed);
86 self.num_alloc_buffers.fetch_sub(1, Ordering::Relaxed);
87 None
88 } else {
89 let arr = free.swap_remove(idx);
90 self.num_free_buffers.fetch_sub(1, Ordering::Relaxed);
91 Some(arr)
92 }
93 } else {
94 None
95 }
96 };
97
98 let mut arr = if let Some(mut reused) = reused {
99 let old_size = reused.data_size;
105 if reused.data.data_type() != data_type {
106 reused.data = NDDataBuffer::zeros(data_type, num_elements);
107 } else {
108 reused.data.resize(num_elements);
109 }
110 let effective_size = if needed_bytes > old_size {
111 let diff = (needed_bytes - old_size) as u64;
112 if self.max_memory > 0 {
117 loop {
118 let current = self.allocated_bytes.load(Ordering::Relaxed);
119 if current + diff > self.max_memory as u64 {
120 let mut free = self.free_list.lock();
123 free.push(reused);
124 self.num_free_buffers.fetch_add(1, Ordering::Relaxed);
125 return Err(ADError::PoolExhausted(needed_bytes, self.max_memory));
126 }
127 if self
128 .allocated_bytes
129 .compare_exchange_weak(
130 current,
131 current + diff,
132 Ordering::Relaxed,
133 Ordering::Relaxed,
134 )
135 .is_ok()
136 {
137 break;
138 }
139 }
140 } else {
141 self.allocated_bytes.fetch_add(diff, Ordering::Relaxed);
142 }
143 needed_bytes
144 } else {
145 old_size
148 };
149 reused.data_size = effective_size;
150 reused.dims = dims;
151 reused.attributes.clear();
152 reused.codec = None;
153 reused
154 } else {
155 if self.max_memory > 0 {
158 loop {
159 let current = self.allocated_bytes.load(Ordering::Relaxed);
160 if current + needed_bytes as u64 > self.max_memory as u64 {
161 let mut freed_enough = false;
162 {
163 let mut free = self.free_list.lock();
164 free.sort_by(|a, b| {
165 b.data.capacity_bytes().cmp(&a.data.capacity_bytes())
166 });
167 let mut reclaimed = 0u64;
168 let over = (current + needed_bytes as u64)
169 .saturating_sub(self.max_memory as u64);
170 while !free.is_empty() && reclaimed < over {
171 let dropped = free.remove(0);
172 let dropped_size = dropped.data_size as u64;
173 self.allocated_bytes
174 .fetch_sub(dropped_size, Ordering::Relaxed);
175 self.num_free_buffers.fetch_sub(1, Ordering::Relaxed);
176 self.num_alloc_buffers.fetch_sub(1, Ordering::Relaxed);
177 reclaimed += dropped_size;
178 }
179 if reclaimed >= over {
180 freed_enough = true;
181 }
182 }
183 if !freed_enough {
184 return Err(ADError::PoolExhausted(needed_bytes, self.max_memory));
185 }
186 continue;
187 }
188 if self
189 .allocated_bytes
190 .compare_exchange_weak(
191 current,
192 current + needed_bytes as u64,
193 Ordering::Relaxed,
194 Ordering::Relaxed,
195 )
196 .is_ok()
197 {
198 break;
199 }
200 }
201 } else {
202 self.allocated_bytes
203 .fetch_add(needed_bytes as u64, Ordering::Relaxed);
204 }
205 self.num_alloc_buffers.fetch_add(1, Ordering::Relaxed);
206 NDArray::new(dims, data_type)
207 };
208
209 arr.unique_id = self.next_unique_id.fetch_add(1, Ordering::Relaxed);
210 arr.pool_id = self.id;
218 Ok(arr)
221 }
222
223 pub fn alloc_copy(&self, source: &NDArray) -> ADResult<NDArray> {
226 let dims = source.dims.clone();
227 let data_type = source.data.data_type();
228 let mut copy = self.alloc(dims, data_type)?;
229 copy.data = source.data.clone();
230 copy.time_stamp = source.time_stamp;
234 copy.timestamp = source.timestamp;
235 copy.attributes = source.attributes.clone();
236 copy.codec = source.codec.clone();
237 Ok(copy)
238 }
239
240 pub fn release(&self, array: NDArray) {
247 if array.pool_id != self.id {
248 return;
252 }
253
254 let mut free = self.free_list.lock();
255 free.push(array);
256 self.num_free_buffers.fetch_add(1, Ordering::Relaxed);
257
258 let total = self.allocated_bytes.load(Ordering::Relaxed) as usize;
262 if self.max_memory > 0 && total > self.max_memory && !free.is_empty() {
263 free.sort_by(|a, b| b.data_size.cmp(&a.data_size));
264 let mut excess = total - self.max_memory;
265 while excess > 0 && !free.is_empty() {
266 let dropped = free.remove(0);
267 let dropped_size = dropped.data_size;
268 self.allocated_bytes
269 .fetch_sub(dropped_size as u64, Ordering::Relaxed);
270 self.num_free_buffers.fetch_sub(1, Ordering::Relaxed);
271 self.num_alloc_buffers.fetch_sub(1, Ordering::Relaxed);
272 if dropped_size >= excess {
273 break;
274 }
275 excess -= dropped_size;
276 }
277 }
278 }
279
280 pub fn empty_free_list(&self) {
282 let mut free = self.free_list.lock();
283 let count = free.len() as u32;
284 for arr in free.drain(..) {
285 self.allocated_bytes
286 .fetch_sub(arr.data_size as u64, Ordering::Relaxed);
287 self.num_alloc_buffers.fetch_sub(1, Ordering::Relaxed);
288 }
289 self.num_free_buffers.fetch_sub(count, Ordering::Relaxed);
290 }
291
292 pub fn allocated_bytes(&self) -> u64 {
293 self.allocated_bytes.load(Ordering::Relaxed)
294 }
295
296 pub fn num_alloc_buffers(&self) -> u32 {
297 self.num_alloc_buffers.load(Ordering::Relaxed)
298 }
299
300 pub fn num_free_buffers(&self) -> u32 {
301 self.num_free_buffers.load(Ordering::Relaxed)
302 }
303
304 pub fn max_memory(&self) -> usize {
305 self.max_memory
306 }
307
308 pub fn alloc_handle(
311 pool: &Arc<Self>,
312 dims: Vec<NDDimension>,
313 data_type: NDDataType,
314 ) -> ADResult<NDArrayHandle> {
315 let array = pool.alloc(dims, data_type)?;
316 Ok(pooled_array(array, pool))
317 }
318
319 pub fn copy(
331 &self,
332 src: &NDArray,
333 out: Option<NDArray>,
334 copy_data: bool,
335 copy_dimensions: bool,
336 copy_data_type: bool,
337 ) -> ADResult<NDArray> {
338 let mut out = match out {
339 Some(o) => o,
340 None => self.alloc(src.dims.clone(), src.data.data_type())?,
341 };
342
343 out.unique_id = src.unique_id;
344 out.time_stamp = src.time_stamp;
345 out.timestamp = src.timestamp;
346 if copy_dimensions {
347 out.dims = src.dims.clone();
348 }
349 out.codec = src.codec.clone();
350
351 if copy_data {
352 if copy_data_type && out.data.data_type() != src.data.data_type() {
353 out.data = src.data.clone();
355 } else if out.data.data_type() == src.data.data_type() {
356 out.data = src.data.clone();
357 } else {
358 out.data = crate::color::convert_data_type(src, out.data.data_type())?.data;
360 }
361 } else if copy_data_type && out.data.data_type() != src.data.data_type() {
362 out.data = NDDataBuffer::zeros(src.data.data_type(), out.data.len());
363 }
364
365 out.attributes.clear();
366 out.attributes.copy_from(&src.attributes);
367 Ok(out)
368 }
369
370 pub fn pre_allocate_buffers(&self, template_array: &NDArray, count: usize) -> ADResult<()> {
376 let mut buffers = Vec::with_capacity(count);
377 for _ in 0..count {
378 buffers.push(self.copy(template_array, None, true, true, true)?);
379 }
380 for arr in buffers {
381 self.release(arr);
382 }
383 Ok(())
384 }
385
386 pub fn convert_type(&self, src: &NDArray, target_type: NDDataType) -> ADResult<NDArray> {
389 if src.codec.is_some() {
391 return Err(ADError::UnsupportedConversion(
392 "convert_type: cannot convert compressed (codec) data".into(),
393 ));
394 }
395 if src.data.data_type() == target_type {
396 return self.alloc_copy(src);
397 }
398 let mut out = self.alloc(src.dims.clone(), target_type)?;
401 let converted = crate::color::convert_data_type(src, target_type)?;
402 out.data = converted.data;
403 out.time_stamp = src.time_stamp;
404 out.timestamp = src.timestamp;
405 out.attributes.copy_from(&src.attributes);
406 Ok(out)
407 }
408
409 pub fn convert(
418 &self,
419 src: &NDArray,
420 dims_out: &[NDDimension],
421 target_type: NDDataType,
422 ) -> ADResult<NDArray> {
423 let converted = crate::convert::convert_dims(src, dims_out, target_type)?;
424
425 let mut arr = self.alloc(converted.dims.clone(), target_type)?;
426 arr.timestamp = src.timestamp;
427 arr.time_stamp = src.time_stamp;
428 arr.attributes.copy_from(&src.attributes);
429 arr.data = converted.data;
430
431 Ok(arr)
432 }
433
434 pub fn report(&self, details: i32) -> String {
438 let mut out = String::new();
439 out.push('\n');
440 out.push_str("NDArrayPool:\n");
441 out.push_str(&format!(
442 " numBuffers={}, numFree={}\n",
443 self.num_alloc_buffers(),
444 self.num_free_buffers()
445 ));
446 out.push_str(&format!(
447 " memorySize={}, maxMemory={}\n",
448 self.allocated_bytes(),
449 self.max_memory
450 ));
451 if details > 5 {
452 let free = self.free_list.lock();
453 out.push_str(" freeList: (index, dataSize, capacity)\n");
454 for (i, arr) in free.iter().enumerate() {
455 out.push_str(&format!(
456 " {} {} {}\n",
457 i,
458 arr.data_size,
459 arr.data.capacity_bytes()
460 ));
461 }
462 if details > 10 {
463 for arr in free.iter() {
464 out.push_str(&arr.report(details));
465 }
466 }
467 }
468 out
469 }
470}
471
472const _: fn() = || {
474 fn assert_send_sync<T: Send + Sync>() {}
475 assert_send_sync::<NDArrayPool>();
476};
477
478#[cfg(test)]
479mod tests {
480 use super::*;
481
482 #[test]
483 fn test_alloc_auto_id() {
484 let pool = NDArrayPool::new(1_000_000);
485 let a1 = pool
486 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
487 .unwrap();
488 let a2 = pool
489 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
490 .unwrap();
491 assert_eq!(a1.unique_id, 1);
492 assert_eq!(a2.unique_id, 2);
493 }
494
495 #[test]
496 fn test_r6_65_alloc_stamps_neither_timestamp() {
497 let pool = NDArrayPool::new(1_000_000);
501 let fresh = pool
502 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
503 .unwrap();
504 assert_eq!(fresh.timestamp, crate::timestamp::EpicsTimestamp::default());
505 assert_eq!(fresh.time_stamp, 0.0);
506
507 let mut used = pool
510 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
511 .unwrap();
512 used.update_time_stamps(crate::timestamp::EpicsTimestamp {
513 sec: 1000,
514 nsec: 250_000_000,
515 });
516 pool.release(used);
517 let reused = pool
518 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
519 .unwrap();
520 assert_eq!(reused.timestamp.sec, 1000);
521 assert_eq!(reused.time_stamp, 1000.25);
522 }
523
524 #[test]
525 fn test_r6_65_update_time_stamps_derives_the_double() {
526 let mut arr = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
529 arr.update_time_stamps(crate::timestamp::EpicsTimestamp {
530 sec: 42,
531 nsec: 500_000_000,
532 });
533 assert_eq!(arr.timestamp.sec, 42);
534 assert_eq!(arr.timestamp.nsec, 500_000_000);
535 assert_eq!(arr.time_stamp, 42.5);
536 }
537
538 #[test]
539 fn test_r6_65_alloc_copy_carries_both_stamps() {
540 let pool = NDArrayPool::new(1_000_000);
542 let mut src = pool
543 .alloc(vec![NDDimension::new(8)], NDDataType::UInt8)
544 .unwrap();
545 src.update_time_stamps(crate::timestamp::EpicsTimestamp {
546 sec: 777,
547 nsec: 125_000_000,
548 });
549
550 let copy = pool.alloc_copy(&src).unwrap();
551 assert_eq!(copy.timestamp, src.timestamp);
552 assert_eq!(copy.time_stamp, src.time_stamp);
553 assert_eq!(copy.time_stamp, 777.125);
554 }
555
556 #[test]
557 fn test_alloc_tracks_bytes() {
558 let pool = NDArrayPool::new(1_000_000);
559 let _ = pool
560 .alloc(vec![NDDimension::new(100)], NDDataType::Float64)
561 .unwrap();
562 assert!(pool.allocated_bytes() >= 800);
563 }
564
565 #[test]
566 fn test_alloc_exceeds_max() {
567 let pool = NDArrayPool::new(100);
568 let result = pool.alloc(vec![NDDimension::new(200)], NDDataType::UInt8);
569 assert!(result.is_err());
570 }
571
572 #[test]
573 fn test_alloc_copy_preserves_data() {
574 let pool = NDArrayPool::new(1_000_000);
575 let mut source = pool
576 .alloc(vec![NDDimension::new(4)], NDDataType::UInt8)
577 .unwrap();
578 if let NDDataBuffer::U8(ref mut v) = source.data {
579 v[0] = 1;
580 v[1] = 2;
581 v[2] = 3;
582 v[3] = 4;
583 }
584
585 let copy = pool.alloc_copy(&source).unwrap();
586 assert_ne!(copy.unique_id, source.unique_id);
587 assert_eq!(copy.dims.len(), source.dims.len());
588 if let NDDataBuffer::U8(ref v) = copy.data {
589 assert_eq!(v, &[1, 2, 3, 4]);
590 } else {
591 panic!("wrong type");
592 }
593 }
594
595 #[test]
596 fn test_alloc_copy_tracks_bytes() {
597 let pool = NDArrayPool::new(1_000_000);
598 let source = pool
599 .alloc(vec![NDDimension::new(10)], NDDataType::UInt16)
600 .unwrap();
601 assert_eq!(pool.allocated_bytes(), 20);
602 let _ = pool.alloc_copy(&source).unwrap();
603 assert!(pool.allocated_bytes() >= 40);
604 }
605
606 #[test]
607 fn test_alloc_copy_exceeds_max() {
608 let pool = NDArrayPool::new(60);
609 let source = pool
610 .alloc(vec![NDDimension::new(50)], NDDataType::UInt8)
611 .unwrap();
612 assert!(pool.alloc_copy(&source).is_err());
613 }
614
615 #[test]
618 fn test_release_and_reuse() {
619 let pool = NDArrayPool::new(1_000_000);
620 let arr = pool
621 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
622 .unwrap();
623 let _alloc_bytes_after_first = pool.allocated_bytes();
624 assert_eq!(pool.num_alloc_buffers(), 1);
625
626 pool.release(arr);
628 assert_eq!(pool.num_free_buffers(), 1);
629
630 let arr2 = pool
632 .alloc(vec![NDDimension::new(80)], NDDataType::UInt8)
633 .unwrap();
634 assert_eq!(arr2.data.len(), 80);
635 }
636
637 #[test]
638 fn test_free_list_prefers_smallest_sufficient() {
639 let pool = NDArrayPool::new(10_000_000);
640 let small = pool
641 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
642 .unwrap();
643 let large = pool
644 .alloc(vec![NDDimension::new(10000)], NDDataType::UInt8)
645 .unwrap();
646 let medium = pool
647 .alloc(vec![NDDimension::new(1000)], NDDataType::UInt8)
648 .unwrap();
649
650 pool.release(large);
651 pool.release(medium);
652 pool.release(small);
653 assert_eq!(pool.num_free_buffers(), 3);
654
655 let reused = pool
657 .alloc(vec![NDDimension::new(900)], NDDataType::UInt8)
658 .unwrap();
659 assert!(reused.data.capacity_bytes() >= 900);
660 }
661
662 #[test]
663 fn test_empty_free_list() {
664 let pool = NDArrayPool::new(1_000_000);
665 let a1 = pool
666 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
667 .unwrap();
668 let a2 = pool
669 .alloc(vec![NDDimension::new(200)], NDDataType::UInt8)
670 .unwrap();
671 pool.release(a1);
672 pool.release(a2);
673 assert_eq!(pool.num_free_buffers(), 2);
674
675 pool.empty_free_list();
676 assert_eq!(pool.num_free_buffers(), 0);
677 assert_eq!(pool.num_alloc_buffers(), 0);
678 }
679
680 #[test]
681 fn test_num_free_buffers_tracking() {
682 let pool = NDArrayPool::new(1_000_000);
683 assert_eq!(pool.num_free_buffers(), 0);
684
685 let a = pool
686 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
687 .unwrap();
688 assert_eq!(pool.num_free_buffers(), 0);
689
690 pool.release(a);
691 assert_eq!(pool.num_free_buffers(), 1);
692
693 let _ = pool
694 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
695 .unwrap();
696 assert_eq!(pool.num_free_buffers(), 0);
697 }
698
699 #[test]
700 fn test_concurrent_alloc_release() {
701 use std::sync::Arc;
702 use std::thread;
703
704 let pool = Arc::new(NDArrayPool::new(10_000_000));
705 let mut handles = Vec::new();
706
707 for _ in 0..4 {
708 let pool = pool.clone();
709 handles.push(thread::spawn(move || {
710 for _ in 0..100 {
711 let arr = pool
712 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
713 .unwrap();
714 pool.release(arr);
715 }
716 }));
717 }
718
719 for h in handles {
720 h.join().unwrap();
721 }
722
723 assert!(pool.num_free_buffers() > 0);
725 }
726
727 #[test]
728 fn test_max_memory() {
729 let pool = NDArrayPool::new(42);
730 assert_eq!(pool.max_memory(), 42);
731 }
732
733 #[test]
736 fn test_convert_type_same_type() {
737 let pool = NDArrayPool::new(1_000_000);
738 let mut src = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
739 if let NDDataBuffer::U8(ref mut v) = src.data {
740 v[0] = 10;
741 v[1] = 20;
742 v[2] = 30;
743 v[3] = 40;
744 }
745
746 let out = pool.convert_type(&src, NDDataType::UInt8).unwrap();
747 assert_eq!(out.data.data_type(), NDDataType::UInt8);
748 if let NDDataBuffer::U8(ref v) = out.data {
749 assert_eq!(v, &[10, 20, 30, 40]);
750 } else {
751 panic!("wrong type");
752 }
753 }
754
755 #[test]
756 fn test_convert_type_u8_to_f32() {
757 let pool = NDArrayPool::new(1_000_000);
758 let mut src = NDArray::new(vec![NDDimension::new(3)], NDDataType::UInt8);
759 if let NDDataBuffer::U8(ref mut v) = src.data {
760 v[0] = 0;
761 v[1] = 128;
762 v[2] = 255;
763 }
764
765 let out = pool.convert_type(&src, NDDataType::Float32).unwrap();
766 assert_eq!(out.data.data_type(), NDDataType::Float32);
767 if let NDDataBuffer::F32(ref v) = out.data {
768 assert_eq!(v[0], 0.0);
769 assert_eq!(v[1], 128.0);
770 assert_eq!(v[2], 255.0);
771 } else {
772 panic!("wrong type");
773 }
774 }
775
776 #[test]
780 fn test_convert_type_u16_to_u8_wraps_like_a_c_cast() {
781 let pool = NDArrayPool::new(1_000_000);
782 let mut src = NDArray::new(vec![NDDimension::new(2)], NDDataType::UInt16);
783 if let NDDataBuffer::U16(ref mut v) = src.data {
784 v[0] = 100;
785 v[1] = 300;
786 }
787
788 let out = pool.convert_type(&src, NDDataType::UInt8).unwrap();
789 if let NDDataBuffer::U8(ref v) = out.data {
790 assert_eq!(v[0], 100);
791 assert_eq!(v[1], 44, "300 % 256 == 44, not a clamp to 255");
792 } else {
793 panic!("wrong type");
794 }
795 }
796
797 fn make_4x4_u8() -> NDArray {
801 let mut arr = NDArray::new(
802 vec![NDDimension::new(4), NDDimension::new(4)],
803 NDDataType::UInt8,
804 );
805 if let NDDataBuffer::U8(ref mut v) = arr.data {
806 for i in 0..16 {
807 v[i] = i as u8;
808 }
809 }
810 arr
811 }
812
813 #[test]
814 fn test_convert_identity() {
815 let pool = NDArrayPool::new(1_000_000);
817 let src = make_4x4_u8();
818 let dims_out = vec![
819 NDDimension {
820 size: 4,
821 offset: 0,
822 binning: 1,
823 reverse: false,
824 },
825 NDDimension {
826 size: 4,
827 offset: 0,
828 binning: 1,
829 reverse: false,
830 },
831 ];
832
833 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
834 assert_eq!(out.dims[0].size, 4);
835 assert_eq!(out.dims[1].size, 4);
836 if let NDDataBuffer::U8(ref v) = out.data {
837 for i in 0..16 {
838 assert_eq!(v[i], i as u8);
839 }
840 } else {
841 panic!("wrong type");
842 }
843 }
844
845 #[test]
846 fn test_convert_offset_extraction() {
847 let pool = NDArrayPool::new(1_000_000);
849 let src = make_4x4_u8();
850 let dims_out = vec![
851 NDDimension {
852 size: 2,
853 offset: 1,
854 binning: 1,
855 reverse: false,
856 },
857 NDDimension {
858 size: 2,
859 offset: 1,
860 binning: 1,
861 reverse: false,
862 },
863 ];
864
865 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
866 assert_eq!(out.dims[0].size, 2);
867 assert_eq!(out.dims[1].size, 2);
868 if let NDDataBuffer::U8(ref v) = out.data {
872 assert_eq!(v[0], 5);
873 assert_eq!(v[1], 6);
874 assert_eq!(v[2], 9);
875 assert_eq!(v[3], 10);
876 } else {
877 panic!("wrong type");
878 }
879
880 assert_eq!(out.dims[0].offset, 1); assert_eq!(out.dims[1].offset, 1);
883 }
884
885 #[test]
886 fn test_convert_binning_2x2() {
887 let pool = NDArrayPool::new(1_000_000);
889 let src = make_4x4_u8();
890 let dims_out = vec![
891 NDDimension {
892 size: 4,
893 offset: 0,
894 binning: 2,
895 reverse: false,
896 },
897 NDDimension {
898 size: 4,
899 offset: 0,
900 binning: 2,
901 reverse: false,
902 },
903 ];
904
905 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
906 assert_eq!(out.dims[0].size, 2);
907 assert_eq!(out.dims[1].size, 2);
908 if let NDDataBuffer::U8(ref v) = out.data {
913 assert_eq!(v[0], 10);
914 assert_eq!(v[1], 18);
915 assert_eq!(v[2], 42);
916 assert_eq!(v[3], 50);
917 } else {
918 panic!("wrong type");
919 }
920
921 assert_eq!(out.dims[0].binning, 2); assert_eq!(out.dims[1].binning, 2);
924 }
925
926 fn make_2x2_filled(data_type: NDDataType, val: i128) -> NDArray {
929 let mut src = NDArray::new(vec![NDDimension::new(2), NDDimension::new(2)], data_type);
930 macro_rules! fill {
931 ($variant:ident, $t:ty) => {
932 if let NDDataBuffer::$variant(ref mut v) = src.data {
933 for e in v.iter_mut() {
934 *e = val as $t;
935 }
936 }
937 };
938 }
939 match data_type {
940 NDDataType::Int8 => fill!(I8, i8),
941 NDDataType::UInt8 => fill!(U8, u8),
942 NDDataType::Int16 => fill!(I16, i16),
943 NDDataType::UInt16 => fill!(U16, u16),
944 NDDataType::Int32 => fill!(I32, i32),
945 NDDataType::UInt32 => fill!(U32, u32),
946 NDDataType::Int64 => fill!(I64, i64),
947 NDDataType::UInt64 => fill!(U64, u64),
948 NDDataType::Float32 => fill!(F32, f32),
949 NDDataType::Float64 => fill!(F64, f64),
950 }
951 src
952 }
953
954 fn bin_2x2_to_one() -> Vec<NDDimension> {
955 vec![
957 NDDimension {
958 size: 2,
959 offset: 0,
960 binning: 2,
961 reverse: false,
962 },
963 NDDimension {
964 size: 2,
965 offset: 0,
966 binning: 2,
967 reverse: false,
968 },
969 ]
970 }
971
972 #[test]
977 fn test_convert_binning_widening_target_keeps_full_sum() {
978 let pool = NDArrayPool::new(1_000_000);
979 let src = make_2x2_filled(NDDataType::UInt8, 200);
980 let out = pool
981 .convert(&src, &bin_2x2_to_one(), NDDataType::UInt16)
982 .unwrap();
983 match out.data {
984 NDDataBuffer::U16(ref v) => assert_eq!(v[0], 800, "C sums in u16: 200*4"),
985 ref other => panic!("expected U16, got {:?}", other.data_type()),
986 }
987 }
988
989 #[test]
993 fn test_convert_binning_integer_overflow_wraps() {
994 let pool = NDArrayPool::new(1_000_000);
995 let src = make_2x2_filled(NDDataType::UInt8, 100);
996 let out = pool
997 .convert(&src, &bin_2x2_to_one(), NDDataType::UInt8)
998 .unwrap();
999 match out.data {
1000 NDDataBuffer::U8(ref v) => assert_eq!(v[0], 144, "400 wraps mod 256 in u8"),
1001 ref other => panic!("expected U8, got {:?}", other.data_type()),
1002 }
1003 }
1004
1005 #[test]
1010 fn test_convert_binning1_int64_above_2pow53_exact() {
1011 let pool = NDArrayPool::new(1_000_000);
1012 let value: i64 = (1i64 << 53) + 1;
1013 let mut src = NDArray::new(
1014 vec![NDDimension::new(1), NDDimension::new(1)],
1015 NDDataType::Int64,
1016 );
1017 if let NDDataBuffer::I64(ref mut v) = src.data {
1018 v[0] = value;
1019 }
1020 let dims_out = vec![
1021 NDDimension {
1022 size: 1,
1023 offset: 0,
1024 binning: 1,
1025 reverse: false,
1026 },
1027 NDDimension {
1028 size: 1,
1029 offset: 0,
1030 binning: 1,
1031 reverse: false,
1032 },
1033 ];
1034 let out = pool.convert(&src, &dims_out, NDDataType::Int64).unwrap();
1035 match out.data {
1036 NDDataBuffer::I64(ref v) => assert_eq!(v[0], value, "2^53+1 kept exactly"),
1037 ref other => panic!("expected I64, got {:?}", other.data_type()),
1038 }
1039 }
1040
1041 #[test]
1042 fn test_convert_reverse_x() {
1043 let pool = NDArrayPool::new(1_000_000);
1045 let mut src = NDArray::new(
1046 vec![NDDimension::new(4), NDDimension::new(1)],
1047 NDDataType::UInt8,
1048 );
1049 if let NDDataBuffer::U8(ref mut v) = src.data {
1050 v[0] = 10;
1051 v[1] = 20;
1052 v[2] = 30;
1053 v[3] = 40;
1054 }
1055
1056 let dims_out = vec![
1057 NDDimension {
1058 size: 4,
1059 offset: 0,
1060 binning: 1,
1061 reverse: true,
1062 },
1063 NDDimension {
1064 size: 1,
1065 offset: 0,
1066 binning: 1,
1067 reverse: false,
1068 },
1069 ];
1070
1071 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1072 if let NDDataBuffer::U8(ref v) = out.data {
1073 assert_eq!(v[0], 40);
1074 assert_eq!(v[1], 30);
1075 assert_eq!(v[2], 20);
1076 assert_eq!(v[3], 10);
1077 } else {
1078 panic!("wrong type");
1079 }
1080 }
1081
1082 #[test]
1083 fn test_convert_reverse_y() {
1084 let pool = NDArrayPool::new(1_000_000);
1086 let mut src = NDArray::new(
1087 vec![NDDimension::new(2), NDDimension::new(2)],
1088 NDDataType::UInt16,
1089 );
1090 if let NDDataBuffer::U16(ref mut v) = src.data {
1091 v[0] = 1;
1093 v[1] = 2;
1094 v[2] = 3;
1095 v[3] = 4;
1096 }
1097
1098 let dims_out = vec![
1099 NDDimension {
1100 size: 2,
1101 offset: 0,
1102 binning: 1,
1103 reverse: false,
1104 },
1105 NDDimension {
1106 size: 2,
1107 offset: 0,
1108 binning: 1,
1109 reverse: true,
1110 },
1111 ];
1112
1113 let out = pool.convert(&src, &dims_out, NDDataType::UInt16).unwrap();
1114 if let NDDataBuffer::U16(ref v) = out.data {
1115 assert_eq!(v[0], 3);
1117 assert_eq!(v[1], 4);
1118 assert_eq!(v[2], 1);
1119 assert_eq!(v[3], 2);
1120 } else {
1121 panic!("wrong type");
1122 }
1123 }
1124
1125 #[test]
1126 fn test_convert_with_type_change() {
1127 let pool = NDArrayPool::new(1_000_000);
1129 let src = make_4x4_u8();
1130 let dims_out = vec![
1131 NDDimension {
1132 size: 4,
1133 offset: 0,
1134 binning: 2,
1135 reverse: false,
1136 },
1137 NDDimension {
1138 size: 4,
1139 offset: 0,
1140 binning: 2,
1141 reverse: false,
1142 },
1143 ];
1144
1145 let out = pool.convert(&src, &dims_out, NDDataType::Float32).unwrap();
1146 assert_eq!(out.data.data_type(), NDDataType::Float32);
1147 assert_eq!(out.dims[0].size, 2);
1148 assert_eq!(out.dims[1].size, 2);
1149 if let NDDataBuffer::F32(ref v) = out.data {
1150 assert_eq!(v[0], 10.0); assert_eq!(v[1], 18.0); } else {
1153 panic!("wrong type");
1154 }
1155 }
1156
1157 #[test]
1158 fn test_convert_cumulative_offset_and_binning() {
1159 let pool = NDArrayPool::new(1_000_000);
1161 let mut src = NDArray::new(
1162 vec![NDDimension::new(4), NDDimension::new(4)],
1163 NDDataType::UInt8,
1164 );
1165 src.dims[0].offset = 10;
1166 src.dims[0].binning = 2;
1167 src.dims[1].offset = 20;
1168 src.dims[1].binning = 3;
1169 if let NDDataBuffer::U8(ref mut v) = src.data {
1170 for i in 0..16 {
1171 v[i] = i as u8;
1172 }
1173 }
1174
1175 let dims_out = vec![
1176 NDDimension {
1177 size: 2,
1178 offset: 1,
1179 binning: 2,
1180 reverse: false,
1181 },
1182 NDDimension {
1183 size: 2,
1184 offset: 1,
1185 binning: 2,
1186 reverse: false,
1187 },
1188 ];
1189
1190 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1191 assert_eq!(out.dims[0].offset, 10 + 1);
1193 assert_eq!(out.dims[1].offset, 20 + 1);
1194 assert_eq!(out.dims[0].binning, 2 * 2);
1196 assert_eq!(out.dims[1].binning, 3 * 2);
1197 }
1198
1199 #[test]
1200 fn test_convert_1d() {
1201 let pool = NDArrayPool::new(1_000_000);
1203 let mut src = NDArray::new(vec![NDDimension::new(8)], NDDataType::UInt16);
1204 if let NDDataBuffer::U16(ref mut v) = src.data {
1205 for i in 0..8 {
1206 v[i] = (i * 10) as u16;
1207 }
1208 }
1210
1211 let dims_out = vec![NDDimension {
1212 size: 4,
1213 offset: 2,
1214 binning: 2,
1215 reverse: false,
1216 }];
1217
1218 let out = pool.convert(&src, &dims_out, NDDataType::UInt16).unwrap();
1219 assert_eq!(out.dims.len(), 1);
1220 assert_eq!(out.dims[0].size, 2);
1221 if let NDDataBuffer::U16(ref v) = out.data {
1222 assert_eq!(v[0], 50);
1225 assert_eq!(v[1], 90);
1226 } else {
1227 panic!("wrong type");
1228 }
1229 }
1230
1231 #[test]
1232 fn test_convert_3d() {
1233 let pool = NDArrayPool::new(1_000_000);
1235 let mut src = NDArray::new(
1236 vec![
1237 NDDimension::new(2),
1238 NDDimension::new(2),
1239 NDDimension::new(2),
1240 ],
1241 NDDataType::UInt8,
1242 );
1243 if let NDDataBuffer::U8(ref mut v) = src.data {
1244 for i in 0..8 {
1245 v[i] = (i + 1) as u8;
1246 }
1247 }
1248
1249 let dims_out = vec![
1250 NDDimension {
1251 size: 2,
1252 offset: 0,
1253 binning: 1,
1254 reverse: false,
1255 },
1256 NDDimension {
1257 size: 2,
1258 offset: 0,
1259 binning: 1,
1260 reverse: false,
1261 },
1262 NDDimension {
1263 size: 2,
1264 offset: 0,
1265 binning: 1,
1266 reverse: false,
1267 },
1268 ];
1269
1270 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1271 if let NDDataBuffer::U8(ref v) = out.data {
1272 for i in 0..8 {
1273 assert_eq!(v[i], (i + 1) as u8);
1274 }
1275 } else {
1276 panic!("wrong type");
1277 }
1278 }
1279
1280 #[test]
1281 fn test_convert_dim_mismatch_error() {
1282 let pool = NDArrayPool::new(1_000_000);
1283 let src = make_4x4_u8();
1284 let dims_out = vec![NDDimension {
1286 size: 4,
1287 offset: 0,
1288 binning: 1,
1289 reverse: false,
1290 }];
1291
1292 let result = pool.convert(&src, &dims_out, NDDataType::UInt8);
1293 assert!(result.is_err());
1294 }
1295
1296 #[test]
1297 fn test_convert_offset_out_of_bounds_error() {
1298 let pool = NDArrayPool::new(1_000_000);
1299 let src = make_4x4_u8();
1300 let dims_out = vec![
1301 NDDimension {
1302 size: 4,
1303 offset: 2,
1304 binning: 1,
1305 reverse: false,
1306 }, NDDimension {
1308 size: 4,
1309 offset: 0,
1310 binning: 1,
1311 reverse: false,
1312 },
1313 ];
1314
1315 let result = pool.convert(&src, &dims_out, NDDataType::UInt8);
1316 assert!(result.is_err());
1317 }
1318
1319 #[test]
1320 fn test_convert_preserves_metadata() {
1321 let pool = NDArrayPool::new(1_000_000);
1322 let mut src = make_4x4_u8();
1323 src.time_stamp = 12345.678;
1324
1325 let dims_out = vec![
1326 NDDimension {
1327 size: 4,
1328 offset: 0,
1329 binning: 1,
1330 reverse: false,
1331 },
1332 NDDimension {
1333 size: 4,
1334 offset: 0,
1335 binning: 1,
1336 reverse: false,
1337 },
1338 ];
1339
1340 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1341 assert_eq!(out.time_stamp, 12345.678);
1342 }
1343
1344 #[test]
1345 fn test_convert_binning_and_reverse_combined() {
1346 let pool = NDArrayPool::new(1_000_000);
1348 let mut src = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt16);
1349 if let NDDataBuffer::U16(ref mut v) = src.data {
1350 v[0] = 1;
1351 v[1] = 2;
1352 v[2] = 3;
1353 v[3] = 4;
1354 }
1355
1356 let dims_out = vec![NDDimension {
1357 size: 4,
1358 offset: 0,
1359 binning: 2,
1360 reverse: true,
1361 }];
1362
1363 let out = pool.convert(&src, &dims_out, NDDataType::UInt16).unwrap();
1364 assert_eq!(out.dims[0].size, 2);
1365 if let NDDataBuffer::U16(ref v) = out.data {
1366 assert_eq!(v[0], 7);
1371 assert_eq!(v[1], 3);
1372 } else {
1373 panic!("wrong type");
1374 }
1375 }
1376
1377 #[test]
1381 fn test_convert_reverse_flag_cumulative() {
1382 let pool = NDArrayPool::new(1_000_000);
1383 let mut src = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
1384 src.dims[0].reverse = true; let dims_out = vec![NDDimension {
1388 size: 4,
1389 offset: 0,
1390 binning: 1,
1391 reverse: true,
1392 }];
1393 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1394 assert!(!out.dims[0].reverse, "true XOR true must be false");
1395
1396 let dims_out2 = vec![NDDimension {
1398 size: 4,
1399 offset: 0,
1400 binning: 1,
1401 reverse: false,
1402 }];
1403 let out2 = pool.convert(&src, &dims_out2, NDDataType::UInt8).unwrap();
1404 assert!(out2.dims[0].reverse, "false XOR true must be true");
1405 }
1406
1407 #[test]
1409 fn test_alloc_tracks_exact_data_size() {
1410 let pool = NDArrayPool::new(0); let a = pool
1412 .alloc(vec![NDDimension::new(333)], NDDataType::UInt16)
1413 .unwrap();
1414 assert_eq!(a.data_size, 666);
1416 assert_eq!(pool.allocated_bytes(), 666);
1417 }
1418
1419 #[test]
1421 fn test_alloc_strict_max_memory_enforcement() {
1422 let pool = NDArrayPool::new(1000);
1423 let _a = pool
1425 .alloc(vec![NDDimension::new(600)], NDDataType::UInt8)
1426 .unwrap();
1427 assert_eq!(pool.allocated_bytes(), 600);
1428 let r = pool.alloc(vec![NDDimension::new(500)], NDDataType::UInt8);
1430 assert!(r.is_err());
1431 assert!(pool.allocated_bytes() <= 1000);
1432 }
1433
1434 #[test]
1436 fn test_convert_output_is_pool_tracked() {
1437 let pool = NDArrayPool::new(1_000_000);
1438 let src = make_4x4_u8();
1439 let before_alloc = pool.num_alloc_buffers();
1440 let dims_out = vec![
1441 NDDimension {
1442 size: 4,
1443 offset: 0,
1444 binning: 1,
1445 reverse: false,
1446 },
1447 NDDimension {
1448 size: 4,
1449 offset: 0,
1450 binning: 1,
1451 reverse: false,
1452 },
1453 ];
1454 let out = pool.convert(&src, &dims_out, NDDataType::UInt8).unwrap();
1455 assert_eq!(out.pool_id, pool.id());
1456 assert_eq!(out.data_size, 16);
1457 assert_eq!(pool.num_alloc_buffers(), before_alloc + 1);
1459 let bytes_with_out = pool.allocated_bytes();
1460 assert_eq!(bytes_with_out, 16);
1461 pool.release(out);
1463 assert_eq!(pool.num_free_buffers(), 1);
1464 assert_eq!(pool.allocated_bytes(), 16);
1465 }
1466
1467 #[test]
1469 fn test_convert_rejects_compressed_input() {
1470 let pool = NDArrayPool::new(1_000_000);
1471 let mut src = make_4x4_u8();
1472 src.codec = Some(crate::codec::Codec {
1473 name: crate::codec::CodecName::LZ4,
1474 compressed_size: 4,
1475 level: 0,
1476 shuffle: 0,
1477 compressor: 0,
1478 original_data_type: crate::ndarray::NDDataType::UInt8,
1479 });
1480 let dims_out = vec![
1481 NDDimension {
1482 size: 4,
1483 offset: 0,
1484 binning: 1,
1485 reverse: false,
1486 },
1487 NDDimension {
1488 size: 4,
1489 offset: 0,
1490 binning: 1,
1491 reverse: false,
1492 },
1493 ];
1494 assert!(pool.convert(&src, &dims_out, NDDataType::UInt8).is_err());
1495 assert!(pool.convert_type(&src, NDDataType::UInt16).is_err());
1496 }
1497
1498 #[test]
1500 fn test_release_foreign_array_rejected() {
1501 let pool_a = NDArrayPool::new(1_000_000);
1502 let pool_b = NDArrayPool::new(1_000_000);
1503 let arr = pool_a
1504 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
1505 .unwrap();
1506 let bytes_b_before = pool_b.allocated_bytes();
1507 let free_b_before = pool_b.num_free_buffers();
1508 pool_b.release(arr);
1510 assert_eq!(pool_b.allocated_bytes(), bytes_b_before);
1511 assert_eq!(pool_b.num_free_buffers(), free_b_before);
1512 }
1513
1514 #[test]
1516 fn test_release_non_pool_array_rejected() {
1517 let pool = NDArrayPool::new(1_000_000);
1518 let arr = NDArray::new(vec![NDDimension::new(10)], NDDataType::UInt8);
1519 assert_eq!(arr.pool_id, 0);
1520 pool.release(arr);
1521 assert_eq!(pool.num_free_buffers(), 0);
1522 assert_eq!(pool.allocated_bytes(), 0);
1523 }
1524
1525 #[test]
1527 fn test_copy_allocates_and_copies() {
1528 let pool = NDArrayPool::new(1_000_000);
1529 let mut src = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
1530 if let NDDataBuffer::U8(ref mut v) = src.data {
1531 v.copy_from_slice(&[9, 8, 7, 6]);
1532 }
1533 let out = pool.copy(&src, None, true, true, true).unwrap();
1534 assert_eq!(out.pool_id, pool.id());
1535 assert_eq!(out.dims.len(), 1);
1536 if let NDDataBuffer::U8(ref v) = out.data {
1537 assert_eq!(v, &[9, 8, 7, 6]);
1538 } else {
1539 panic!("wrong type");
1540 }
1541 }
1542
1543 #[test]
1545 fn test_pre_allocate_buffers_warms_free_list() {
1546 let pool = NDArrayPool::new(10_000_000);
1547 let template = pool
1548 .alloc(vec![NDDimension::new(256)], NDDataType::UInt16)
1549 .unwrap();
1550 pool.pre_allocate_buffers(&template, 3).unwrap();
1551 assert_eq!(pool.num_free_buffers(), 3);
1552 }
1553
1554 #[test]
1565 fn test_concurrent_reuse_grow_does_not_overshoot_max_memory() {
1566 use std::sync::Arc;
1567 use std::sync::atomic::AtomicBool;
1568 use std::thread;
1569
1570 const N: usize = 16;
1571 for _ in 0..50 {
1573 let pool = Arc::new(NDArrayPool::new(2000));
1580 let mut warm = Vec::with_capacity(N);
1588 for _ in 0..N {
1589 let mut a = pool
1590 .alloc(vec![NDDimension::new(100)], NDDataType::UInt8)
1591 .unwrap();
1592 if let NDDataBuffer::U8(ref mut v) = a.data {
1595 let mut big = Vec::with_capacity(200);
1596 big.resize(100, 0u8);
1597 *v = big;
1598 }
1599 assert_eq!(a.data_size, 100);
1600 assert!(a.data.capacity_bytes() >= 200);
1601 warm.push(a);
1602 }
1603 for a in warm {
1604 pool.release(a);
1605 }
1606 assert_eq!(pool.allocated_bytes(), 1600);
1607 assert_eq!(pool.num_free_buffers(), N as u32);
1608
1609 let overshoot = Arc::new(AtomicBool::new(false));
1610 let mut handles = Vec::new();
1611 for _ in 0..N {
1612 let pool = pool.clone();
1613 let overshoot = overshoot.clone();
1614 handles.push(thread::spawn(move || {
1615 let res = pool.alloc(vec![NDDimension::new(200)], NDDataType::UInt8);
1617 if res.is_ok() && pool.allocated_bytes() > pool.max_memory() as u64 {
1618 overshoot.store(true, Ordering::Relaxed);
1619 }
1620 }));
1621 }
1622 for h in handles {
1623 h.join().unwrap();
1624 }
1625
1626 assert!(
1627 !overshoot.load(Ordering::Relaxed),
1628 "allocated_bytes overshot max_memory during concurrent reuse-grow"
1629 );
1630 assert!(
1631 pool.allocated_bytes() <= pool.max_memory() as u64,
1632 "final allocated_bytes {} > max_memory {}",
1633 pool.allocated_bytes(),
1634 pool.max_memory()
1635 );
1636 }
1637 }
1638
1639 #[test]
1641 fn test_pool_report_nonempty() {
1642 let pool = NDArrayPool::new(1_000_000);
1643 let _ = pool
1644 .alloc(vec![NDDimension::new(10)], NDDataType::UInt8)
1645 .unwrap();
1646 let r = pool.report(10);
1647 assert!(r.contains("NDArrayPool"));
1648 assert!(r.contains("numBuffers"));
1649 }
1650}