Skip to main content

sklears_utils/
memory.rs

1//! Memory management utilities for high-performance ML workloads
2//!
3//! This module provides memory management utilities including custom allocators,
4//! memory pools, leak detection, memory-mapped file utilities, bounds checking,
5//! and safe memory management helpers.
6
7use crate::{UtilsError, UtilsResult};
8use std::alloc::{GlobalAlloc, Layout};
9use std::collections::HashMap;
10use std::fs::File;
11use std::sync::{Arc, Mutex, RwLock};
12use std::time::{Duration, Instant};
13
14/// Custom allocator that tracks memory usage
15pub struct TrackingAllocator<A: GlobalAlloc> {
16    inner: A,
17    stats: Arc<RwLock<AllocationStats>>,
18}
19
20/// Allocation statistics
21#[derive(Debug, Clone, Default)]
22pub struct AllocationStats {
23    pub total_allocated: u64,
24    pub total_deallocated: u64,
25    pub current_allocated: u64,
26    pub peak_allocated: u64,
27    pub allocation_count: u64,
28    pub deallocation_count: u64,
29    pub leak_count: u64,
30}
31
32impl<A: GlobalAlloc> TrackingAllocator<A> {
33    pub fn new(inner: A) -> Self {
34        Self {
35            inner,
36            stats: Arc::new(RwLock::new(AllocationStats::default())),
37        }
38    }
39
40    pub fn stats(&self) -> AllocationStats {
41        self.stats.read().expect("operation should succeed").clone()
42    }
43
44    pub fn reset_stats(&self) {
45        let mut stats = self.stats.write().expect("operation should succeed");
46        *stats = AllocationStats::default();
47    }
48}
49
50unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
51    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
52        let ptr = self.inner.alloc(layout);
53        if !ptr.is_null() {
54            let mut stats = self.stats.write().expect("operation should succeed");
55            stats.total_allocated += layout.size() as u64;
56            stats.current_allocated += layout.size() as u64;
57            stats.allocation_count += 1;
58            if stats.current_allocated > stats.peak_allocated {
59                stats.peak_allocated = stats.current_allocated;
60            }
61        }
62        ptr
63    }
64
65    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
66        self.inner.dealloc(ptr, layout);
67        let mut stats = self.stats.write().expect("operation should succeed");
68        stats.total_deallocated += layout.size() as u64;
69        stats.current_allocated = stats.current_allocated.saturating_sub(layout.size() as u64);
70        stats.deallocation_count += 1;
71    }
72}
73
74/// Memory pool for efficient allocation of fixed-size objects
75pub struct MemoryPool<T> {
76    blocks: Vec<Box<[T]>>,
77    free_list: Vec<*mut T>,
78    block_size: usize,
79    stats: AllocationStats,
80}
81
82impl<T: Default + Clone> MemoryPool<T> {
83    pub fn new(block_size: usize) -> Self {
84        Self {
85            blocks: Vec::new(),
86            free_list: Vec::new(),
87            block_size,
88            stats: AllocationStats::default(),
89        }
90    }
91
92    pub fn allocate(&mut self) -> Option<&mut T> {
93        if self.free_list.is_empty() {
94            self.add_block();
95        }
96
97        if let Some(ptr) = self.free_list.pop() {
98            self.stats.allocation_count += 1;
99            self.stats.current_allocated += std::mem::size_of::<T>() as u64;
100            unsafe { Some(&mut *ptr) }
101        } else {
102            None
103        }
104    }
105
106    pub fn deallocate(&mut self, item: &mut T) {
107        let ptr = item as *mut T;
108        self.free_list.push(ptr);
109        self.stats.deallocation_count += 1;
110        self.stats.current_allocated = self
111            .stats
112            .current_allocated
113            .saturating_sub(std::mem::size_of::<T>() as u64);
114    }
115
116    fn add_block(&mut self) {
117        let block = vec![T::default(); self.block_size].into_boxed_slice();
118        // Move the block into its final resting place *before* handing out any
119        // raw pointers into its heap buffer. Passing a `Box<[T]>` by value into
120        // `Vec::push` counts, under Stacked Borrows, as forming a fresh Unique
121        // reborrow over the box's *entire* pointee (a `Box` is treated as
122        // `noalias`, much like `&mut T`). If pointers into the buffer were
123        // taken beforehand (as the old code did via `block.iter_mut()`), that
124        // move retroactively invalidates them, which is genuine UB (confirmed
125        // by Miri). Growing the outer `Vec<Box<[T]>>` itself later on is fine:
126        // that only memcpy's the 16-byte fat pointers, never the box's
127        // contents, so pointers derived here remain valid for the pool's
128        // lifetime.
129        self.blocks.push(block);
130        let stored_block = self.blocks.last_mut().expect("block was just pushed above");
131        for item in stored_block.iter_mut() {
132            self.free_list.push(item as *mut T);
133        }
134        self.stats.total_allocated += (self.block_size * std::mem::size_of::<T>()) as u64;
135    }
136
137    pub fn stats(&self) -> &AllocationStats {
138        &self.stats
139    }
140
141    pub fn capacity(&self) -> usize {
142        self.blocks.len() * self.block_size
143    }
144
145    pub fn used(&self) -> usize {
146        self.capacity() - self.free_list.len()
147    }
148}
149
150/// Memory leak detector
151pub struct LeakDetector {
152    allocations: Arc<Mutex<HashMap<usize, AllocationInfo>>>,
153    enabled: bool,
154}
155
156#[derive(Debug, Clone)]
157pub struct AllocationInfo {
158    pub size: usize,
159    pub timestamp: Instant,
160    pub backtrace: String,
161}
162
163impl LeakDetector {
164    pub fn new() -> Self {
165        Self {
166            allocations: Arc::new(Mutex::new(HashMap::new())),
167            enabled: true,
168        }
169    }
170
171    pub fn enable(&mut self) {
172        self.enabled = true;
173    }
174
175    pub fn disable(&mut self) {
176        self.enabled = false;
177    }
178
179    pub fn track_allocation(&self, ptr: *mut u8, size: usize) {
180        if !self.enabled {
181            return;
182        }
183
184        let mut allocations = self.allocations.lock().expect("operation should succeed");
185        allocations.insert(
186            ptr as usize,
187            AllocationInfo {
188                size,
189                timestamp: Instant::now(),
190                backtrace: format!("Allocation at {ptr:p}"), // In real implementation, use backtrace crate
191            },
192        );
193    }
194
195    pub fn track_deallocation(&self, ptr: *mut u8) {
196        if !self.enabled {
197            return;
198        }
199
200        let mut allocations = self.allocations.lock().expect("operation should succeed");
201        allocations.remove(&(ptr as usize));
202    }
203
204    pub fn check_leaks(&self) -> Vec<AllocationInfo> {
205        let allocations = self.allocations.lock().expect("operation should succeed");
206        allocations.values().cloned().collect()
207    }
208
209    pub fn check_leaks_older_than(&self, duration: Duration) -> Vec<AllocationInfo> {
210        let allocations = self.allocations.lock().expect("operation should succeed");
211        let now = Instant::now();
212        allocations
213            .values()
214            .filter(|info| now.duration_since(info.timestamp) > duration)
215            .cloned()
216            .collect()
217    }
218
219    pub fn total_leaked_bytes(&self) -> usize {
220        let allocations = self.allocations.lock().expect("operation should succeed");
221        allocations.values().map(|info| info.size).sum()
222    }
223
224    pub fn clear(&self) {
225        let mut allocations = self.allocations.lock().expect("operation should succeed");
226        allocations.clear();
227    }
228}
229
230impl Default for LeakDetector {
231    fn default() -> Self {
232        Self::new()
233    }
234}
235
236/// Memory-mapped file utilities
237pub struct MemoryMappedFile {
238    #[allow(dead_code)]
239    file: File,
240    ptr: *mut u8,
241    size: usize,
242}
243
244impl MemoryMappedFile {
245    #[cfg(unix)]
246    pub fn new(file: File, writable: bool) -> Result<Self, std::io::Error> {
247        use std::os::unix::io::AsRawFd;
248
249        let size = file.metadata()?.len() as usize;
250        let prot = if writable {
251            libc::PROT_READ | libc::PROT_WRITE
252        } else {
253            libc::PROT_READ
254        };
255
256        let ptr = unsafe {
257            libc::mmap(
258                std::ptr::null_mut(),
259                size,
260                prot,
261                libc::MAP_SHARED,
262                file.as_raw_fd(),
263                0,
264            )
265        };
266
267        if ptr == libc::MAP_FAILED {
268            return Err(std::io::Error::last_os_error());
269        }
270
271        Ok(Self {
272            file,
273            ptr: ptr as *mut u8,
274            size,
275        })
276    }
277
278    #[cfg(windows)]
279    pub fn new(file: File, writable: bool) -> Result<Self, std::io::Error> {
280        use std::os::windows::io::AsRawHandle;
281        use winapi::um::handleapi::CloseHandle;
282        use winapi::um::memoryapi::{
283            CreateFileMappingW, MapViewOfFile, FILE_MAP_READ, FILE_MAP_WRITE,
284        };
285        use winapi::um::winnt::{PAGE_READONLY, PAGE_READWRITE};
286
287        let size = file.metadata()?.len() as usize;
288        let protect = if writable {
289            PAGE_READWRITE
290        } else {
291            PAGE_READONLY
292        };
293        let access = if writable {
294            FILE_MAP_WRITE
295        } else {
296            FILE_MAP_READ
297        };
298
299        let mapping = unsafe {
300            CreateFileMappingW(
301                file.as_raw_handle() as _,
302                std::ptr::null_mut(),
303                protect,
304                0,
305                0,
306                std::ptr::null(),
307            )
308        };
309
310        if mapping.is_null() {
311            return Err(std::io::Error::last_os_error());
312        }
313
314        let ptr = unsafe { MapViewOfFile(mapping, access, 0, 0, 0) };
315        unsafe { CloseHandle(mapping) };
316
317        if ptr.is_null() {
318            return Err(std::io::Error::last_os_error());
319        }
320
321        Ok(Self {
322            file,
323            ptr: ptr as *mut u8,
324            size,
325        })
326    }
327
328    #[cfg(not(any(unix, windows)))]
329    pub fn new(_file: File, _writable: bool) -> Result<Self, std::io::Error> {
330        Err(std::io::Error::new(
331            std::io::ErrorKind::Unsupported,
332            "Memory mapping not supported on this platform",
333        ))
334    }
335
336    pub fn as_slice(&self) -> &[u8] {
337        unsafe { std::slice::from_raw_parts(self.ptr, self.size) }
338    }
339
340    pub fn as_mut_slice(&mut self) -> &mut [u8] {
341        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size) }
342    }
343
344    pub fn size(&self) -> usize {
345        self.size
346    }
347}
348
349impl Drop for MemoryMappedFile {
350    fn drop(&mut self) {
351        if !self.ptr.is_null() {
352            #[cfg(unix)]
353            unsafe {
354                libc::munmap(self.ptr as *mut libc::c_void, self.size);
355            }
356
357            #[cfg(windows)]
358            unsafe {
359                winapi::um::memoryapi::UnmapViewOfFile(self.ptr as *mut winapi::ctypes::c_void);
360            }
361        }
362    }
363}
364
365unsafe impl Send for MemoryMappedFile {}
366unsafe impl Sync for MemoryMappedFile {}
367
368/// Garbage collection helper for reference counting
369pub struct GcHelper<T> {
370    data: Arc<T>,
371    weak_refs: Arc<Mutex<Vec<std::sync::Weak<T>>>>,
372}
373
374impl<T> GcHelper<T> {
375    pub fn new(data: T) -> Self {
376        Self {
377            data: Arc::new(data),
378            weak_refs: Arc::new(Mutex::new(Vec::new())),
379        }
380    }
381
382    pub fn get_ref(&self) -> Arc<T> {
383        self.data.clone()
384    }
385
386    pub fn get_weak_ref(&self) -> std::sync::Weak<T> {
387        let weak = Arc::downgrade(&self.data);
388        let mut refs = self.weak_refs.lock().expect("operation should succeed");
389        refs.push(weak.clone());
390        weak
391    }
392
393    pub fn collect_garbage(&self) {
394        let mut refs = self.weak_refs.lock().expect("operation should succeed");
395        refs.retain(|weak_ref| weak_ref.upgrade().is_some());
396    }
397
398    pub fn ref_count(&self) -> usize {
399        Arc::strong_count(&self.data)
400    }
401
402    pub fn weak_ref_count(&self) -> usize {
403        let refs = self.weak_refs.lock().expect("operation should succeed");
404        refs.len()
405    }
406}
407
408impl<T> Clone for GcHelper<T> {
409    fn clone(&self) -> Self {
410        Self {
411            data: self.data.clone(),
412            weak_refs: self.weak_refs.clone(),
413        }
414    }
415}
416
417/// Memory usage monitor
418pub struct MemoryMonitor {
419    start_time: Instant,
420    peak_memory: u64,
421    current_memory: u64,
422    samples: Vec<(Instant, u64)>,
423}
424
425impl MemoryMonitor {
426    pub fn new() -> Self {
427        Self {
428            start_time: Instant::now(),
429            peak_memory: 0,
430            current_memory: 0,
431            samples: Vec::new(),
432        }
433    }
434
435    pub fn update(&mut self, memory_usage: u64) {
436        self.current_memory = memory_usage;
437        if memory_usage > self.peak_memory {
438            self.peak_memory = memory_usage;
439        }
440        self.samples.push((Instant::now(), memory_usage));
441    }
442
443    pub fn peak_memory(&self) -> u64 {
444        self.peak_memory
445    }
446
447    pub fn current_memory(&self) -> u64 {
448        self.current_memory
449    }
450
451    pub fn average_memory(&self) -> f64 {
452        if self.samples.is_empty() {
453            return 0.0;
454        }
455        let sum: u64 = self.samples.iter().map(|(_, mem)| *mem).sum();
456        sum as f64 / self.samples.len() as f64
457    }
458
459    pub fn memory_over_time(&self) -> &[(Instant, u64)] {
460        &self.samples
461    }
462
463    pub fn duration(&self) -> Duration {
464        Instant::now().duration_since(self.start_time)
465    }
466}
467
468impl Default for MemoryMonitor {
469    fn default() -> Self {
470        Self::new()
471    }
472}
473
474// ===== BOUNDS CHECKING HELPERS =====
475
476/// Safe wrapper around vectors with bounds checking
477#[derive(Debug, Clone)]
478pub struct SafeVec<T> {
479    data: Vec<T>,
480    bounds_check: bool,
481}
482
483impl<T> SafeVec<T> {
484    /// Create a new safe vector with bounds checking enabled
485    pub fn new() -> Self {
486        Self {
487            data: Vec::new(),
488            bounds_check: true,
489        }
490    }
491
492    /// Create a safe vector with specified capacity
493    pub fn with_capacity(capacity: usize) -> Self {
494        Self {
495            data: Vec::with_capacity(capacity),
496            bounds_check: true,
497        }
498    }
499
500    /// Create from existing vector
501    pub fn from_vec(vec: Vec<T>) -> Self {
502        Self {
503            data: vec,
504            bounds_check: true,
505        }
506    }
507
508    /// Disable bounds checking for performance (unsafe)
509    pub fn disable_bounds_check(mut self) -> Self {
510        self.bounds_check = false;
511        self
512    }
513
514    /// Safe get with bounds checking
515    pub fn get(&self, index: usize) -> UtilsResult<&T> {
516        if self.bounds_check && index >= self.data.len() {
517            return Err(UtilsError::InvalidParameter(format!(
518                "Index {} out of bounds for vector of length {}",
519                index,
520                self.data.len()
521            )));
522        }
523        self.data
524            .get(index)
525            .ok_or_else(|| UtilsError::InvalidParameter(format!("Index {index} out of bounds")))
526    }
527
528    /// Safe mutable get with bounds checking
529    pub fn get_mut(&mut self, index: usize) -> UtilsResult<&mut T> {
530        if self.bounds_check && index >= self.data.len() {
531            return Err(UtilsError::InvalidParameter(format!(
532                "Index {} out of bounds for vector of length {}",
533                index,
534                self.data.len()
535            )));
536        }
537        let len = self.data.len();
538        self.data.get_mut(index).ok_or_else(|| {
539            UtilsError::InvalidParameter(format!(
540                "Index {index} out of bounds for vector of length {len}"
541            ))
542        })
543    }
544
545    /// Safe slice access
546    pub fn safe_slice(&self, start: usize, end: usize) -> UtilsResult<&[T]> {
547        if self.bounds_check {
548            if start > end {
549                return Err(UtilsError::InvalidParameter(
550                    "Start index cannot be greater than end index".to_string(),
551                ));
552            }
553            if end > self.data.len() {
554                return Err(UtilsError::InvalidParameter(format!(
555                    "End index {end} out of bounds for vector of length {}",
556                    self.data.len()
557                )));
558            }
559        }
560        Ok(&self.data[start..end])
561    }
562
563    /// Push element
564    pub fn push(&mut self, item: T) {
565        self.data.push(item);
566    }
567
568    /// Pop element
569    pub fn pop(&mut self) -> Option<T> {
570        self.data.pop()
571    }
572
573    /// Get length
574    pub fn len(&self) -> usize {
575        self.data.len()
576    }
577
578    /// Check if empty
579    pub fn is_empty(&self) -> bool {
580        self.data.is_empty()
581    }
582
583    /// Get capacity
584    pub fn capacity(&self) -> usize {
585        self.data.capacity()
586    }
587
588    /// Reserve additional capacity
589    pub fn reserve(&mut self, additional: usize) {
590        self.data.reserve(additional);
591    }
592
593    /// Access underlying vector (unsafe)
594    ///
595    /// # Safety
596    ///
597    /// This function exposes the underlying `Vec<T>` directly, bypassing all
598    /// bounds checking and overflow protection mechanisms. The caller must
599    /// ensure that any modifications to the returned vector do not violate
600    /// the buffer's safety guarantees and internal invariants.
601    pub unsafe fn as_vec(&self) -> &Vec<T> {
602        &self.data
603    }
604
605    /// Convert to vector
606    pub fn into_vec(self) -> Vec<T> {
607        self.data
608    }
609}
610
611impl<T> Default for SafeVec<T> {
612    fn default() -> Self {
613        Self::new()
614    }
615}
616
617/// Safe buffer with automatic bounds checking and buffer overflow protection
618#[derive(Debug, Clone)]
619pub struct SafeBuffer<T> {
620    data: Vec<T>,
621    capacity: usize,
622    size: usize,
623    overflow_protection: bool,
624}
625
626impl<T: Clone> SafeBuffer<T> {
627    /// Create a new safe buffer with fixed capacity
628    pub fn new(capacity: usize, default_value: T) -> Self {
629        Self {
630            data: vec![default_value; capacity],
631            capacity,
632            size: 0,
633            overflow_protection: true,
634        }
635    }
636
637    /// Write to buffer with bounds checking
638    pub fn write(&mut self, index: usize, value: T) -> UtilsResult<()> {
639        if self.overflow_protection && index >= self.capacity {
640            return Err(UtilsError::InvalidParameter(format!(
641                "Buffer overflow: index {} exceeds capacity {}",
642                index, self.capacity
643            )));
644        }
645
646        if index < self.data.len() {
647            self.data[index] = value;
648            self.size = self.size.max(index + 1);
649            Ok(())
650        } else {
651            Err(UtilsError::InvalidParameter(format!(
652                "Index {} out of bounds for buffer of capacity {}",
653                index, self.capacity
654            )))
655        }
656    }
657
658    /// Read from buffer with bounds checking
659    pub fn read(&self, index: usize) -> UtilsResult<&T> {
660        if index >= self.size {
661            return Err(UtilsError::InvalidParameter(format!(
662                "Index {} out of bounds for buffer of size {}",
663                index, self.size
664            )));
665        }
666
667        self.data
668            .get(index)
669            .ok_or_else(|| UtilsError::InvalidParameter(format!("Index {index} out of bounds")))
670    }
671
672    /// Append to buffer
673    pub fn append(&mut self, value: T) -> UtilsResult<()> {
674        if self.size >= self.capacity {
675            return Err(UtilsError::InvalidParameter(
676                "Buffer overflow: cannot append to full buffer".to_string(),
677            ));
678        }
679
680        self.data[self.size] = value;
681        self.size += 1;
682        Ok(())
683    }
684
685    /// Get current size
686    pub fn size(&self) -> usize {
687        self.size
688    }
689
690    /// Get capacity
691    pub fn capacity(&self) -> usize {
692        self.capacity
693    }
694
695    /// Check if buffer is full
696    pub fn is_full(&self) -> bool {
697        self.size >= self.capacity
698    }
699
700    /// Clear buffer
701    pub fn clear(&mut self) {
702        self.size = 0;
703    }
704
705    /// Disable overflow protection (unsafe)
706    ///
707    /// # Safety
708    ///
709    /// This function disables the buffer's overflow protection mechanism,
710    /// allowing operations that could potentially lead to buffer overflows
711    /// or memory corruption. The caller must ensure that all subsequent
712    /// operations on the buffer are within proper bounds.
713    pub unsafe fn disable_overflow_protection(&mut self) {
714        self.overflow_protection = false;
715    }
716}
717
718/// Memory-safe smart pointer with reference counting and automatic cleanup
719pub struct SafePtr<T> {
720    data: Arc<RwLock<Option<T>>>,
721    cleanup_fn: Option<Box<dyn Fn() + Send + Sync>>,
722}
723
724impl<T> std::fmt::Debug for SafePtr<T> {
725    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
726        f.debug_struct("SafePtr")
727            .field("data", &"Arc<RwLock<Option<T>>>")
728            .field(
729                "cleanup_fn",
730                &self.cleanup_fn.as_ref().map(|_| "Some(cleanup_fn)"),
731            )
732            .finish()
733    }
734}
735
736impl<T> SafePtr<T> {
737    /// Create a new safe pointer
738    pub fn new(value: T) -> Self {
739        Self {
740            data: Arc::new(RwLock::new(Some(value))),
741            cleanup_fn: None,
742        }
743    }
744
745    /// Create with cleanup function
746    pub fn with_cleanup<F>(value: T, cleanup: F) -> Self
747    where
748        F: Fn() + Send + Sync + 'static,
749    {
750        Self {
751            data: Arc::new(RwLock::new(Some(value))),
752            cleanup_fn: Some(Box::new(cleanup)),
753        }
754    }
755
756    /// Try to read the value
757    pub fn try_read(&self) -> UtilsResult<std::sync::RwLockReadGuard<'_, Option<T>>> {
758        self.data
759            .read()
760            .map_err(|e| UtilsError::InvalidParameter(format!("Failed to acquire read lock: {e}")))
761    }
762
763    /// Try to write the value
764    pub fn try_write(&self) -> UtilsResult<std::sync::RwLockWriteGuard<'_, Option<T>>> {
765        self.data
766            .write()
767            .map_err(|e| UtilsError::InvalidParameter(format!("Failed to acquire write lock: {e}")))
768    }
769
770    /// Check if the pointer is still valid
771    pub fn is_valid(&self) -> bool {
772        if let Ok(guard) = self.data.read() {
773            guard.is_some()
774        } else {
775            false
776        }
777    }
778
779    /// Take the value, leaving None
780    pub fn take(&self) -> UtilsResult<Option<T>> {
781        let mut guard = self.try_write()?;
782        Ok(guard.take())
783    }
784
785    /// Get reference count
786    pub fn ref_count(&self) -> usize {
787        Arc::strong_count(&self.data)
788    }
789}
790
791impl<T> Clone for SafePtr<T> {
792    fn clone(&self) -> Self {
793        Self {
794            data: self.data.clone(),
795            cleanup_fn: None, // Cleanup functions are not cloned
796        }
797    }
798}
799
800impl<T> Drop for SafePtr<T> {
801    fn drop(&mut self) {
802        // Run cleanup function if this is the last reference
803        if Arc::strong_count(&self.data) == 1 {
804            if let Some(cleanup) = &self.cleanup_fn {
805                cleanup();
806            }
807        }
808    }
809}
810
811/// Memory alignment utilities
812pub struct MemoryAlignment;
813
814impl MemoryAlignment {
815    /// Check if a pointer is aligned to the specified boundary
816    pub fn is_aligned<T>(ptr: *const T, alignment: usize) -> bool {
817        (ptr as usize).is_multiple_of(alignment)
818    }
819
820    /// Get the alignment of a type
821    pub fn alignment_of<T>() -> usize {
822        std::mem::align_of::<T>()
823    }
824
825    /// Calculate aligned size
826    pub fn aligned_size(size: usize, alignment: usize) -> usize {
827        (size + alignment - 1) & !(alignment - 1)
828    }
829
830    /// Create aligned memory layout
831    pub fn aligned_layout(
832        size: usize,
833        alignment: usize,
834    ) -> Result<Layout, std::alloc::LayoutError> {
835        Layout::from_size_align(size, alignment)
836    }
837}
838
839/// Stack-based memory guard for automatic cleanup
840pub struct StackGuard<F: FnOnce()> {
841    cleanup: Option<F>,
842}
843
844impl<F: FnOnce()> StackGuard<F> {
845    /// Create a new stack guard with cleanup function
846    pub fn new(cleanup: F) -> Self {
847        Self {
848            cleanup: Some(cleanup),
849        }
850    }
851
852    /// Manually trigger cleanup (consumes the guard)
853    pub fn cleanup(mut self) {
854        if let Some(cleanup) = self.cleanup.take() {
855            cleanup();
856        }
857    }
858}
859
860impl<F: FnOnce()> Drop for StackGuard<F> {
861    fn drop(&mut self) {
862        if let Some(cleanup) = self.cleanup.take() {
863            cleanup();
864        }
865    }
866}
867
868/// Macro for creating stack guards
869#[macro_export]
870macro_rules! defer {
871    ($cleanup:expr) => {
872        let _guard = $crate::memory::StackGuard::new(|| $cleanup);
873    };
874}
875
876/// Memory validation utilities
877pub struct MemoryValidator;
878
879impl MemoryValidator {
880    /// Validate that a memory range is accessible
881    ///
882    /// # Safety
883    ///
884    /// This function performs raw pointer arithmetic and validation. The caller
885    /// must ensure that the provided pointer was obtained through safe means and
886    /// that the memory range \[ptr, ptr + count * `size_of::<T>`()\] is within valid
887    /// allocated memory boundaries. Incorrect usage can lead to undefined behavior.
888    pub unsafe fn validate_range<T>(ptr: *const T, count: usize) -> UtilsResult<()> {
889        if ptr.is_null() {
890            return Err(UtilsError::InvalidParameter("Null pointer".to_string()));
891        }
892
893        // Basic overflow check
894        let end_ptr = unsafe { ptr.add(count) };
895        if end_ptr < ptr {
896            return Err(UtilsError::InvalidParameter("Pointer overflow".to_string()));
897        }
898
899        Ok(())
900    }
901
902    /// Validate memory alignment
903    pub fn validate_alignment<T>(ptr: *const T, required_alignment: usize) -> UtilsResult<()> {
904        if !MemoryAlignment::is_aligned(ptr, required_alignment) {
905            return Err(UtilsError::InvalidParameter(format!(
906                "Pointer not aligned to {required_alignment} byte boundary"
907            )));
908        }
909        Ok(())
910    }
911
912    /// Validate buffer bounds
913    pub fn validate_buffer_access(
914        buffer_size: usize,
915        offset: usize,
916        access_size: usize,
917    ) -> UtilsResult<()> {
918        if offset >= buffer_size {
919            return Err(UtilsError::InvalidParameter(format!(
920                "Offset {offset} exceeds buffer size {buffer_size}"
921            )));
922        }
923
924        if offset + access_size > buffer_size {
925            return Err(UtilsError::InvalidParameter(format!(
926                "Access range {}..{} exceeds buffer size {}",
927                offset,
928                offset + access_size,
929                buffer_size
930            )));
931        }
932
933        Ok(())
934    }
935}
936
937#[allow(non_snake_case)]
938#[cfg(test)]
939mod tests {
940    use super::*;
941    use std::alloc::System;
942
943    #[test]
944    fn test_tracking_allocator_stats() {
945        let allocator = TrackingAllocator::new(System);
946        let initial_stats = allocator.stats();
947        assert_eq!(initial_stats.allocation_count, 0);
948        assert_eq!(initial_stats.current_allocated, 0);
949    }
950
951    #[test]
952    fn test_memory_pool() {
953        let mut pool: MemoryPool<u64> = MemoryPool::new(10);
954
955        {
956            let item1 = pool.allocate().expect("operation should succeed");
957            *item1 = 42;
958            assert_eq!(*item1, 42);
959        }
960        assert_eq!(pool.used(), 1);
961
962        {
963            let item2 = pool.allocate().expect("operation should succeed");
964            *item2 = 84;
965            assert_eq!(*item2, 84);
966        }
967        assert_eq!(pool.used(), 2);
968
969        // Note: In a real implementation, deallocate would need the item reference
970        // For this test, we'll just check that the pool tracks allocation correctly
971        assert_eq!(pool.capacity(), 10);
972    }
973
974    #[test]
975    fn test_leak_detector() {
976        let detector = LeakDetector::new();
977        let ptr = Box::into_raw(Box::new(42u64));
978
979        detector.track_allocation(ptr as *mut u8, 8);
980        assert_eq!(detector.total_leaked_bytes(), 8);
981
982        detector.track_deallocation(ptr as *mut u8);
983        assert_eq!(detector.total_leaked_bytes(), 0);
984
985        unsafe { drop(Box::from_raw(ptr)) };
986    }
987
988    #[test]
989    fn test_gc_helper() {
990        let gc = GcHelper::new(42u64);
991        assert_eq!(gc.ref_count(), 1);
992
993        let strong_ref = gc.get_ref();
994        assert_eq!(gc.ref_count(), 2);
995        assert_eq!(*strong_ref, 42);
996
997        let weak_ref = gc.get_weak_ref();
998        assert!(weak_ref.upgrade().is_some());
999
1000        drop(strong_ref);
1001        assert_eq!(gc.ref_count(), 1);
1002    }
1003
1004    #[test]
1005    fn test_memory_monitor() {
1006        let mut monitor = MemoryMonitor::new();
1007        assert_eq!(monitor.peak_memory(), 0);
1008        assert_eq!(monitor.current_memory(), 0);
1009
1010        monitor.update(1024);
1011        assert_eq!(monitor.peak_memory(), 1024);
1012        assert_eq!(monitor.current_memory(), 1024);
1013
1014        monitor.update(512);
1015        assert_eq!(monitor.peak_memory(), 1024);
1016        assert_eq!(monitor.current_memory(), 512);
1017
1018        monitor.update(2048);
1019        assert_eq!(monitor.peak_memory(), 2048);
1020        assert_eq!(monitor.current_memory(), 2048);
1021
1022        assert_eq!(monitor.average_memory(), (1024.0 + 512.0 + 2048.0) / 3.0);
1023    }
1024
1025    #[test]
1026    fn test_safe_vec() {
1027        let mut safe_vec = SafeVec::new();
1028        safe_vec.push(1);
1029        safe_vec.push(2);
1030        safe_vec.push(3);
1031
1032        // Valid access
1033        assert_eq!(*safe_vec.get(0).expect("operation should succeed"), 1);
1034        assert_eq!(*safe_vec.get(2).expect("operation should succeed"), 3);
1035
1036        // Invalid access
1037        assert!(safe_vec.get(5).is_err());
1038
1039        // Safe slice
1040        let slice = safe_vec.safe_slice(1, 3).expect("operation should succeed");
1041        assert_eq!(slice, &[2, 3]);
1042
1043        // Invalid slice
1044        assert!(safe_vec.safe_slice(2, 5).is_err());
1045        assert!(safe_vec.safe_slice(3, 2).is_err());
1046    }
1047
1048    #[test]
1049    fn test_safe_buffer() {
1050        let mut buffer = SafeBuffer::new(5, 0);
1051
1052        // Write and read
1053        buffer.write(0, 42).expect("operation should succeed");
1054        buffer.write(1, 84).expect("operation should succeed");
1055
1056        assert_eq!(*buffer.read(0).expect("operation should succeed"), 42);
1057        assert_eq!(*buffer.read(1).expect("operation should succeed"), 84);
1058        assert_eq!(buffer.size(), 2);
1059
1060        // Append
1061        buffer.append(100).expect("operation should succeed");
1062        buffer.append(200).expect("operation should succeed");
1063        buffer.append(300).expect("operation should succeed");
1064
1065        assert!(buffer.is_full());
1066        assert!(buffer.append(400).is_err()); // Should fail - buffer full
1067
1068        // Buffer overflow protection
1069        assert!(buffer.write(10, 500).is_err()); // Should fail - out of bounds
1070    }
1071
1072    #[test]
1073    fn test_safe_ptr() {
1074        let ptr = SafePtr::new(42);
1075        assert!(ptr.is_valid());
1076        assert_eq!(ptr.ref_count(), 1);
1077
1078        // Clone increases ref count
1079        let _ptr2 = ptr.clone();
1080        assert_eq!(ptr.ref_count(), 2);
1081
1082        // Read value
1083        {
1084            let guard = ptr.try_read().expect("operation should succeed");
1085            assert_eq!(*guard, Some(42));
1086        }
1087
1088        // Take value
1089        let value = ptr.take().expect("operation should succeed");
1090        assert_eq!(value, Some(42));
1091        assert!(!ptr.is_valid());
1092    }
1093
1094    #[test]
1095    fn test_memory_alignment() {
1096        // Test alignment checking
1097        let data = 42u64;
1098        let ptr = &data as *const u64;
1099
1100        assert!(MemoryAlignment::is_aligned(ptr, 8)); // u64 should be 8-byte aligned
1101        assert_eq!(MemoryAlignment::alignment_of::<u64>(), 8);
1102
1103        // Test aligned size calculation
1104        assert_eq!(MemoryAlignment::aligned_size(10, 8), 16);
1105        assert_eq!(MemoryAlignment::aligned_size(16, 8), 16);
1106        assert_eq!(MemoryAlignment::aligned_size(17, 8), 24);
1107    }
1108
1109    #[test]
1110    fn test_stack_guard() {
1111        use std::sync::Arc;
1112
1113        let cleanup_called = Arc::new(Mutex::new(false));
1114        let cleanup_called_clone = cleanup_called.clone();
1115
1116        {
1117            let _guard = StackGuard::new(|| {
1118                *cleanup_called_clone
1119                    .lock()
1120                    .expect("operation should succeed") = true;
1121            });
1122
1123            // Cleanup not called yet
1124            assert!(!*cleanup_called.lock().expect("operation should succeed"));
1125        } // Guard drops here
1126
1127        // Cleanup should be called now
1128        assert!(*cleanup_called.lock().expect("operation should succeed"));
1129    }
1130
1131    #[test]
1132    fn test_memory_validator() {
1133        // Test null pointer validation
1134        let null_ptr: *const u8 = std::ptr::null();
1135        assert!(unsafe { MemoryValidator::validate_range(null_ptr, 10) }.is_err());
1136
1137        // Test valid pointer
1138        let data = [1u8, 2, 3, 4, 5];
1139        let ptr = data.as_ptr();
1140        assert!(unsafe { MemoryValidator::validate_range(ptr, 5) }.is_ok());
1141
1142        // Test alignment validation
1143        let aligned_ptr = &42u64 as *const u64;
1144        assert!(MemoryValidator::validate_alignment(aligned_ptr, 8).is_ok());
1145
1146        // Test buffer access validation
1147        assert!(MemoryValidator::validate_buffer_access(10, 0, 5).is_ok());
1148        assert!(MemoryValidator::validate_buffer_access(10, 5, 5).is_ok());
1149        assert!(MemoryValidator::validate_buffer_access(10, 10, 1).is_err()); // Offset too large
1150        assert!(MemoryValidator::validate_buffer_access(10, 8, 5).is_err()); // Access exceeds buffer
1151    }
1152
1153    #[test]
1154    fn test_defer_macro() {
1155        use std::sync::Arc;
1156
1157        let cleanup_called = Arc::new(Mutex::new(false));
1158        let cleanup_called_clone = cleanup_called.clone();
1159
1160        {
1161            defer!({
1162                *cleanup_called_clone
1163                    .lock()
1164                    .expect("operation should succeed") = true;
1165            });
1166
1167            // Cleanup not called yet
1168            assert!(!*cleanup_called.lock().expect("operation should succeed"));
1169        } // Deferred cleanup happens here
1170
1171        // Cleanup should be called now
1172        assert!(*cleanup_called.lock().expect("operation should succeed"));
1173    }
1174}