memscope-rs 0.2.6

//! Allocation reconstruction from raw memory events.
//!
//! This module provides the `rebuild_allocations_from_events` function which
//! converts the raw `MemoryEvent` stream into the per-allocation records the
//! dashboard consumes.
//!
//! # Single source of truth
//!
//! All binding matching — address reuse, release matching, reallocation and byte
//! accounting — is delegated to [`crate::snapshot::bindings::BindingLedger`], the
//! same ledger that builds [`crate::snapshot::MemorySnapshot`]. This module only
//! *projects* ledger bindings into the dashboard DTO plus the dashboard-only
//! decorations (smart-pointer classification, clone links, container nodes).
//!
//! # Known limits of the legacy event model
//!
//! * Explicit binding serials match exactly; legacy events use address/size heuristics.
//! * Migrating reallocations carry `new_ptr` and retain the original identity.
//! * `AllocationInfo::thread_id` (`std::thread::ThreadId`) cannot be rebuilt from a
//!   `u64`. `thread_id_u64` is authoritative; the `ThreadId` field keeps whatever
//!   the observing thread is and must not be used for aggregation.

use crate::capture::types::AllocationInfo;
use crate::event_store::{MemoryEvent, MemoryEventType};
use crate::snapshot::bindings::BindingRecord;
use crate::snapshot::BindingLedger;
use std::collections::HashMap;

/// Rebuild all allocations from event_store (unified data source).
///
/// Processes a slice of `MemoryEvent` and reconstructs the set of live
/// allocations along with their metadata (type, variable, thread, clone info).
pub fn rebuild_allocations_from_events(events: &[MemoryEvent]) -> Vec<AllocationInfo> {
    let mut ledger = BindingLedger::new();
    for event in events {
        ledger.apply(event);
    }

    project_allocations(&ledger, events)
}

/// Project a previously reduced capture cut without replaying its transitions.
pub(super) fn project_allocations(
    ledger: &BindingLedger,
    events: &[MemoryEvent],
) -> Vec<AllocationInfo> {
    let mut all_allocations: Vec<AllocationInfo> = ledger
        .live()
        .into_iter()
        .map(binding_to_allocation_info)
        .collect();

    let mut clone_info_map: HashMap<usize, crate::capture::types::CloneInfo> = HashMap::new();
    let mut container_allocations: Vec<AllocationInfo> = Vec::new();

    for event in events {
        match event.event_type {
            MemoryEventType::Metadata if event.ptr != 0 => {
                // A pointer-carrying metadata event annotates an existing heap binding
                // (its type was observed after the allocation was recorded). The ledger
                // already applied it, so it must not create a container allocation.
            }
            MemoryEventType::Metadata => {
                let stack_trace = event
                    .source_file
                    .as_ref()
                    .map(|file| format!("{}:{}", file, event.source_line.unwrap_or(0)));
                let mut alloc = AllocationInfo::new(0, event.size);
                alloc.timestamp_alloc = event.timestamp;
                alloc.var_name = event.var_name.clone();
                alloc.type_name = event.type_name.clone();
                alloc.thread_id_u64 = event.thread_id;
                alloc.task_id = event.task_id;
                alloc.stack_trace = stack_trace.map(|s| vec![s]);
                alloc.module_path = event.module_path.clone();
                container_allocations.push(alloc);
            }
            MemoryEventType::Clone => {
                if let (Some(source_ptr), Some(target_ptr)) =
                    (event.clone_source_ptr, event.clone_target_ptr)
                {
                    clone_info_map
                        .entry(source_ptr)
                        .and_modify(|info| info.clone_count += 1)
                        .or_insert(crate::capture::types::CloneInfo {
                            clone_count: 1,
                            is_clone: false,
                            original_ptr: None,
                            _source: None,
                            _confidence: None,
                        });

                    clone_info_map
                        .entry(target_ptr)
                        .and_modify(|info| {
                            info.is_clone = true;
                            info.original_ptr = Some(source_ptr);
                        })
                        .or_insert(crate::capture::types::CloneInfo {
                            clone_count: 0,
                            is_clone: true,
                            original_ptr: Some(source_ptr),
                            _source: None,
                            _confidence: None,
                        });
                }
            }
            // Allocation life-cycle transitions are already handled by the ledger.
            MemoryEventType::Allocate
            | MemoryEventType::Deallocate
            | MemoryEventType::Reallocate
            | MemoryEventType::Move
            | MemoryEventType::Borrow
            | MemoryEventType::Return => {}
        }
    }

    decorate_smart_pointers(&mut all_allocations, &clone_info_map);

    for alloc in &mut all_allocations {
        if let Some(clone_info) = clone_info_map.get(&alloc.ptr) {
            alloc.clone_info = Some(clone_info.clone());
        }
    }

    let container_allocations_with_virtual_ptrs: Vec<_> = container_allocations
        .into_iter()
        .enumerate()
        .map(|(index, mut alloc)| {
            alloc.ptr = crate::analysis::VIRTUAL_PTR_BASE + index;
            alloc
        })
        .collect();

    all_allocations.extend(container_allocations_with_virtual_ptrs);
    all_allocations
}

