1use scirs2_core::ndarray::{Array1, Array2, ArrayViewMut1, ArrayViewMut2};
8use std::alloc::{alloc, dealloc, Layout};
9use std::collections::{HashMap, VecDeque};
10use std::ptr::NonNull;
11use std::sync::{
12 atomic::{AtomicUsize, Ordering},
13 Arc, Mutex, RwLock,
14};
15use thiserror::Error;
16
17static NEXT_POOL_ID: AtomicUsize = AtomicUsize::new(1);
19
20#[derive(Error, Debug)]
22pub enum MemoryPoolError {
23 #[error("Pool exhausted: no available blocks of size {0}")]
24 PoolExhausted(usize),
25 #[error("Invalid block size: {0}")]
26 InvalidBlockSize(usize),
27 #[error("Allocation failed: {0}")]
28 AllocationFailed(String),
29 #[error("Block not found: {0:?}")]
30 BlockNotFound(usize),
31 #[error("Double free detected: {0:?}")]
32 DoubleFree(usize),
33 #[error("Pool corruption detected: {0}")]
34 PoolCorruption(String),
35 #[error("Alignment error: required {required}, got {actual}")]
36 AlignmentError { required: usize, actual: usize },
37}
38
39pub type MemoryPoolResult<T> = Result<T, MemoryPoolError>;
40
41#[derive(Debug, Clone)]
43struct MemoryBlock {
44 ptr: NonNull<u8>,
45 size: usize,
46 layout: Layout,
47 allocated_at: std::time::Instant,
48 #[allow(dead_code)] pool_id: usize,
50 magic: u64, }
52
53const MAGIC_VALUE: u64 = 0xDEADBEEFCAFEBABE;
54
55impl MemoryBlock {
56 fn new(size: usize, align: usize, pool_id: usize) -> MemoryPoolResult<Self> {
57 let layout = Layout::from_size_align(size, align)
58 .map_err(|e| MemoryPoolError::AllocationFailed(e.to_string()))?;
59
60 let ptr = unsafe { alloc(layout) };
61 if ptr.is_null() {
62 return Err(MemoryPoolError::AllocationFailed(
63 "System allocation failed".to_string(),
64 ));
65 }
66
67 Ok(Self {
68 ptr: NonNull::new(ptr).ok_or_else(|| {
69 MemoryPoolError::AllocationFailed("NonNull creation failed".to_string())
70 })?,
71 size,
72 layout,
73 allocated_at: std::time::Instant::now(),
74 pool_id,
75 magic: MAGIC_VALUE,
76 })
77 }
78
79 fn is_valid(&self) -> bool {
80 self.magic == MAGIC_VALUE
81 }
82
83 fn invalidate(&mut self) {
84 self.magic = 0;
85 }
86
87 fn as_slice_mut(&mut self) -> &mut [u8] {
88 unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.size) }
89 }
90
91 #[allow(dead_code)] fn as_f64_slice_mut(&mut self) -> MemoryPoolResult<&mut [f64]> {
93 if !self.size.is_multiple_of(std::mem::size_of::<f64>()) {
94 return Err(MemoryPoolError::InvalidBlockSize(self.size));
95 }
96
97 let len = self.size / std::mem::size_of::<f64>();
98 Ok(unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut f64, len) })
99 }
100}
101
102impl Drop for MemoryBlock {
103 fn drop(&mut self) {
104 if self.is_valid() {
105 unsafe {
106 dealloc(self.ptr.as_ptr(), self.layout);
107 }
108 self.invalidate();
109 }
110 }
111}
112
113unsafe impl Send for MemoryBlock {}
118unsafe impl Sync for MemoryBlock {}
119
120#[derive(Debug, Clone)]
122pub struct MemoryPoolConfig {
123 pub initial_blocks_per_size: usize,
125 pub max_blocks_per_size: usize,
127 pub growth_factor: f64,
129 pub alignment: usize,
131 pub enable_reuse: bool,
133 pub max_idle_time: std::time::Duration,
135 pub enable_stats: bool,
137}
138
139impl Default for MemoryPoolConfig {
140 fn default() -> Self {
141 Self {
142 initial_blocks_per_size: 4,
143 max_blocks_per_size: 64,
144 growth_factor: 1.5,
145 alignment: 64, enable_reuse: true,
147 max_idle_time: std::time::Duration::from_secs(300), enable_stats: true,
149 }
150 }
151}
152
153#[derive(Debug, Clone, Default)]
155pub struct MemoryPoolStats {
156 pub total_allocations: usize,
157 pub total_deallocations: usize,
158 pub current_allocations: usize,
159 pub peak_allocations: usize,
160 pub total_bytes_allocated: usize,
161 pub current_bytes_allocated: usize,
162 pub peak_bytes_allocated: usize,
163 pub pool_hits: usize,
164 pub pool_misses: usize,
165 pub blocks_created: usize,
166 pub blocks_reused: usize,
167 pub blocks_released: usize,
168}
169
170impl MemoryPoolStats {
171 pub fn hit_rate(&self) -> f64 {
172 if self.pool_hits + self.pool_misses == 0 {
173 0.0
174 } else {
175 self.pool_hits as f64 / (self.pool_hits + self.pool_misses) as f64
176 }
177 }
178
179 pub fn reuse_rate(&self) -> f64 {
180 if self.blocks_created == 0 {
181 0.0
182 } else {
183 self.blocks_reused as f64 / self.blocks_created as f64
184 }
185 }
186}
187
188struct SizeBucket {
190 size: usize,
191 available: VecDeque<MemoryBlock>,
192 allocated: HashMap<usize, MemoryBlock>,
193 total_created: usize,
194}
195
196impl SizeBucket {
197 fn new(size: usize) -> Self {
198 Self {
199 size,
200 available: VecDeque::new(),
201 allocated: HashMap::new(),
202 total_created: 0,
203 }
204 }
205
206 fn allocate(
207 &mut self,
208 pool_id: usize,
209 config: &MemoryPoolConfig,
210 ) -> MemoryPoolResult<(usize, *mut u8, bool)> {
211 let (block, was_reused) = if let Some(block) = self.available.pop_front() {
212 (block, true)
213 } else {
214 let block = MemoryBlock::new(self.size, config.alignment, pool_id)?;
216 self.total_created += 1;
217 (block, false)
218 };
219
220 let ptr = block.ptr.as_ptr();
221 let block_id = ptr as usize;
222 self.allocated.insert(block_id, block);
223
224 Ok((block_id, ptr, was_reused))
225 }
226
227 fn deallocate(&mut self, block_id: usize, enable_reuse: bool) -> MemoryPoolResult<()> {
228 if let Some(mut block) = self.allocated.remove(&block_id) {
229 if !block.is_valid() {
230 return Err(MemoryPoolError::PoolCorruption(
231 "Block magic value corrupted".to_string(),
232 ));
233 }
234
235 if enable_reuse {
236 let slice = block.as_slice_mut();
238 slice.fill(0);
239
240 self.available.push_back(block);
241 } else {
242 }
250
251 Ok(())
252 } else {
253 Err(MemoryPoolError::BlockNotFound(block_id))
254 }
255 }
256
257 fn cleanup_idle(&mut self, max_idle_time: std::time::Duration) -> usize {
258 let now = std::time::Instant::now();
259 let mut removed = 0;
260
261 self.available.retain(|block| {
262 if now.duration_since(block.allocated_at) > max_idle_time {
263 removed += 1;
264 false
265 } else {
266 true
267 }
268 });
269
270 removed
271 }
272}
273
274pub struct MemoryPool {
276 config: MemoryPoolConfig,
277 buckets: RwLock<HashMap<usize, SizeBucket>>,
278 stats: RwLock<MemoryPoolStats>,
279 pool_id: usize,
280 next_cleanup: Mutex<std::time::Instant>,
281}
282
283impl MemoryPool {
284 pub fn new(config: MemoryPoolConfig) -> Self {
286 let max_idle_time = config.max_idle_time;
287 Self {
288 config,
289 buckets: RwLock::new(HashMap::new()),
290 stats: RwLock::new(MemoryPoolStats::default()),
291 pool_id: NEXT_POOL_ID.fetch_add(1, Ordering::Relaxed),
292 next_cleanup: Mutex::new(std::time::Instant::now() + max_idle_time),
293 }
294 }
295
296 pub fn pool_id(&self) -> usize {
298 self.pool_id
299 }
300
301 pub fn allocate(&self, size: usize) -> MemoryPoolResult<PooledBlock> {
303 if size == 0 {
304 return Err(MemoryPoolError::InvalidBlockSize(size));
305 }
306
307 let bucket_size = size.next_power_of_two().max(64);
309
310 let (block_id, ptr, was_reused) = {
311 let mut buckets = self
312 .buckets
313 .write()
314 .unwrap_or_else(|poisoned| poisoned.into_inner());
315 let bucket = buckets
316 .entry(bucket_size)
317 .or_insert_with(|| SizeBucket::new(bucket_size));
318
319 if bucket.allocated.len() >= self.config.max_blocks_per_size {
321 return Err(MemoryPoolError::PoolExhausted(bucket_size));
322 }
323
324 bucket.allocate(self.pool_id, &self.config)?
325 };
326
327 if self.config.enable_stats {
329 let mut stats = self
330 .stats
331 .write()
332 .unwrap_or_else(|poisoned| poisoned.into_inner());
333 stats.total_allocations += 1;
334 stats.current_allocations += 1;
335 stats.peak_allocations = stats.peak_allocations.max(stats.current_allocations);
336 stats.total_bytes_allocated += bucket_size;
337 stats.current_bytes_allocated += bucket_size;
338 stats.peak_bytes_allocated = stats
339 .peak_bytes_allocated
340 .max(stats.current_bytes_allocated);
341
342 if was_reused {
343 stats.blocks_reused += 1;
344 } else {
345 stats.blocks_created += 1;
346 }
347
348 if bucket_size != size {
349 stats.pool_hits += 1;
350 } else {
351 stats.pool_misses += 1;
352 }
353 }
354
355 self.maybe_cleanup();
357
358 Ok(PooledBlock {
359 ptr,
360 size: bucket_size,
361 block_id,
362 pool: self as *const Self,
363 })
364 }
365
366 pub fn deallocate(&self, block_id: usize, size: usize) -> MemoryPoolResult<()> {
368 let bucket_size = size.next_power_of_two().max(64);
369
370 {
371 let mut buckets = self
372 .buckets
373 .write()
374 .unwrap_or_else(|poisoned| poisoned.into_inner());
375 if let Some(bucket) = buckets.get_mut(&bucket_size) {
376 bucket.deallocate(block_id, self.config.enable_reuse)?;
377 } else {
378 return Err(MemoryPoolError::BlockNotFound(block_id));
379 }
380 }
381
382 if self.config.enable_stats {
384 let mut stats = self
385 .stats
386 .write()
387 .unwrap_or_else(|poisoned| poisoned.into_inner());
388 stats.total_deallocations += 1;
389 stats.current_allocations = stats.current_allocations.saturating_sub(1);
390 stats.current_bytes_allocated =
391 stats.current_bytes_allocated.saturating_sub(bucket_size);
392 }
393
394 Ok(())
395 }
396
397 pub fn stats(&self) -> MemoryPoolStats {
399 if self.config.enable_stats {
400 self.stats
401 .read()
402 .unwrap_or_else(|poisoned| poisoned.into_inner())
403 .clone()
404 } else {
405 MemoryPoolStats::default()
406 }
407 }
408
409 pub fn cleanup(&self) -> usize {
411 let mut total_removed = 0;
412 let mut buckets = self
413 .buckets
414 .write()
415 .unwrap_or_else(|poisoned| poisoned.into_inner());
416
417 for bucket in buckets.values_mut() {
418 total_removed += bucket.cleanup_idle(self.config.max_idle_time);
419 }
420
421 *self
423 .next_cleanup
424 .lock()
425 .unwrap_or_else(|poisoned| poisoned.into_inner()) =
426 std::time::Instant::now() + self.config.max_idle_time;
427
428 total_removed
429 }
430
431 fn maybe_cleanup(&self) {
433 let now = std::time::Instant::now();
434 let should_cleanup = {
435 let next_cleanup = self
436 .next_cleanup
437 .lock()
438 .unwrap_or_else(|poisoned| poisoned.into_inner());
439 now >= *next_cleanup
440 };
441
442 if should_cleanup {
443 self.cleanup();
444 }
445 }
446
447 pub fn allocate_array2(&self, nrows: usize, ncols: usize) -> MemoryPoolResult<PooledArray2> {
449 let size = nrows * ncols * std::mem::size_of::<f64>();
450 let block = self.allocate(size)?;
451
452 Ok(PooledArray2 {
453 block,
454 nrows,
455 ncols,
456 })
457 }
458
459 pub fn allocate_array1(&self, len: usize) -> MemoryPoolResult<PooledArray1> {
461 let size = len * std::mem::size_of::<f64>();
462 let block = self.allocate(size)?;
463
464 Ok(PooledArray1 { block, len })
465 }
466
467 pub fn bucket_info(&self) -> Vec<(usize, usize, usize)> {
469 let buckets = self
470 .buckets
471 .read()
472 .unwrap_or_else(|poisoned| poisoned.into_inner());
473 buckets
474 .iter()
475 .map(|(size, bucket)| (*size, bucket.available.len(), bucket.allocated.len()))
476 .collect()
477 }
478}
479
480impl Default for MemoryPool {
481 fn default() -> Self {
482 Self::new(MemoryPoolConfig::default())
483 }
484}
485
486pub struct PooledBlock {
488 ptr: *mut u8,
489 size: usize,
490 block_id: usize,
491 pool: *const MemoryPool,
492}
493
494impl PooledBlock {
495 pub fn as_slice_mut(&mut self) -> &mut [u8] {
497 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size) }
498 }
499
500 pub fn size(&self) -> usize {
502 self.size
503 }
504
505 pub fn as_f64_slice_mut(&mut self) -> MemoryPoolResult<&mut [f64]> {
507 if !self.size.is_multiple_of(std::mem::size_of::<f64>()) {
508 return Err(MemoryPoolError::InvalidBlockSize(self.size));
509 }
510
511 let len = self.size / std::mem::size_of::<f64>();
512 Ok(unsafe { std::slice::from_raw_parts_mut(self.ptr as *mut f64, len) })
513 }
514}
515
516impl Drop for PooledBlock {
517 fn drop(&mut self) {
518 if !self.pool.is_null() {
519 unsafe {
520 if let Err(e) = (*self.pool).deallocate(self.block_id, self.size) {
521 eprintln!("Error deallocating pooled block: {}", e);
522 }
523 }
524 }
525 }
526}
527
528unsafe impl Send for PooledBlock {}
529unsafe impl Sync for PooledBlock {}
530
531pub struct PooledArray2 {
533 block: PooledBlock,
534 nrows: usize,
535 ncols: usize,
536}
537
538impl PooledArray2 {
539 pub fn dim(&self) -> (usize, usize) {
541 (self.nrows, self.ncols)
542 }
543
544 pub fn view_mut(&mut self) -> MemoryPoolResult<ArrayViewMut2<'_, f64>> {
546 let slice = self.block.as_f64_slice_mut()?;
547 ArrayViewMut2::from_shape((self.nrows, self.ncols), slice)
548 .map_err(|e| MemoryPoolError::AllocationFailed(format!("Shape error: {}", e)))
549 }
550
551 pub fn to_array(mut self) -> MemoryPoolResult<Array2<f64>> {
553 let slice = self.block.as_f64_slice_mut()?;
554 let total_elements = self.nrows * self.ncols;
555 let data = slice[..total_elements].to_vec();
557 let array = Array2::from_shape_vec((self.nrows, self.ncols), data).map_err(|e| {
558 MemoryPoolError::AllocationFailed(format!("Array creation error: {}", e))
559 })?;
560 Ok(array)
561 }
562}
563
564pub struct PooledArray1 {
566 block: PooledBlock,
567 len: usize,
568}
569
570impl PooledArray1 {
571 pub fn len(&self) -> usize {
573 self.len
574 }
575
576 pub fn is_empty(&self) -> bool {
578 self.len == 0
579 }
580
581 pub fn view_mut(&mut self) -> MemoryPoolResult<ArrayViewMut1<'_, f64>> {
583 let slice = self.block.as_f64_slice_mut()?;
584 ArrayViewMut1::from_shape(self.len, slice)
585 .map_err(|e| MemoryPoolError::AllocationFailed(format!("Shape error: {}", e)))
586 }
587
588 pub fn to_array(mut self) -> MemoryPoolResult<Array1<f64>> {
590 let slice = self.block.as_f64_slice_mut()?;
591 let array = Array1::from_vec(slice[..self.len].to_vec());
593 Ok(array)
594 }
595}
596
597pub type SharedMemoryPool = Arc<MemoryPool>;
599
600pub fn create_shared_pool(config: MemoryPoolConfig) -> SharedMemoryPool {
602 Arc::new(MemoryPool::new(config))
603}
604
605lazy_static::lazy_static! {
606 static ref GLOBAL_POOL: SharedMemoryPool = create_shared_pool(MemoryPoolConfig::default());
608}
609
610pub fn global_pool() -> &'static SharedMemoryPool {
612 &GLOBAL_POOL
613}
614
615pub fn allocate_global(size: usize) -> MemoryPoolResult<PooledBlock> {
617 global_pool().allocate(size)
618}
619
620pub fn allocate_array2_global(nrows: usize, ncols: usize) -> MemoryPoolResult<PooledArray2> {
622 global_pool().allocate_array2(nrows, ncols)
623}
624
625pub fn allocate_array1_global(len: usize) -> MemoryPoolResult<PooledArray1> {
627 global_pool().allocate_array1(len)
628}
629
630#[allow(non_snake_case)]
631#[cfg(test)]
632mod tests {
633 use super::*;
634
635 #[test]
636 fn test_memory_pool_basic() -> MemoryPoolResult<()> {
637 let pool = MemoryPool::new(MemoryPoolConfig::default());
638
639 let mut block = pool.allocate(1024)?;
641 assert_eq!(block.size(), 1024);
642
643 let slice = block.as_slice_mut();
645 slice[0] = 42;
646 slice[1023] = 24;
647
648 assert_eq!(slice[0], 42);
650 assert_eq!(slice[1023], 24);
651
652 Ok(())
653 }
654
655 #[test]
656 fn test_memory_pool_reuse() -> MemoryPoolResult<()> {
657 let config = MemoryPoolConfig {
658 enable_reuse: true,
659 ..Default::default()
660 };
661 let pool = MemoryPool::new(config);
662
663 {
665 let _block1 = pool.allocate(512)?;
666 } let _block2 = pool.allocate(512)?;
670
671 let stats = pool.stats();
672 assert!(stats.blocks_reused > 0 || stats.pool_hits > 0);
673
674 Ok(())
675 }
676
677 #[test]
678 fn test_pooled_array2() -> MemoryPoolResult<()> {
679 let pool = MemoryPool::new(MemoryPoolConfig::default());
680
681 let mut array = pool.allocate_array2(100, 50)?;
682 assert_eq!(array.dim(), (100, 50));
683
684 {
686 let mut view = array.view_mut()?;
687 view.fill(std::f64::consts::PI);
688 }
689
690 let owned = array.to_array()?;
692 assert_eq!(owned.dim(), (100, 50));
693 assert!((owned[[0, 0]] - std::f64::consts::PI).abs() < f64::EPSILON);
694
695 Ok(())
696 }
697
698 #[test]
699 fn test_pooled_array1() -> MemoryPoolResult<()> {
700 let pool = MemoryPool::new(MemoryPoolConfig::default());
701
702 let mut array = pool.allocate_array1(1000)?;
703 assert_eq!(array.len(), 1000);
704
705 {
707 let mut view = array.view_mut()?;
708 view.fill(2.71);
709 }
710
711 let owned = array.to_array()?;
713 assert_eq!(owned.len(), 1000);
714 assert!((owned[0] - 2.71).abs() < f64::EPSILON);
715
716 Ok(())
717 }
718
719 #[test]
720 fn test_memory_pool_stats() -> MemoryPoolResult<()> {
721 let config = MemoryPoolConfig {
722 enable_stats: true,
723 ..Default::default()
724 };
725 let pool = MemoryPool::new(config);
726
727 let initial_stats = pool.stats();
729 assert_eq!(initial_stats.total_allocations, 0);
730
731 let _block1 = pool.allocate(256)?;
733 let _block2 = pool.allocate(512)?;
734
735 let stats = pool.stats();
736 assert_eq!(stats.total_allocations, 2);
737 assert_eq!(stats.current_allocations, 2);
738 assert!(stats.current_bytes_allocated > 0);
739
740 Ok(())
741 }
742
743 #[test]
744 fn test_global_pool() -> MemoryPoolResult<()> {
745 let mut block = allocate_global(128)?;
746 assert_eq!(block.size(), 128);
747
748 let slice = block.as_slice_mut();
749 slice[0] = 255;
750 assert_eq!(slice[0], 255);
751
752 let array2 = allocate_array2_global(10, 20)?;
753 assert_eq!(array2.dim(), (10, 20));
754
755 let array1 = allocate_array1_global(100)?;
756 assert_eq!(array1.len(), 100);
757
758 Ok(())
759 }
760
761 #[test]
762 fn test_pool_exhaustion() {
763 let config = MemoryPoolConfig {
764 max_blocks_per_size: 2,
765 enable_reuse: false,
766 ..Default::default()
767 };
768 let pool = MemoryPool::new(config);
769
770 let _block1 = pool.allocate(64).expect("operation should succeed");
772 let _block2 = pool.allocate(64).expect("operation should succeed");
773
774 assert!(pool.allocate(64).is_err());
776 }
777
778 #[test]
779 fn test_cleanup() -> MemoryPoolResult<()> {
780 let config = MemoryPoolConfig {
781 max_idle_time: std::time::Duration::from_millis(1),
782 enable_reuse: true,
783 ..Default::default()
784 };
785 let pool = MemoryPool::new(config);
786
787 {
789 let _block = pool.allocate(256)?;
790 }
791
792 std::thread::sleep(std::time::Duration::from_millis(10));
794
795 let removed = pool.cleanup();
797 assert!(removed > 0);
798
799 Ok(())
800 }
801
802 #[test]
803 fn test_pool_ids_are_unique_and_nonzero() {
804 let p1 = MemoryPool::new(MemoryPoolConfig::default());
806 let p2 = MemoryPool::new(MemoryPoolConfig::default());
807 let p3 = MemoryPool::new(MemoryPoolConfig::default());
808
809 assert_ne!(p1.pool_id(), 0, "pool_id must not be zero");
810 assert_ne!(p1.pool_id(), p2.pool_id(), "pool IDs must be unique");
811 assert_ne!(p2.pool_id(), p3.pool_id(), "pool IDs must be unique");
812 assert_ne!(p1.pool_id(), p3.pool_id(), "pool IDs must be unique");
813 }
814
815 #[test]
816 fn test_bucket_alignment() -> MemoryPoolResult<()> {
817 let pool = MemoryPool::new(MemoryPoolConfig::default());
818
819 let block1 = pool.allocate(100)?; let block2 = pool.allocate(200)?; assert!(block1.size() >= 100);
824 assert!(block2.size() >= 200);
825
826 assert!(block1.size().is_power_of_two());
828 assert!(block2.size().is_power_of_two());
829
830 Ok(())
831 }
832}