memscope_rs/export/binary/
writer.rs

1//! Binary data writer for efficient allocation record serialization
2
3use crate::core::types::AllocationInfo;
4use crate::export::binary::config::BinaryExportConfig;
5use crate::export::binary::error::BinaryExportError;
6use crate::export::binary::format::{
7    AdvancedMetricsHeader, BinaryExportMode, FileHeader, MetricsBitmapFlags, ALLOCATION_RECORD_TYPE,
8};
9use crate::export::binary::serializable::BinarySerializable;
10use crate::export::binary::string_table::{StringTable, StringTableBuilder};
11use std::fs::File;
12use std::io::{BufWriter, Write};
13use std::path::Path;
14
15/// Binary writer for allocation records using buffered I/O
16pub struct BinaryWriter {
17    writer: BufWriter<File>,
18    config: BinaryExportConfig,
19    string_table: Option<StringTable>,
20}
21
22impl BinaryWriter {
23    /// Create new binary writer for the specified file path with default config
24    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, BinaryExportError> {
25        Self::new_with_config(path, &BinaryExportConfig::default())
26    }
27
28    /// Create new binary writer with custom configuration
29    pub fn new_with_config<P: AsRef<Path>>(
30        path: P,
31        config: &BinaryExportConfig,
32    ) -> Result<Self, BinaryExportError> {
33        let file = File::create(path)?;
34        let writer = BufWriter::new(file);
35
36        Ok(Self {
37            writer,
38            config: config.clone(),
39            string_table: None,
40        })
41    }
42
43    /// Build string table from allocation data for optimization
44    pub fn build_string_table(
45        &mut self,
46        allocations: &[AllocationInfo],
47    ) -> Result<(), BinaryExportError> {
48        if !self.config.string_table_optimization {
49            return Ok(()); // String table optimization disabled
50        }
51
52        // Use frequency threshold based on data size
53        let min_frequency = if allocations.len() > 1000 { 3 } else { 2 };
54        let mut builder = StringTableBuilder::new(min_frequency);
55
56        // Collect string frequencies from all allocations
57        for alloc in allocations {
58            // Record frequently repeated strings
59            if let Some(ref type_name) = alloc.type_name {
60                builder.record_string(type_name);
61            }
62            if let Some(ref var_name) = alloc.var_name {
63                builder.record_string(var_name);
64            }
65            if let Some(ref scope_name) = alloc.scope_name {
66                builder.record_string(scope_name);
67            }
68            builder.record_string(&alloc.thread_id);
69
70            // Record stack trace strings (function names are often repeated)
71            if let Some(ref stack_trace) = alloc.stack_trace {
72                for frame in stack_trace {
73                    builder.record_string(frame);
74                }
75            }
76        }
77
78        let table = builder.build()?;
79        let stats = table.compression_stats();
80
81        // Only use string table if it provides meaningful compression
82        if stats.space_saved() > 0 && !table.is_empty() {
83            tracing::debug!(
84                "String table built: {} strings, {:.1}% space savings",
85                table.len(),
86                stats.space_saved_percent()
87            );
88            self.string_table = Some(table);
89        }
90
91        Ok(())
92    }
93
94    /// Write file header with allocation count and optional string table (legacy compatibility)
95    pub fn write_header(&mut self, count: u32) -> Result<(), BinaryExportError> {
96        let header = FileHeader::new_legacy(count);
97        let header_bytes = header.to_bytes();
98
99        self.writer.write_all(&header_bytes)?;
100
101        self.write_string_table_if_present()
102    }
103
104    /// Write enhanced file header with export mode and allocation counts
105    pub fn write_enhanced_header(
106        &mut self,
107        total_count: u32,
108        export_mode: BinaryExportMode,
109        user_count: u16,
110        system_count: u16,
111    ) -> Result<(), BinaryExportError> {
112        let header = FileHeader::new(total_count, export_mode, user_count, system_count);
113        let header_bytes = header.to_bytes();
114
115        self.writer.write_all(&header_bytes)?;
116
117        self.write_string_table_if_present()
118    }
119
120    /// Write string table if present (helper method)
121    fn write_string_table_if_present(&mut self) -> Result<(), BinaryExportError> {
122        // Write string table if present
123        if let Some(ref string_table) = self.string_table {
124            // Write string table marker and size
125            self.writer.write_all(b"STBL")?; // String table marker
126            let table_size = string_table.serialized_size() as u32;
127            self.writer.write_all(&table_size.to_le_bytes())?;
128
129            // Write the string table
130            string_table.write_binary(&mut self.writer)?;
131        } else {
132            // Write empty string table marker
133            self.writer.write_all(b"NONE")?; // No string table
134            self.writer.write_all(&0u32.to_le_bytes())?; // Size 0
135        }
136
137        Ok(())
138    }
139
140    /// Write single allocation record in TLV format
141    pub fn write_allocation(&mut self, alloc: &AllocationInfo) -> Result<(), BinaryExportError> {
142        // Calculate value size (excluding Type and Length fields)
143        let value_size = self.calculate_value_size(alloc);
144
145        // Write Type (1 byte)
146        self.writer.write_all(&[ALLOCATION_RECORD_TYPE])?;
147
148        // Write Length (4 bytes, Little Endian) - only the Value part
149        self.writer.write_all(&(value_size as u32).to_le_bytes())?;
150
151        // Write Value: basic fields (ptr, size, timestamps)
152        self.writer.write_all(&(alloc.ptr as u64).to_le_bytes())?;
153        self.writer.write_all(&(alloc.size as u64).to_le_bytes())?;
154        self.writer
155            .write_all(&alloc.timestamp_alloc.to_le_bytes())?;
156
157        // Write optional timestamp_dealloc
158        match alloc.timestamp_dealloc {
159            Some(ts) => {
160                self.writer.write_all(&1u8.to_le_bytes())?; // has value
161                self.writer.write_all(&ts.to_le_bytes())?;
162            }
163            None => {
164                self.writer.write_all(&0u8.to_le_bytes())?; // no value
165            }
166        }
167
168        // Write string fields
169        self.write_optional_string(&alloc.var_name)?;
170        self.write_optional_string(&alloc.type_name)?;
171        self.write_optional_string(&alloc.scope_name)?;
172        self.write_string(&alloc.thread_id)?;
173
174        // Write stack trace
175        self.write_optional_string_vec(&alloc.stack_trace)?;
176
177        // Write numeric fields
178        self.writer
179            .write_all(&(alloc.borrow_count as u32).to_le_bytes())?;
180        self.writer
181            .write_all(&(alloc.is_leaked as u8).to_le_bytes())?;
182
183        // Write optional lifetime_ms
184        match alloc.lifetime_ms {
185            Some(ms) => {
186                self.writer.write_all(&1u8.to_le_bytes())?; // has value
187                self.writer.write_all(&ms.to_le_bytes())?;
188            }
189            None => {
190                self.writer.write_all(&0u8.to_le_bytes())?; // no value
191            }
192        }
193
194        // Write improve.md extensions: borrow_info
195        match &alloc.borrow_info {
196            Some(borrow_info) => {
197                self.writer.write_all(&1u8.to_le_bytes())?; // has value
198                self.writer
199                    .write_all(&(borrow_info.immutable_borrows as u32).to_le_bytes())?;
200                self.writer
201                    .write_all(&(borrow_info.mutable_borrows as u32).to_le_bytes())?;
202                self.writer
203                    .write_all(&(borrow_info.max_concurrent_borrows as u32).to_le_bytes())?;
204                match borrow_info.last_borrow_timestamp {
205                    Some(ts) => {
206                        self.writer.write_all(&1u8.to_le_bytes())?; // has timestamp
207                        self.writer.write_all(&ts.to_le_bytes())?;
208                    }
209                    None => {
210                        self.writer.write_all(&0u8.to_le_bytes())?; // no timestamp
211                    }
212                }
213            }
214            None => {
215                self.writer.write_all(&0u8.to_le_bytes())?; // no borrow_info
216            }
217        }
218
219        // Write improve.md extensions: clone_info
220        match &alloc.clone_info {
221            Some(clone_info) => {
222                self.writer.write_all(&1u8.to_le_bytes())?; // has value
223                self.writer
224                    .write_all(&(clone_info.clone_count as u32).to_le_bytes())?;
225                self.writer
226                    .write_all(&(clone_info.is_clone as u8).to_le_bytes())?;
227                match clone_info.original_ptr {
228                    Some(ptr) => {
229                        self.writer.write_all(&1u8.to_le_bytes())?; // has original_ptr
230                        self.writer.write_all(&(ptr as u64).to_le_bytes())?;
231                    }
232                    None => {
233                        self.writer.write_all(&0u8.to_le_bytes())?; // no original_ptr
234                    }
235                }
236            }
237            None => {
238                self.writer.write_all(&0u8.to_le_bytes())?; // no clone_info
239            }
240        }
241
242        // Write improve.md extensions: ownership_history_available
243        self.writer
244            .write_all(&(alloc.ownership_history_available as u8).to_le_bytes())?;
245
246        // Write complex fields using binary serialization
247        self.write_optional_binary_field(&alloc.smart_pointer_info)?;
248        self.write_optional_binary_field(&alloc.memory_layout)?;
249        self.write_optional_json_field(&alloc.generic_info)?;
250        self.write_optional_json_field(&alloc.dynamic_type_info)?;
251        self.write_optional_json_field(&alloc.runtime_state)?;
252        self.write_optional_json_field(&alloc.stack_allocation)?;
253        self.write_optional_json_field(&alloc.temporary_object)?;
254        self.write_optional_json_field(&alloc.fragmentation_analysis)?;
255        self.write_optional_json_field(&alloc.generic_instantiation)?;
256        self.write_optional_json_field(&alloc.type_relationships)?;
257        self.write_optional_json_field(&alloc.type_usage)?;
258        self.write_optional_json_field(&alloc.function_call_tracking)?;
259        self.write_optional_json_field(&alloc.lifecycle_tracking)?;
260        self.write_optional_json_field(&alloc.access_tracking)?;
261
262        Ok(())
263    }
264
265    /// Write advanced metrics segment if enabled in config
266    pub fn write_advanced_metrics_segment(
267        &mut self,
268        allocations: &[AllocationInfo],
269    ) -> Result<(), BinaryExportError> {
270        if !self.config.has_advanced_metrics() {
271            return Ok(()); // Skip if no advanced metrics enabled
272        }
273
274        // Calculate metrics bitmap based on config
275        let mut metrics_bitmap = 0u32;
276
277        if self.config.lifecycle_timeline {
278            metrics_bitmap =
279                MetricsBitmapFlags::enable(metrics_bitmap, MetricsBitmapFlags::LifecycleAnalysis);
280        }
281        if self.config.container_analysis {
282            metrics_bitmap =
283                MetricsBitmapFlags::enable(metrics_bitmap, MetricsBitmapFlags::ContainerAnalysis);
284        }
285        if self.config.source_analysis {
286            metrics_bitmap =
287                MetricsBitmapFlags::enable(metrics_bitmap, MetricsBitmapFlags::SourceAnalysis);
288        }
289        if self.config.fragmentation_analysis {
290            metrics_bitmap = MetricsBitmapFlags::enable(
291                metrics_bitmap,
292                MetricsBitmapFlags::FragmentationAnalysis,
293            );
294        }
295        if self.config.thread_context_tracking {
296            metrics_bitmap =
297                MetricsBitmapFlags::enable(metrics_bitmap, MetricsBitmapFlags::ThreadContext);
298        }
299        if self.config.drop_chain_analysis {
300            metrics_bitmap =
301                MetricsBitmapFlags::enable(metrics_bitmap, MetricsBitmapFlags::DropChainAnalysis);
302        }
303        if self.config.zst_analysis {
304            metrics_bitmap =
305                MetricsBitmapFlags::enable(metrics_bitmap, MetricsBitmapFlags::ZstAnalysis);
306        }
307        if self.config.health_scoring {
308            metrics_bitmap =
309                MetricsBitmapFlags::enable(metrics_bitmap, MetricsBitmapFlags::HealthScoring);
310        }
311        if self.config.performance_benchmarking {
312            metrics_bitmap = MetricsBitmapFlags::enable(
313                metrics_bitmap,
314                MetricsBitmapFlags::PerformanceBenchmarks,
315            );
316        }
317
318        // Calculate segment size (header + data)
319        let data_size = self.calculate_advanced_metrics_data_size(allocations, metrics_bitmap);
320        let segment_size = 16 + data_size; // 16 bytes for AdvancedMetricsHeader
321
322        // Write advanced metrics header
323        let header = AdvancedMetricsHeader::new(segment_size as u32, metrics_bitmap);
324        let header_bytes = header.to_bytes();
325        self.writer.write_all(&header_bytes)?;
326
327        // Write metrics data based on enabled flags
328        self.write_metrics_data(allocations, metrics_bitmap)?;
329
330        Ok(())
331    }
332
333    /// Calculate size needed for advanced metrics data
334    fn calculate_advanced_metrics_data_size(
335        &self,
336        allocations: &[AllocationInfo],
337        metrics_bitmap: u32,
338    ) -> usize {
339        let mut size = 0;
340
341        // For each enabled metric, calculate its data size
342        if MetricsBitmapFlags::is_enabled(metrics_bitmap, MetricsBitmapFlags::LifecycleAnalysis) {
343            size += self.calculate_lifecycle_data_size(allocations);
344        }
345        if MetricsBitmapFlags::is_enabled(metrics_bitmap, MetricsBitmapFlags::ContainerAnalysis) {
346            size += self.calculate_container_data_size(allocations);
347        }
348        if MetricsBitmapFlags::is_enabled(metrics_bitmap, MetricsBitmapFlags::TypeUsageStats) {
349            size += self.calculate_type_usage_data_size(allocations);
350        }
351        // Add other metrics as needed...
352
353        size
354    }
355
356    /// Write metrics data based on enabled bitmap flags
357    fn write_metrics_data(
358        &mut self,
359        allocations: &[AllocationInfo],
360        metrics_bitmap: u32,
361    ) -> Result<(), BinaryExportError> {
362        // Write lifecycle analysis data if enabled
363        if MetricsBitmapFlags::is_enabled(metrics_bitmap, MetricsBitmapFlags::LifecycleAnalysis) {
364            self.write_lifecycle_metrics(allocations)?;
365        }
366
367        // Write container analysis data if enabled
368        if MetricsBitmapFlags::is_enabled(metrics_bitmap, MetricsBitmapFlags::ContainerAnalysis) {
369            self.write_container_metrics(allocations)?;
370        }
371
372        // Write type usage statistics if enabled
373        if MetricsBitmapFlags::is_enabled(metrics_bitmap, MetricsBitmapFlags::TypeUsageStats) {
374            self.write_type_usage_metrics(allocations)?;
375        }
376
377        // Add other metrics as they are implemented...
378
379        Ok(())
380    }
381
382    /// Write lifecycle analysis metrics
383    fn write_lifecycle_metrics(
384        &mut self,
385        allocations: &[AllocationInfo],
386    ) -> Result<(), BinaryExportError> {
387        // Write count of allocations with lifecycle data
388        let lifecycle_count = allocations
389            .iter()
390            .filter(|a| a.lifetime_ms.is_some())
391            .count();
392
393        self.writer
394            .write_all(&(lifecycle_count as u32).to_le_bytes())?;
395
396        // Write lifecycle data for each allocation
397        for alloc in allocations {
398            if let Some(lifetime) = alloc.lifetime_ms {
399                self.writer.write_all(&(alloc.ptr as u64).to_le_bytes())?; // allocation ID
400                self.writer.write_all(&lifetime.to_le_bytes())?; // lifetime in ms
401
402                // Write lifecycle phase information if available
403                if let Some(ref lifecycle_tracking) = alloc.lifecycle_tracking {
404                    self.writer.write_all(&1u8.to_le_bytes())?; // has lifecycle tracking
405                    let json_str = serde_json::to_string(lifecycle_tracking).map_err(|e| {
406                        BinaryExportError::CorruptedData(format!(
407                            "Lifecycle JSON serialization failed: {e}"
408                        ))
409                    })?;
410                    self.write_string(&json_str)?;
411                } else {
412                    self.writer.write_all(&0u8.to_le_bytes())?; // no lifecycle tracking
413                }
414            }
415        }
416
417        Ok(())
418    }
419
420    /// Write container analysis metrics
421    fn write_container_metrics(
422        &mut self,
423        allocations: &[AllocationInfo],
424    ) -> Result<(), BinaryExportError> {
425        // Write count of allocations with container data
426        let container_count = allocations
427            .iter()
428            .filter(|a| a.memory_layout.is_some())
429            .count();
430
431        self.writer
432            .write_all(&(container_count as u32).to_le_bytes())?;
433
434        // Write container data for each allocation
435        for alloc in allocations {
436            if let Some(ref memory_layout) = alloc.memory_layout {
437                self.writer.write_all(&(alloc.ptr as u64).to_le_bytes())?; // allocation ID
438
439                // Serialize memory layout as JSON for now (can be optimized later)
440                let json_str = serde_json::to_string(memory_layout).map_err(|e| {
441                    BinaryExportError::CorruptedData(format!(
442                        "Memory layout JSON serialization failed: {e}"
443                    ))
444                })?;
445                self.write_string(&json_str)?;
446            }
447        }
448
449        Ok(())
450    }
451
452    /// Write type usage statistics
453    fn write_type_usage_metrics(
454        &mut self,
455        allocations: &[AllocationInfo],
456    ) -> Result<(), BinaryExportError> {
457        // Write count of allocations with type usage data
458        let type_usage_count = allocations
459            .iter()
460            .filter(|a| a.type_usage.is_some())
461            .count();
462
463        self.writer
464            .write_all(&(type_usage_count as u32).to_le_bytes())?;
465
466        // Write type usage data for each allocation
467        for alloc in allocations {
468            if let Some(ref type_usage) = alloc.type_usage {
469                self.writer.write_all(&(alloc.ptr as u64).to_le_bytes())?; // allocation ID
470
471                // Serialize type usage as JSON for now (can be optimized later)
472                let json_str = serde_json::to_string(type_usage).map_err(|e| {
473                    BinaryExportError::CorruptedData(format!(
474                        "Type usage JSON serialization failed: {e}"
475                    ))
476                })?;
477                self.write_string(&json_str)?;
478            }
479        }
480
481        Ok(())
482    }
483
484    /// Calculate size for lifecycle data
485    fn calculate_lifecycle_data_size(&self, allocations: &[AllocationInfo]) -> usize {
486        let mut size = 4; // count field
487
488        for alloc in allocations {
489            if alloc.lifetime_ms.is_some() {
490                size += 8; // ptr
491                size += 8; // lifetime_ms
492                size += 1; // lifecycle_tracking flag
493
494                if let Some(ref lifecycle_tracking) = alloc.lifecycle_tracking {
495                    if let Ok(json_str) = serde_json::to_string(lifecycle_tracking) {
496                        size += 4 + json_str.len(); // string length + content
497                    }
498                }
499            }
500        }
501
502        size
503    }
504
505    /// Calculate size for container data
506    fn calculate_container_data_size(&self, allocations: &[AllocationInfo]) -> usize {
507        let mut size = 4; // count field
508
509        for alloc in allocations {
510            if let Some(ref memory_layout) = alloc.memory_layout {
511                size += 8; // ptr
512                if let Ok(json_str) = serde_json::to_string(memory_layout) {
513                    size += 4 + json_str.len(); // string length + content
514                }
515            }
516        }
517
518        size
519    }
520
521    /// Calculate size for type usage data
522    fn calculate_type_usage_data_size(&self, allocations: &[AllocationInfo]) -> usize {
523        let mut size = 4; // count field
524
525        for alloc in allocations {
526            if let Some(ref type_usage) = alloc.type_usage {
527                size += 8; // ptr
528                if let Ok(json_str) = serde_json::to_string(type_usage) {
529                    size += 4 + json_str.len(); // string length + content
530                }
531            }
532        }
533
534        size
535    }
536
537    /// Finish writing and flush all data to disk
538    pub fn finish(mut self) -> Result<(), BinaryExportError> {
539        self.writer.flush()?;
540        Ok(())
541    }
542
543    /// Calculate size needed for the Value part of TLV (excluding Type and Length)
544    fn calculate_value_size(&self, alloc: &AllocationInfo) -> usize {
545        let mut size = 8 + 8 + 8; // ptr + size + timestamp_alloc
546
547        // timestamp_dealloc: 1 byte flag + optional 8 bytes
548        size += 1;
549        if alloc.timestamp_dealloc.is_some() {
550            size += 8;
551        }
552
553        // String fields: length (4 bytes) + content
554        size += 4; // var_name length
555        if let Some(ref name) = alloc.var_name {
556            size += name.len();
557        }
558
559        size += 4; // type_name length
560        if let Some(ref name) = alloc.type_name {
561            size += name.len();
562        }
563
564        size += 4; // scope_name length
565        if let Some(ref name) = alloc.scope_name {
566            size += name.len();
567        }
568
569        size += 4 + alloc.thread_id.len(); // thread_id
570
571        // Stack trace
572        size += 4; // stack_trace count
573        if let Some(ref stack_trace) = alloc.stack_trace {
574            for frame in stack_trace {
575                size += 4 + frame.len(); // length + content for each frame
576            }
577        }
578
579        // Numeric fields
580        size += 4; // borrow_count
581        size += 1; // is_leaked
582
583        // lifetime_ms: 1 byte flag + optional 8 bytes
584        size += 1;
585        if alloc.lifetime_ms.is_some() {
586            size += 8;
587        }
588
589        // improve.md extensions: borrow_info
590        size += 1; // presence flag
591        if let Some(ref borrow_info) = alloc.borrow_info {
592            size += 4 + 4 + 4; // immutable_borrows + mutable_borrows + max_concurrent_borrows
593            size += 1; // timestamp presence flag
594            if borrow_info.last_borrow_timestamp.is_some() {
595                size += 8; // timestamp
596            }
597        }
598
599        // improve.md extensions: clone_info
600        size += 1; // presence flag
601        if let Some(ref clone_info) = alloc.clone_info {
602            size += 4 + 1; // clone_count + is_clone
603            size += 1; // original_ptr presence flag
604            if clone_info.original_ptr.is_some() {
605                size += 8; // original_ptr
606            }
607        }
608
609        // improve.md extensions: ownership_history_available
610        size += 1; // boolean flag
611
612        // Binary fields
613        size += self.calculate_binary_field_size(&alloc.smart_pointer_info);
614        size += self.calculate_binary_field_size(&alloc.memory_layout);
615        size += self.calculate_json_field_size(&alloc.generic_info);
616        size += self.calculate_json_field_size(&alloc.dynamic_type_info);
617        size += self.calculate_json_field_size(&alloc.runtime_state);
618        size += self.calculate_json_field_size(&alloc.stack_allocation);
619        size += self.calculate_json_field_size(&alloc.temporary_object);
620        size += self.calculate_json_field_size(&alloc.fragmentation_analysis);
621        size += self.calculate_json_field_size(&alloc.generic_instantiation);
622        size += self.calculate_json_field_size(&alloc.type_relationships);
623        size += self.calculate_json_field_size(&alloc.type_usage);
624        size += self.calculate_json_field_size(&alloc.function_call_tracking);
625        size += self.calculate_json_field_size(&alloc.lifecycle_tracking);
626        size += self.calculate_json_field_size(&alloc.access_tracking);
627
628        size
629    }
630
631    /// Calculate size needed for optional binary field
632    fn calculate_binary_field_size<T: BinarySerializable>(&self, field: &Option<T>) -> usize {
633        let mut size = 1; // flag byte
634        if let Some(value) = field {
635            size += value.binary_size();
636        }
637        size
638    }
639
640    /// Calculate size needed for optional JSON field
641    fn calculate_json_field_size<T: serde::Serialize>(&self, field: &Option<T>) -> usize {
642        let mut size = 1; // flag byte
643        if let Some(value) = field {
644            // Pre-serialize to get exact size
645            if let Ok(json_str) = serde_json::to_string(value) {
646                size += 4 + json_str.len(); // length + content
647            } else {
648                // If serialization fails, just account for the flag
649                size = 1;
650            }
651        }
652        size
653    }
654
655    /// Write optional string field, using string table if available
656    fn write_optional_string(&mut self, opt_str: &Option<String>) -> Result<(), BinaryExportError> {
657        match opt_str {
658            Some(s) => {
659                if let Some(ref string_table) = self.string_table {
660                    // Try to find string in table
661                    if let Some(index) = self.find_string_index(s) {
662                        // Write as string table reference: 0xFFFF followed by compressed index
663                        self.writer.write_all(&0xFFFFu32.to_le_bytes())?;
664                        string_table.write_compressed_index(&mut self.writer, index)?;
665                    } else {
666                        // Write as inline string
667                        self.write_inline_string(s)?;
668                    }
669                } else {
670                    // No string table, write inline
671                    self.write_inline_string(s)?;
672                }
673            }
674            None => {
675                // Use 0xFFFFFFFE as None marker to distinguish from empty string (length 0)
676                self.writer.write_all(&0xFFFFFFFEu32.to_le_bytes())?;
677            }
678        }
679        Ok(())
680    }
681
682    /// Write string field, using string table if available
683    fn write_string(&mut self, s: &str) -> Result<(), BinaryExportError> {
684        if let Some(ref string_table) = self.string_table {
685            // Try to find string in table
686            if let Some(index) = self.find_string_index(s) {
687                // Write as string table reference: 0xFFFF followed by compressed index
688                self.writer.write_all(&0xFFFFu32.to_le_bytes())?;
689                string_table.write_compressed_index(&mut self.writer, index)?;
690            } else {
691                // Write as inline string
692                self.write_inline_string(s)?;
693            }
694        } else {
695            // No string table, write inline
696            self.write_inline_string(s)?;
697        }
698        Ok(())
699    }
700
701    /// Write string inline with length prefix
702    fn write_inline_string(&mut self, s: &str) -> Result<(), BinaryExportError> {
703        self.writer.write_all(&(s.len() as u32).to_le_bytes())?;
704        self.writer.write_all(s.as_bytes())?;
705        Ok(())
706    }
707
708    /// Find string index in string table
709    fn find_string_index(&self, s: &str) -> Option<u16> {
710        // Don't use string table for empty strings
711        if s.is_empty() {
712            return None;
713        }
714
715        if let Some(ref string_table) = self.string_table {
716            string_table.get_index(s)
717        } else {
718            None
719        }
720    }
721
722    /// Write an optional vector of strings, using string table if available
723    fn write_optional_string_vec(
724        &mut self,
725        vec: &Option<Vec<String>>,
726    ) -> Result<(), BinaryExportError> {
727        match vec {
728            Some(strings) => {
729                // Write count
730                self.writer
731                    .write_all(&(strings.len() as u32).to_le_bytes())?;
732                // Write each string using string table optimization
733                for string in strings {
734                    self.write_string(string)?;
735                }
736            }
737            None => {
738                self.writer.write_all(&0u32.to_le_bytes())?; // 0 count indicates None
739            }
740        }
741        Ok(())
742    }
743
744    /// Write optional binary field using BinarySerializable trait
745    fn write_optional_binary_field<T: BinarySerializable>(
746        &mut self,
747        field: &Option<T>,
748    ) -> Result<(), BinaryExportError> {
749        match field {
750            Some(value) => {
751                self.writer.write_all(&1u8.to_le_bytes())?; // has value
752                value.write_binary(&mut self.writer)?;
753            }
754            None => {
755                self.writer.write_all(&0u8.to_le_bytes())?; // no value
756            }
757        }
758        Ok(())
759    }
760
761    /// Write optional JSON field (serialize to JSON string)
762    fn write_optional_json_field<T: serde::Serialize>(
763        &mut self,
764        field: &Option<T>,
765    ) -> Result<(), BinaryExportError> {
766        match field {
767            Some(value) => {
768                let json_str = serde_json::to_string(value).map_err(|e| {
769                    BinaryExportError::CorruptedData(format!("JSON serialization failed: {e}"))
770                })?;
771                self.writer.write_all(&1u8.to_le_bytes())?; // has value
772                self.write_string(&json_str)?;
773            }
774            None => {
775                self.writer.write_all(&0u8.to_le_bytes())?; // no value
776            }
777        }
778        Ok(())
779    }
780}
781
782#[cfg(test)]
783mod tests {
784    use super::*;
785    use std::fs;
786    use tempfile::NamedTempFile;
787
788    fn create_test_allocation() -> AllocationInfo {
789        AllocationInfo {
790            ptr: 0x1000,
791            size: 1024,
792            var_name: Some("test_var".to_string()),
793            type_name: Some("i32".to_string()),
794            scope_name: None,
795            timestamp_alloc: 1234567890,
796            timestamp_dealloc: None,
797            thread_id: "main".to_string(),
798            borrow_count: 0,
799            stack_trace: None,
800            is_leaked: false,
801            lifetime_ms: None,
802            borrow_info: None,
803            clone_info: None,
804            ownership_history_available: false,
805            smart_pointer_info: None,
806            memory_layout: None,
807            generic_info: None,
808            dynamic_type_info: None,
809            runtime_state: None,
810            stack_allocation: None,
811            temporary_object: None,
812            fragmentation_analysis: None,
813            generic_instantiation: None,
814            type_relationships: None,
815            type_usage: None,
816            function_call_tracking: None,
817            lifecycle_tracking: None,
818            access_tracking: None,
819            drop_chain_analysis: None,
820        }
821    }
822
823    #[test]
824    fn test_writer_creation() {
825        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
826        let writer = BinaryWriter::new(temp_file.path());
827        assert!(writer.is_ok());
828    }
829
830    #[test]
831    fn test_header_writing() {
832        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
833        let mut writer = BinaryWriter::new(temp_file.path()).expect("Failed to create temp file");
834
835        let result = writer.write_header(42);
836        assert!(result.is_ok());
837
838        writer.finish().expect("Failed to finish writing");
839
840        // Verify file size is at least header size
841        let metadata = fs::metadata(temp_file.path()).expect("Failed to create temp file");
842        assert!(metadata.len() >= 16);
843    }
844
845    #[test]
846    fn test_allocation_writing() {
847        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
848        let mut writer = BinaryWriter::new(temp_file.path()).expect("Failed to create temp file");
849
850        writer.write_header(1).expect("Failed to write header");
851
852        let alloc = create_test_allocation();
853        let result = writer.write_allocation(&alloc);
854        assert!(result.is_ok());
855
856        writer.finish().expect("Failed to finish writing");
857
858        // Verify file has content beyond header
859        let metadata = fs::metadata(temp_file.path()).expect("Failed to create temp file");
860        assert!(metadata.len() > 16);
861    }
862
863    #[test]
864    fn test_record_size_calculation() {
865        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
866        let writer = BinaryWriter::new(temp_file.path()).expect("Failed to create temp file");
867
868        let alloc = create_test_allocation();
869        let size = writer.calculate_value_size(&alloc);
870
871        // Basic fields: 8 + 8 + 8 = 24 (ptr + size + timestamp_alloc)
872        // timestamp_dealloc: 1 byte flag (None) = 1
873        // var_name: 4 + 8 = 12 ("test_var")
874        // type_name: 4 + 3 = 7 ("i32")
875        // scope_name: 4 + 0 = 4 (None)
876        // thread_id: 4 + 4 = 8 ("main")
877        // stack_trace: 4 + 0 = 4 (None)
878        // borrow_count: 4
879        // is_leaked: 1
880        // lifetime_ms: 1 byte flag (None) = 1
881        // borrow_info: 1 byte flag (None) = 1
882        // clone_info: 1 byte flag (None) = 1
883        // ownership_history_available: 1 byte = 1
884        // JSON fields (14 fields * 1 byte flag each): 14
885        // Total: 24 + 1 + 12 + 7 + 4 + 8 + 4 + 4 + 1 + 1 + 1 + 1 + 1 + 14 = 83
886        assert_eq!(size, 83);
887    }
888
889    #[test]
890    fn test_advanced_metrics_segment_writing() {
891        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
892        let config = BinaryExportConfig::debug_comprehensive();
893        let mut writer = BinaryWriter::new_with_config(temp_file.path(), &config)
894            .expect("Failed to create temp file");
895
896        writer.write_header(1).expect("Failed to write header");
897
898        let mut alloc = create_test_allocation();
899        alloc.lifetime_ms = Some(1500); // Add some lifecycle data
900
901        writer
902            .write_allocation(&alloc)
903            .expect("Failed to write allocation");
904        writer
905            .write_advanced_metrics_segment(&[alloc])
906            .expect("Failed to write metrics segment");
907        writer.finish().expect("Failed to finish writing");
908
909        // Verify file has content beyond basic allocation data
910        let metadata = std::fs::metadata(temp_file.path()).expect("Failed to create temp file");
911        assert!(metadata.len() > 100); // Should be larger with advanced metrics
912    }
913
914    #[test]
915    fn test_advanced_metrics_disabled() {
916        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
917        let config = BinaryExportConfig::minimal(); // No advanced metrics
918        let mut writer = BinaryWriter::new_with_config(temp_file.path(), &config)
919            .expect("Failed to create temp file");
920
921        writer.write_header(1).expect("Failed to write header");
922        let alloc = create_test_allocation();
923        writer
924            .write_allocation(&alloc)
925            .expect("Failed to write allocation");
926
927        // Should not write advanced metrics segment
928        writer
929            .write_advanced_metrics_segment(&[alloc])
930            .expect("Failed to write metrics segment");
931        writer.finish().expect("Failed to finish writing");
932
933        // File should be smaller without advanced metrics
934        let metadata = std::fs::metadata(temp_file.path()).expect("Failed to create temp file");
935        assert!(metadata.len() < 150); // Should be smaller without advanced metrics
936    }
937}
memscope_rs/export/binary/writer.rs

memscope_rs/export/binary/
writer.rs
