Skip to main content

mir_extractor/
memory_profiler.rs

1//! Memory profiling utilities for debugging memory usage during analysis
2//!
3//! Enable with RUSTCOLA_MEMORY_PROFILE=1 environment variable
4
5use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
6use sysinfo::{Pid, System};
7
8static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);
9static PEAK_MEMORY_MB: AtomicU64 = AtomicU64::new(0);
10
11/// Initialize memory profiling based on environment variable
12pub fn init() {
13    if std::env::var("RUSTCOLA_MEMORY_PROFILE").is_ok() {
14        PROFILING_ENABLED.store(true, Ordering::SeqCst);
15        eprintln!("[MEMORY] Profiling enabled");
16    }
17}
18
19/// Check if profiling is enabled
20pub fn is_enabled() -> bool {
21    PROFILING_ENABLED.load(Ordering::SeqCst)
22}
23
24/// Get current process memory usage in MB
25pub fn current_memory_mb() -> f64 {
26    let mut sys = System::new();
27    let pid = Pid::from(std::process::id() as usize);
28    sys.refresh_process(pid);
29
30    if let Some(process) = sys.process(pid) {
31        // process.memory() returns bytes
32        let bytes = process.memory();
33        bytes as f64 / (1024.0 * 1024.0)
34    } else {
35        0.0
36    }
37}
38
39fn update_peak(current_mb: f64) -> f64 {
40    let current_mb_int = current_mb as u64;
41    let old_peak = PEAK_MEMORY_MB.fetch_max(current_mb_int, Ordering::SeqCst);
42    old_peak.max(current_mb_int) as f64
43}
44
45/// Log memory usage at a checkpoint with context
46pub fn checkpoint(label: &str) {
47    if !is_enabled() {
48        return;
49    }
50
51    let mb = current_memory_mb();
52    let peak = update_peak(mb);
53
54    eprintln!("[MEMORY] {}: {:.1} MB (peak: {:.1} MB)", label, mb, peak);
55}
56
57/// Log memory usage with additional context about what's being processed
58pub fn checkpoint_with_context(label: &str, context: &str) {
59    if !is_enabled() {
60        return;
61    }
62
63    let mb = current_memory_mb();
64    let peak = update_peak(mb);
65
66    eprintln!(
67        "[MEMORY] {} [{}]: {:.1} MB (peak: {:.1} MB)",
68        label, context, mb, peak
69    );
70}
71
72/// Track memory delta for a scope
73pub struct MemoryScope {
74    label: String,
75    start_mb: f64,
76}
77
78impl MemoryScope {
79    pub fn new(label: &str) -> Self {
80        let start_mb = if is_enabled() {
81            current_memory_mb()
82        } else {
83            0.0
84        };
85
86        Self {
87            label: label.to_string(),
88            start_mb,
89        }
90    }
91}
92
93impl Drop for MemoryScope {
94    fn drop(&mut self) {
95        if is_enabled() {
96            let end_mb = current_memory_mb();
97            let delta = end_mb - self.start_mb;
98            let sign = if delta >= 0.0 { "+" } else { "" };
99            eprintln!(
100                "[MEMORY] {} completed: {:.1} MB ({}{:.1} MB)",
101                self.label, end_mb, sign, delta
102            );
103        }
104    }
105}
106
107/// Print final memory report
108pub fn final_report() {
109    if !is_enabled() {
110        return;
111    }
112
113    let current = current_memory_mb();
114    let peak = PEAK_MEMORY_MB.load(Ordering::SeqCst) as f64;
115
116    eprintln!("\n[MEMORY] === Final Report ===");
117    eprintln!("[MEMORY] Current: {:.1} MB", current);
118    eprintln!("[MEMORY] Peak:    {:.1} MB", peak);
119    eprintln!("[MEMORY] ====================\n");
120}