/// Project a ledger binding into the dashboard's allocation DTO.
fn binding_to_allocation_info(binding: &BindingRecord) -> AllocationInfo {
    let mut alloc = AllocationInfo::new(binding.ptr, binding.size);
    alloc.timestamp_alloc = binding.allocated_at;
    alloc.var_name = binding.var_name.clone();
    alloc.type_name = binding.type_name.clone();
    // `ThreadId` cannot be rebuilt from a `u64`; the numeric id is authoritative.
    alloc.thread_id_u64 = binding.thread_id;
    alloc.task_id = binding.task_id;
    alloc.stack_trace = match (binding.source_file.clone(), binding.source_line) {
        (Some(file), Some(line)) => Some(vec![format!("{file}:{line}")]),
        (Some(file), None) => Some(vec![file]),
        (None, _) => None,
    };
    alloc.module_path = binding.module_path.clone();
    alloc.stack_ptr = binding.stack_ptr;
    alloc.generation_id = usize::try_from(binding.generation).unwrap_or(usize::MAX);
    alloc
}

/// Classify smart pointers and attach clone relationships.
fn decorate_smart_pointers(
    allocations: &mut [AllocationInfo],
    clone_info_map: &HashMap<usize, crate::capture::types::CloneInfo>,
) {
    for alloc in allocations.iter_mut() {
        let type_name_lower = alloc
            .type_name
            .as_ref()
            .map(|s| s.to_lowercase())
            .unwrap_or_default();

        let (is_smart_pointer, pointer_type) = if type_name_lower.contains("arc") {
            (
                true,
                crate::capture::types::smart_pointer::SmartPointerType::Arc,
            )
        } else if type_name_lower.contains("rc") {
            (
                true,
                crate::capture::types::smart_pointer::SmartPointerType::Rc,
            )
        } else if type_name_lower.contains("box") {
            (
                true,
                crate::capture::types::smart_pointer::SmartPointerType::Box,
            )
        } else {
            (
                false,
                crate::capture::types::smart_pointer::SmartPointerType::Rc,
            )
        };

        if !is_smart_pointer {
            continue;
        }

        let clone_info = clone_info_map.get(&alloc.ptr);
        let is_data_owner = clone_info.map(|info| !info.is_clone).unwrap_or(true);
        alloc.smart_pointer_info = Some(crate::capture::types::smart_pointer::SmartPointerInfo {
            data_ptr: alloc.ptr,
            pointer_type,
            is_data_owner,
            ref_count_history: vec![],
            weak_count: None,
            cloned_from: clone_info.and_then(|info| info.original_ptr),
            clones: vec![],
            is_implicitly_deallocated: false,
            is_weak_reference: false,
        });
    }
}

/// Allocation sampling metadata for dashboard transparency.
///
/// When the allocation set is large, the dashboard may sample the data
/// to keep render times manageable. This struct records what was done
/// so the frontend can display a sampling notice.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SamplingMetadata {
    /// Total number of raw allocations before sampling
    pub original_count: usize,
    /// Number of allocations actually included in the dashboard
    pub displayed_count: usize,
    /// Whether sampling was applied
    pub is_sampled: bool,
    /// Description of the sampling strategy used
    pub sampling_strategy: String,
}

impl SamplingMetadata {
    /// Create sampling metadata for an unsampled (complete) dataset
    pub fn complete(count: usize) -> Self {
        Self {
            original_count: count,
            displayed_count: count,
            is_sampled: false,
            sampling_strategy: "complete".to_string(),
        }
    }

    /// Create sampling metadata for a truncated dataset
    pub fn truncated(original_count: usize, displayed_count: usize, strategy: &str) -> Self {
        Self {
            original_count,
            displayed_count,
            is_sampled: true,
            sampling_strategy: strategy.to_string(),
        }
    }
}

/// Maximum allocations to include in dashboard output before sampling.
/// Set to `None` to disable sampling entirely.
pub const MAX_DASHBOARD_ALLOCATIONS: Option<usize> = None;

#[cfg(test)]
mod tests {
    use super::*;

    /// Objective: Verify that rebuild_allocations_from_events correctly
    /// reconstructs allocate/deallocate pairs.
    /// Invariants: After allocate+deallocate, no active allocations remain.
    #[test]
    fn test_rebuild_allocate_deallocate() {
        let events = vec![
            MemoryEvent::allocate(0x100, 64, 1),
            MemoryEvent::deallocate(0x100, 64, 1),
        ];

        let result = rebuild_allocations_from_events(&events);
        assert!(
            result.is_empty(),
            "Allocations should be empty after deallocate"
        );
    }

