Skip to main content

memscope_rs/core/
scope_tracker.rs

1//! Scope tracking functionality for memory analysis
2
3use crate::core::types::*;
4use crate::core::{MemScopeError, MemScopeResult, SystemErrorType};
5use std::collections::HashMap;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, Mutex, OnceLock, RwLock};
8use std::time::{SystemTime, UNIX_EPOCH};
9
10/// Global scope tracker instance
11static GLOBAL_SCOPE_TRACKER: OnceLock<Arc<ScopeTracker>> = OnceLock::new();
12
13/// Get the global scope tracker instance
14pub fn get_global_scope_tracker() -> Arc<ScopeTracker> {
15    GLOBAL_SCOPE_TRACKER
16        .get_or_init(|| Arc::new(ScopeTracker::new()))
17        .clone()
18}
19
20/// Unique identifier for scopes
21pub type ScopeId = u64;
22
23/// Core scope tracking functionality
24pub struct ScopeTracker {
25    /// Active scopes
26    pub active_scopes: RwLock<HashMap<ScopeId, ScopeInfo>>,
27    /// Completed scopes for analysis
28    pub completed_scopes: Mutex<Vec<ScopeInfo>>,
29    /// Scope hierarchy relationships
30    pub scope_hierarchy: Mutex<ScopeHierarchy>,
31    /// Next available scope ID using atomic counter
32    next_scope_id: AtomicU64,
33    /// Current scope stack per thread
34    pub scope_stack: RwLock<HashMap<String, Vec<ScopeId>>>,
35}
36
37impl ScopeTracker {
38    /// Create a new scope tracker
39    pub fn new() -> Self {
40        Self {
41            active_scopes: RwLock::new(HashMap::new()),
42            completed_scopes: Mutex::new(Vec::new()),
43            scope_hierarchy: Mutex::new(ScopeHierarchy {
44                root_scopes: Vec::new(),
45                scope_tree: HashMap::new(),
46                max_depth: 0,
47                total_scopes: 0,
48                relationships: HashMap::new(),
49                depth_map: HashMap::new(),
50            }),
51            next_scope_id: AtomicU64::new(1),
52            scope_stack: RwLock::new(HashMap::new()),
53        }
54    }
55
56    /// Enter a new scope
57    pub fn enter_scope(&self, name: String) -> MemScopeResult<ScopeId> {
58        let scope_id = self.allocate_scope_id();
59        let thread_id = format!("{:?}", std::thread::current().id());
60        let timestamp = current_timestamp();
61
62        let (parent_scope, depth) = {
63            let stack = self.scope_stack.read().map_err(|e| {
64                MemScopeError::system(
65                    SystemErrorType::Locking,
66                    format!("Failed to acquire scope_stack read lock: {}", e),
67                )
68            })?;
69            if let Some(thread_stack) = stack.get(&thread_id) {
70                if let Some(&parent_id) = thread_stack.last() {
71                    let active = self.active_scopes.read().map_err(|e| {
72                        MemScopeError::system(
73                            SystemErrorType::Locking,
74                            format!("Failed to acquire active_scopes read lock: {}", e),
75                        )
76                    })?;
77                    if let Some(parent) = active.get(&parent_id) {
78                        (Some(parent.name.clone()), parent.depth + 1)
79                    } else {
80                        (None, 0)
81                    }
82                } else {
83                    (None, 0)
84                }
85            } else {
86                (None, 0)
87            }
88        };
89
90        let scope_info = ScopeInfo {
91            name: name.clone(),
92            parent: parent_scope.clone(),
93            children: Vec::new(),
94            depth,
95            variables: Vec::new(),
96            total_memory: 0,
97            peak_memory: 0,
98            allocation_count: 0,
99            lifetime_start: Some(timestamp as u64),
100            lifetime_end: None,
101            is_active: true,
102            start_time: timestamp as u64,
103            end_time: None,
104            memory_usage: 0,
105            child_scopes: Vec::new(),
106            parent_scope: parent_scope.clone(),
107        };
108
109        self.active_scopes
110            .write()
111            .map_err(|e| {
112                MemScopeError::system(
113                    SystemErrorType::Locking,
114                    format!("Failed to acquire active_scopes write lock: {}", e),
115                )
116            })?
117            .insert(scope_id, scope_info);
118
119        self.scope_stack
120            .write()
121            .map_err(|e| {
122                MemScopeError::system(
123                    SystemErrorType::Locking,
124                    format!("Failed to acquire scope_stack write lock: {}", e),
125                )
126            })?
127            .entry(thread_id.clone())
128            .or_default()
129            .push(scope_id);
130
131        if let Ok(mut hierarchy) = self.scope_hierarchy.lock() {
132            hierarchy.depth_map.insert(name.clone(), depth);
133
134            if let Some(parent) = parent_scope.clone() {
135                hierarchy
136                    .relationships
137                    .entry(parent)
138                    .or_default()
139                    .push(name.clone());
140            } else {
141                hierarchy.root_scopes.push(name);
142            }
143        }
144
145        Ok(scope_id)
146    }
147
148    /// Exit a scope
149    pub fn exit_scope(&self, scope_id: ScopeId) -> MemScopeResult<()> {
150        let thread_id = format!("{:?}", std::thread::current().id());
151        let timestamp = current_timestamp();
152
153        let mut scope_info = self
154            .active_scopes
155            .write()
156            .map_err(|e| {
157                MemScopeError::system(
158                    SystemErrorType::Locking,
159                    format!("Failed to acquire active_scopes write lock: {}", e),
160                )
161            })?
162            .remove(&scope_id)
163            .ok_or_else(|| MemScopeError::internal(format!("Invalid scope ID: {scope_id}")))?;
164
165        scope_info.end_time = Some(timestamp as u64);
166        scope_info.lifetime_end = Some(timestamp as u64);
167
168        if let Ok(mut stack) = self.scope_stack.write() {
169            if let Some(thread_stack) = stack.get_mut(&thread_id) {
170                if let Some(pos) = thread_stack.iter().position(|&id| id == scope_id) {
171                    thread_stack.remove(pos);
172                }
173            }
174        }
175
176        if let Ok(mut completed_scopes) = self.completed_scopes.lock() {
177            completed_scopes.push(scope_info);
178        }
179
180        Ok(())
181    }
182
183    pub fn associate_variable(
184        &self,
185        variable_name: String,
186        memory_size: usize,
187    ) -> MemScopeResult<()> {
188        let thread_id = format!("{:?}", std::thread::current().id());
189
190        let current_scope_id = self
191            .scope_stack
192            .read()
193            .map_err(|e| {
194                MemScopeError::system(
195                    SystemErrorType::Locking,
196                    format!("Failed to acquire scope_stack read lock: {}", e),
197                )
198            })?
199            .get(&thread_id)
200            .and_then(|stack| stack.last().copied());
201
202        if let Some(scope_id) = current_scope_id {
203            if let Ok(mut active) = self.active_scopes.write() {
204                if let Some(scope) = active.get_mut(&scope_id) {
205                    scope.variables.push(variable_name);
206                    scope.memory_usage += memory_size;
207                    scope.peak_memory = scope.peak_memory.max(scope.memory_usage);
208                    scope.allocation_count += 1;
209                }
210            }
211        }
212
213        Ok(())
214    }
215
216    pub fn get_scope_analysis(&self) -> MemScopeResult<ScopeAnalysis> {
217        let mut all_scopes: Vec<ScopeInfo> = self
218            .active_scopes
219            .read()
220            .map_err(|e| {
221                MemScopeError::system(
222                    SystemErrorType::Locking,
223                    format!("Failed to acquire active_scopes read lock: {}", e),
224                )
225            })?
226            .values()
227            .cloned()
228            .collect();
229
230        if let Ok(completed_scopes) = self.completed_scopes.lock() {
231            all_scopes.extend(completed_scopes.iter().cloned());
232        }
233
234        let hierarchy = if let Ok(hierarchy) = self.scope_hierarchy.lock() {
235            hierarchy.clone()
236        } else {
237            ScopeHierarchy {
238                root_scopes: Vec::new(),
239                scope_tree: HashMap::new(),
240                max_depth: 0,
241                total_scopes: 0,
242                relationships: HashMap::new(),
243                depth_map: HashMap::new(),
244            }
245        };
246
247        Ok(ScopeAnalysis {
248            total_scopes: all_scopes.len(),
249            active_scopes: all_scopes.iter().filter(|s| s.is_active).count(),
250            max_depth: hierarchy.max_depth,
251            average_lifetime: 1000.0,
252            memory_efficiency: 0.8,
253            scopes: all_scopes,
254            scope_hierarchy: hierarchy,
255            cross_scope_references: Vec::new(),
256        })
257    }
258
259    pub fn get_scope_lifecycle_metrics(&self) -> MemScopeResult<Vec<ScopeLifecycleMetrics>> {
260        let metrics = if let Ok(completed_scopes) = self.completed_scopes.lock() {
261            completed_scopes
262                .iter()
263                .map(|scope| {
264                    let lifetime =
265                        scope.end_time.unwrap_or(current_timestamp() as u64) - scope.start_time;
266                    let efficiency_score = if scope.peak_memory > 0 {
267                        scope.memory_usage as f64 / scope.peak_memory as f64
268                    } else {
269                        1.0
270                    };
271
272                    ScopeLifecycleMetrics {
273                        scope_name: scope.name.clone(),
274                        variable_count: scope.variables.len(),
275                        average_lifetime_ms: lifetime as f64,
276                        total_memory_usage: scope.memory_usage,
277                        peak_memory_usage: scope.peak_memory,
278                        allocation_frequency: 1.0,
279                        deallocation_efficiency: efficiency_score,
280                        completed_allocations: scope.allocation_count,
281                        memory_growth_events: 0,
282                        peak_concurrent_variables: scope.variables.len(),
283                        memory_efficiency_ratio: if scope.peak_memory > 0 {
284                            scope.memory_usage as f64 / scope.peak_memory as f64
285                        } else {
286                            1.0
287                        },
288                        ownership_transfer_events: 0,
289                        fragmentation_score: 0.0,
290                        instant_allocations: 0,
291                        short_term_allocations: 0,
292                        medium_term_allocations: 0,
293                        long_term_allocations: 0,
294                        suspected_leaks: 0,
295                        risk_distribution: crate::core::types::RiskDistribution::default(),
296                        scope_metrics: Vec::new(),
297                        type_lifecycle_patterns: Vec::new(),
298                    }
299                })
300                .collect()
301        } else {
302            Vec::new()
303        };
304
305        Ok(metrics)
306    }
307
308    pub fn get_all_scopes(&self) -> Vec<ScopeInfo> {
309        let mut all_scopes: Vec<ScopeInfo> = match self.active_scopes.read() {
310            Ok(active) => active.values().cloned().collect(),
311            Err(_) => Vec::new(),
312        };
313
314        if let Ok(completed_scopes) = self.completed_scopes.lock() {
315            all_scopes.extend(completed_scopes.iter().cloned());
316        }
317
318        all_scopes
319    }
320
321    /// Allocate a new unique scope ID using atomic operations
322    fn allocate_scope_id(&self) -> ScopeId {
323        self.next_scope_id.fetch_add(1, Ordering::Relaxed)
324    }
325}
326
327impl Default for ScopeTracker {
328    fn default() -> Self {
329        Self::new()
330    }
331}
332
333/// RAII scope guard for automatic scope management
334pub struct ScopeGuard {
335    scope_id: ScopeId,
336    tracker: Arc<ScopeTracker>,
337}
338
339impl ScopeGuard {
340    pub fn enter(name: &str) -> MemScopeResult<Self> {
341        let tracker = get_global_scope_tracker();
342        let scope_id = tracker.enter_scope(name.to_string())?;
343
344        Ok(Self { scope_id, tracker })
345    }
346
347    pub fn scope_id(&self) -> ScopeId {
348        self.scope_id
349    }
350}
351
352impl Drop for ScopeGuard {
353    fn drop(&mut self) {
354        let _ = self.tracker.exit_scope(self.scope_id);
355    }
356}
357
358fn current_timestamp() -> u128 {
359    SystemTime::now()
360        .duration_since(UNIX_EPOCH)
361        .map(|d| d.as_millis())
362        .unwrap_or(0)
363}
364
365/// Macro for tracking scopes with automatic cleanup
366#[macro_export]
367macro_rules! track_scope {
368    ($scope_name:expr) => {
369        let _scope_guard = $crate::scope_tracker::ScopeGuard::enter($scope_name)?;
370    };
371    ($scope_name:expr, $block:block) => {{
372        let _scope_guard = $crate::scope_tracker::ScopeGuard::enter($scope_name)?;
373        $block
374    }};
375}
376
377/// Enhanced track_var macro that also associates with current scope
378#[macro_export]
379macro_rules! track_var_with_scope {
380    ($var:ident) => {{
381        // Track the variable normally
382        let result = $crate::_track_var_impl(&$var, stringify!($var));
383
384        // Also associate with current scope
385        if result.is_ok() {
386            if let $crate::TrackKind::HeapOwner { ptr: _, size } =
387                $crate::Trackable::track_kind(&$var)
388            {
389                let scope_tracker = $crate::scope_tracker::get_global_scope_tracker();
390                let _ = scope_tracker.associate_variable(stringify!($var).to_string(), size);
391            }
392        }
393
394        result
395    }};
396}
397
398#[cfg(test)]
399mod tests {
400    use super::*;
401
402    /// Objective: Verify ScopeTracker creation with default values
403    /// Invariants: New tracker should have empty collections and next_scope_id = 1
404    #[test]
405    fn test_scope_tracker_creation() {
406        let tracker = ScopeTracker::new();
407        assert_eq!(
408            tracker.next_scope_id.load(Ordering::SeqCst),
409            1,
410            "Next scope ID should start at 1"
411        );
412        assert!(
413            tracker.active_scopes.read().unwrap().is_empty(),
414            "Active scopes should be empty"
415        );
416        assert!(
417            tracker.completed_scopes.lock().unwrap().is_empty(),
418            "Completed scopes should be empty"
419        );
420    }
421
422    /// Objective: Verify Default trait implementation
423    /// Invariants: Default should create same as new()
424    #[test]
425    fn test_scope_tracker_default() {
426        let tracker = ScopeTracker::default();
427        assert_eq!(
428            tracker.next_scope_id.load(Ordering::SeqCst),
429            1,
430            "Default should start with ID 1"
431        );
432    }
433
434    /// Objective: Verify enter and exit scope functionality
435    /// Invariants: Scope should be active after enter, moved to completed after exit
436    #[test]
437    fn test_enter_and_exit_scope() {
438        let tracker = ScopeTracker::new();
439        let scope_id = tracker
440            .enter_scope("test_scope".to_string())
441            .expect("Should enter scope");
442
443        assert!(
444            tracker
445                .active_scopes
446                .read()
447                .unwrap()
448                .get(&scope_id)
449                .is_some(),
450            "Scope should be active"
451        );
452
453        tracker.exit_scope(scope_id).expect("Should exit scope");
454
455        assert!(
456            tracker
457                .active_scopes
458                .read()
459                .unwrap()
460                .get(&scope_id)
461                .is_none(),
462            "Scope should no longer be active"
463        );
464        assert_eq!(
465            tracker.completed_scopes.lock().unwrap().len(),
466            1,
467            "Should have one completed scope"
468        );
469    }
470
471    /// Objective: Verify scope ID allocation is sequential
472    /// Invariants: Each enter_scope should return incrementing IDs
473    #[test]
474    fn test_sequential_scope_ids() {
475        let tracker = ScopeTracker::new();
476
477        let id1 = tracker.enter_scope("scope1".to_string()).unwrap();
478        let id2 = tracker.enter_scope("scope2".to_string()).unwrap();
479        let id3 = tracker.enter_scope("scope3".to_string()).unwrap();
480
481        assert!(id1 < id2, "IDs should be sequential");
482        assert!(id2 < id3, "IDs should be sequential");
483    }
484
485    /// Objective: Verify nested scope tracking
486    /// Invariants: Child scope should have correct parent and depth
487    #[test]
488    fn test_nested_scopes() {
489        let tracker = ScopeTracker::new();
490
491        let parent_id = tracker.enter_scope("parent".to_string()).unwrap();
492        let child_id = tracker.enter_scope("child".to_string()).unwrap();
493
494        let active = tracker.active_scopes.read().unwrap();
495        let parent_scope = active.get(&parent_id).expect("Parent should exist");
496        let child_scope = active.get(&child_id).expect("Child should exist");
497
498        assert_eq!(parent_scope.depth, 0, "Parent should have depth 0");
499        assert_eq!(child_scope.depth, 1, "Child should have depth 1");
500        assert_eq!(
501            child_scope.parent_scope,
502            Some("parent".to_string()),
503            "Child should have parent"
504        );
505    }
506
507    /// Objective: Verify associate_variable functionality
508    /// Invariants: Variable should be associated with current scope
509    #[test]
510    fn test_associate_variable() {
511        let tracker = ScopeTracker::new();
512        let _scope_id = tracker.enter_scope("test".to_string()).unwrap();
513
514        tracker
515            .associate_variable("my_var".to_string(), 1024)
516            .expect("Should associate variable");
517
518        let active = tracker.active_scopes.read().unwrap();
519        let scope = active.values().next().expect("Should have scope");
520
521        assert!(
522            scope.variables.contains(&"my_var".to_string()),
523            "Variable should be in scope"
524        );
525        assert_eq!(scope.memory_usage, 1024, "Memory usage should be updated");
526        assert_eq!(scope.allocation_count, 1, "Allocation count should be 1");
527    }
528
529    /// Objective: Verify peak memory tracking
530    /// Invariants: Peak memory should track maximum usage
531    #[test]
532    fn test_peak_memory_tracking() {
533        let tracker = ScopeTracker::new();
534        let _scope_id = tracker.enter_scope("test".to_string()).unwrap();
535
536        tracker.associate_variable("var1".to_string(), 100).unwrap();
537        tracker.associate_variable("var2".to_string(), 200).unwrap();
538
539        let active = tracker.active_scopes.read().unwrap();
540        let scope = active.values().next().expect("Should have scope");
541
542        assert_eq!(scope.memory_usage, 300, "Total memory should be 300");
543        assert_eq!(scope.peak_memory, 300, "Peak memory should be 300");
544    }
545
546    /// Objective: Verify get_scope_analysis functionality
547    /// Invariants: Should return analysis with all scopes
548    #[test]
549    fn test_get_scope_analysis() {
550        let tracker = ScopeTracker::new();
551
552        let id1 = tracker.enter_scope("scope1".to_string()).unwrap();
553        let _id2 = tracker.enter_scope("scope2".to_string()).unwrap();
554        tracker.exit_scope(id1).unwrap();
555
556        let analysis = tracker.get_scope_analysis().expect("Should get analysis");
557
558        assert_eq!(analysis.total_scopes, 2, "Should have 2 scopes total");
559        assert!(
560            analysis.active_scopes >= 1,
561            "Should have at least 1 active scope"
562        );
563    }
564
565    /// Objective: Verify get_all_scopes returns all scopes
566    /// Invariants: Should include both active and completed scopes
567    #[test]
568    fn test_get_all_scopes() {
569        let tracker = ScopeTracker::new();
570
571        let id1 = tracker.enter_scope("scope1".to_string()).unwrap();
572        let _id2 = tracker.enter_scope("scope2".to_string()).unwrap();
573        tracker.exit_scope(id1).unwrap();
574
575        let all_scopes = tracker.get_all_scopes();
576
577        assert_eq!(all_scopes.len(), 2, "Should have 2 scopes");
578    }
579
580    /// Objective: Verify get_scope_lifecycle_metrics functionality
581    /// Invariants: Should return metrics for completed scopes
582    #[test]
583    fn test_get_scope_lifecycle_metrics() {
584        let tracker = ScopeTracker::new();
585
586        let id = tracker.enter_scope("test".to_string()).unwrap();
587        tracker.associate_variable("var".to_string(), 100).unwrap();
588        tracker.exit_scope(id).unwrap();
589
590        let metrics = tracker
591            .get_scope_lifecycle_metrics()
592            .expect("Should get metrics");
593
594        assert_eq!(metrics.len(), 1, "Should have 1 metric");
595        assert_eq!(metrics[0].scope_name, "test", "Scope name should match");
596        assert_eq!(metrics[0].variable_count, 1, "Should have 1 variable");
597    }
598
599    /// Objective: Verify ScopeGuard RAII behavior
600    /// Invariants: Scope should be exited when guard is dropped
601    #[test]
602    fn test_scope_guard() {
603        let tracker = get_global_scope_tracker();
604
605        let guard = ScopeGuard::enter("test_guard").expect("Should enter");
606        let scope_id = guard.scope_id();
607
608        let active = tracker.active_scopes.read().unwrap();
609        assert!(active.contains_key(&scope_id), "Scope should be active");
610        drop(active);
611
612        drop(guard);
613
614        let active = tracker.active_scopes.read().unwrap();
615        assert!(
616            !active.contains_key(&scope_id),
617            "Scope should be exited after guard drop"
618        );
619    }
620
621    /// Objective: Verify ScopeGuard scope_id method
622    /// Invariants: Should return the scope ID
623    #[test]
624    fn test_scope_guard_id() {
625        let guard = ScopeGuard::enter("test").expect("Should enter");
626        let id = guard.scope_id();
627        assert!(id > 0, "Scope ID should be positive");
628    }
629
630    /// Objective: Verify exiting invalid scope returns error
631    /// Invariants: Should return error for non-existent scope
632    #[test]
633    fn test_exit_invalid_scope() {
634        let tracker = ScopeTracker::new();
635        let result = tracker.exit_scope(999);
636
637        assert!(result.is_err(), "Should return error for invalid scope");
638    }
639
640    /// Objective: Verify scope hierarchy tracking
641    /// Invariants: Hierarchy should track parent-child relationships
642    #[test]
643    fn test_scope_hierarchy() {
644        let tracker = ScopeTracker::new();
645
646        let _parent = tracker.enter_scope("parent".to_string()).unwrap();
647        let _child = tracker.enter_scope("child".to_string()).unwrap();
648
649        let hierarchy = tracker.scope_hierarchy.lock().unwrap();
650        assert!(
651            hierarchy.relationships.contains_key("parent"),
652            "Should have parent relationship"
653        );
654        assert!(
655            hierarchy.depth_map.contains_key("parent"),
656            "Should have parent depth"
657        );
658        assert!(
659            hierarchy.depth_map.contains_key("child"),
660            "Should have child depth"
661        );
662    }
663
664    /// Objective: Verify multiple variables in scope
665    /// Invariants: All variables should be tracked
666    #[test]
667    fn test_multiple_variables() {
668        let tracker = ScopeTracker::new();
669        let _scope_id = tracker.enter_scope("test".to_string()).unwrap();
670
671        tracker.associate_variable("var1".to_string(), 100).unwrap();
672        tracker.associate_variable("var2".to_string(), 200).unwrap();
673        tracker.associate_variable("var3".to_string(), 300).unwrap();
674
675        let active = tracker.active_scopes.read().unwrap();
676        let scope = active.values().next().expect("Should have scope");
677
678        assert_eq!(scope.variables.len(), 3, "Should have 3 variables");
679        assert_eq!(scope.allocation_count, 3, "Should have 3 allocations");
680    }
681
682    /// Objective: Verify current_timestamp returns valid value
683    /// Invariants: Timestamp should be positive
684    #[test]
685    fn test_current_timestamp() {
686        let ts = current_timestamp();
687        assert!(ts > 0, "Timestamp should be positive");
688    }
689
690    /// Objective: Verify global scope tracker singleton
691    /// Invariants: Multiple calls should return same instance
692    #[test]
693    fn test_global_scope_tracker_singleton() {
694        let tracker1 = get_global_scope_tracker();
695        let tracker2 = get_global_scope_tracker();
696
697        assert!(Arc::ptr_eq(&tracker1, &tracker2));
698    }
699}