    /// Objective: Verify that rebuild_allocations_from_events preserves
    /// allocation attributes from events.
    /// Invariants: Variable name, type name, and thread_id should be
    /// carried forward to the reconstructed allocation.
    #[test]
    fn test_rebuild_preserves_attributes() {
        let mut event = MemoryEvent::allocate(0x200, 128, 42);
        event.var_name = Some("my_var".to_string());
        event.type_name = Some("Vec<u8>".to_string());

        let events = vec![event];
        let result = rebuild_allocations_from_events(&events);

        assert_eq!(result.len(), 1, "Should have one allocation");
        assert_eq!(result[0].var_name.as_deref(), Some("my_var"));
        assert_eq!(result[0].type_name.as_deref(), Some("Vec<u8>"));
        assert_eq!(result[0].thread_id_u64, 42);
    }

    /// Objective: Verify that address reuse does not cause false
    /// duplicate allocations. After deallocate, a new allocation at
    /// the same pointer should be treated as a new allocation.
    /// Invariants: Two allocations with the same ptr but separated by
    /// a deallocation should both appear in the final set (the second one).
    #[test]
    fn test_rebuild_address_reuse() {
        let events = vec![
            MemoryEvent::allocate(0x100, 64, 1),
            MemoryEvent::deallocate(0x100, 64, 1),
            MemoryEvent::allocate(0x100, 128, 1),
        ];

        let result = rebuild_allocations_from_events(&events);
        assert_eq!(
            result.len(),
            1,
            "Should have one active allocation after reuse"
        );
        assert_eq!(
            result[0].size, 128,
            "Should reflect the newer allocation size"
        );
        assert_eq!(
            result[0].generation_id, 2,
            "the reused address must report its second binding generation"
        );
    }

    /// Objective: Verify that a late release of the previous address generation
    /// does not delete the newer allocation.
    /// Invariants: the newer binding survives with its own size and generation.
    #[test]
    fn test_rebuild_late_release_keeps_newer_generation() {
        let events = vec![
            MemoryEvent::allocate(0x300, 64, 1),
            MemoryEvent::allocate(0x300, 512, 1),
            MemoryEvent::deallocate(0x300, 64, 2),
        ];

        let result = rebuild_allocations_from_events(&events);
        assert_eq!(
            result.len(),
            1,
            "one release must close exactly one binding, got {} allocations",
            result.len()
        );
        assert_eq!(
            result[0].size, 512,
            "the surviving allocation must be the newer, larger one"
        );
    }

    /// Objective: Verify that a reallocation refreshes the allocation size instead
    /// of being silently dropped.
    /// Invariants: one allocation remains, with the new size and the same address.
    #[test]
    fn test_rebuild_applies_reallocation() {
        let events = vec![
            MemoryEvent::allocate(0x400, 32, 1),
            MemoryEvent::reallocate(0x400, 32, 256, 1),
        ];

        let result = rebuild_allocations_from_events(&events);
        assert_eq!(result.len(), 1, "reallocation must not duplicate rows");
        assert_eq!(result[0].ptr, 0x400, "the address must be preserved");
        assert_eq!(
            result[0].size, 256,
            "the reallocated size must replace the stale one"
        );
    }

    /// Objective: Verify that clone events produce correct clone info.
    /// Invariants: Clone source should have clone_count > 0,
    /// clone target should be is_clone = true.
    #[test]
    fn test_rebuild_clone_relationship() {
        let events = vec![
            MemoryEvent::allocate(0x100, 64, 1),
            MemoryEvent::allocate(0x200, 64, 1),
            MemoryEvent::clone_event(0x100, 0x200, 64, 1, None, None),
        ];

        let result = rebuild_allocations_from_events(&events);
        let source = result.iter().find(|a| a.ptr == 0x100).unwrap();
        let target = result.iter().find(|a| a.ptr == 0x200).unwrap();

        assert_eq!(
            source
                .clone_info
                .as_ref()
                .map(|c| c.clone_count)
                .unwrap_or(0),
            1,
            "Source should have clone count 1"
        );
        assert!(
            target
                .clone_info
                .as_ref()
                .map(|c| c.is_clone)
                .unwrap_or(false),
            "Target should be marked as clone"
        );
    }

    /// Objective: Verify that SamplingMetadata indicates unsampled for complete data.
    /// Invariants: is_sampled is false, original_count == displayed_count.
    #[test]
    fn test_sampling_metadata_complete() {
        let meta = SamplingMetadata::complete(100);
        assert!(!meta.is_sampled);
        assert_eq!(meta.original_count, 100);
        assert_eq!(meta.displayed_count, 100);
        assert_eq!(meta.sampling_strategy, "complete");
    }

    /// Objective: Verify that SamplingMetadata indicates sampled for truncated data.
    /// Invariants: is_sampled is true, original_count > displayed_count.
    #[test]
    fn test_sampling_metadata_truncated() {
        let meta = SamplingMetadata::truncated(1000, 100, "top-100-by-size");
        assert!(meta.is_sampled);
        assert_eq!(meta.original_count, 1000);
        assert_eq!(meta.displayed_count, 100);
        assert_eq!(meta.sampling_strategy, "top-100-by-size");
    }
}