1use crate::export::binary::error::BinaryExportError;
10use crate::export::binary::format::FileHeader;
11use bincode::{Decode, Encode};
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14use std::path::PathBuf;
15use std::time::SystemTime;
16
17#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
19pub struct BinaryIndex {
20 pub version: u32,
22
23 pub file_hash: u64,
25
26 pub file_path: PathBuf,
28
29 pub header: FileHeader,
31
32 pub string_table: StringTableIndex,
34
35 pub allocations: CompactAllocationIndex,
37
38 pub advanced_metrics: Option<AdvancedMetricsIndex>,
40
41 pub created_at: SystemTime,
43
44 pub file_size: u64,
46}
47
48impl BinaryIndex {
49 pub const INDEX_VERSION: u32 = 1;
51
52 pub fn new(file_path: PathBuf, file_hash: u64, file_size: u64, header: FileHeader) -> Self {
54 Self {
55 version: Self::INDEX_VERSION,
56 file_hash,
57 file_path,
58 header,
59 string_table: StringTableIndex::default(),
60 allocations: CompactAllocationIndex::default(),
61 advanced_metrics: None,
62 created_at: SystemTime::now(),
63 file_size,
64 }
65 }
66
67 pub fn is_valid_for_file(&self, file_path: &std::path::Path, file_hash: u64) -> bool {
69 self.file_path == file_path && self.file_hash == file_hash
70 }
71
72 pub fn record_count(&self) -> u32 {
74 self.allocations.count
75 }
76
77 pub fn get_record_offset(&self, record_index: usize) -> Option<u64> {
79 if record_index >= self.allocations.count as usize {
80 return None;
81 }
82
83 let relative_offset = self.allocations.relative_offsets.get(record_index)?;
84 Some(self.allocations.records_start_offset + *relative_offset as u64)
85 }
86
87 pub fn get_record_size(&self, record_index: usize) -> Option<u16> {
89 self.allocations.record_sizes.get(record_index).copied()
90 }
91
92 pub fn has_quick_filter_data(&self) -> bool {
94 self.allocations.quick_filter_data.is_some()
95 }
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize, Default, Encode, Decode)]
100pub struct StringTableIndex {
101 pub offset: u64,
103
104 pub size: u64,
106
107 pub string_count: u32,
109
110 pub uses_compression: bool,
112}
113
114#[derive(Debug, Clone, Serialize, Deserialize, Default, Encode, Decode)]
116pub struct CompactAllocationIndex {
117 pub count: u32,
119
120 pub records_start_offset: u64,
122
123 pub relative_offsets: Vec<u32>,
125
126 pub record_sizes: Vec<u16>,
128
129 pub quick_filter_data: Option<QuickFilterData>,
131}
132
133impl CompactAllocationIndex {
134 pub fn new(records_start_offset: u64) -> Self {
136 Self {
137 count: 0,
138 records_start_offset,
139 relative_offsets: Vec::new(),
140 record_sizes: Vec::new(),
141 quick_filter_data: None,
142 }
143 }
144
145 pub fn add_record(&mut self, absolute_offset: u64, size: u16) -> Result<(), BinaryExportError> {
147 if absolute_offset < self.records_start_offset {
148 return Err(BinaryExportError::CorruptedData(
149 "Record offset is before records start".to_string(),
150 ));
151 }
152
153 let relative_offset = absolute_offset - self.records_start_offset;
154 if relative_offset > u32::MAX as u64 {
155 return Err(BinaryExportError::CorruptedData(
156 "Record offset too large for relative addressing".to_string(),
157 ));
158 }
159
160 self.relative_offsets.push(relative_offset as u32);
161 self.record_sizes.push(size);
162 self.count += 1;
163
164 Ok(())
165 }
166
167 pub fn reserve(&mut self, capacity: usize) {
169 self.relative_offsets.reserve(capacity);
170 self.record_sizes.reserve(capacity);
171 }
172
173 pub fn memory_usage(&self) -> usize {
175 std::mem::size_of::<Self>()
176 + self.relative_offsets.capacity() * std::mem::size_of::<u32>()
177 + self.record_sizes.capacity() * std::mem::size_of::<u16>()
178 + self
179 .quick_filter_data
180 .as_ref()
181 .map_or(0, |qfd| qfd.memory_usage())
182 }
183}
184
185#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
187pub struct QuickFilterData {
188 pub batch_size: usize,
190
191 pub ptr_ranges: Vec<(usize, usize)>,
193
194 pub size_ranges: Vec<(usize, usize)>,
196
197 pub timestamp_ranges: Vec<(u64, u64)>,
199
200 pub thread_bloom_filter: Vec<u8>,
202
203 pub type_bloom_filter: Vec<u8>,
205
206 pub bloom_filter_params: BloomFilterParams,
208}
209
210impl QuickFilterData {
211 pub fn new(batch_size: usize) -> Self {
213 Self {
214 batch_size,
215 ptr_ranges: Vec::new(),
216 size_ranges: Vec::new(),
217 timestamp_ranges: Vec::new(),
218 thread_bloom_filter: Vec::new(),
219 type_bloom_filter: Vec::new(),
220 bloom_filter_params: BloomFilterParams::default(),
221 }
222 }
223
224 pub fn memory_usage(&self) -> usize {
226 std::mem::size_of::<Self>()
227 + self.ptr_ranges.capacity() * std::mem::size_of::<(usize, usize)>()
228 + self.size_ranges.capacity() * std::mem::size_of::<(usize, usize)>()
229 + self.timestamp_ranges.capacity() * std::mem::size_of::<(u64, u64)>()
230 + self.thread_bloom_filter.capacity()
231 + self.type_bloom_filter.capacity()
232 }
233
234 pub fn ptr_might_be_in_batch(&self, batch_index: usize, ptr: usize) -> bool {
236 if let Some(&(min, max)) = self.ptr_ranges.get(batch_index) {
237 ptr >= min && ptr <= max
238 } else {
239 false
240 }
241 }
242
243 pub fn size_might_be_in_batch(&self, batch_index: usize, size: usize) -> bool {
245 if let Some(&(min, max)) = self.size_ranges.get(batch_index) {
246 size >= min && size <= max
247 } else {
248 false
249 }
250 }
251
252 pub fn timestamp_might_be_in_batch(&self, batch_index: usize, timestamp: u64) -> bool {
254 if let Some(&(min, max)) = self.timestamp_ranges.get(batch_index) {
255 timestamp >= min && timestamp <= max
256 } else {
257 false
258 }
259 }
260}
261
262#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
264pub struct BloomFilterParams {
265 pub hash_functions: u32,
267
268 pub filter_size_bits: u32,
270
271 pub expected_elements: u32,
273
274 pub false_positive_rate: f64,
276}
277
278impl Default for BloomFilterParams {
279 fn default() -> Self {
280 Self {
281 hash_functions: 3,
282 filter_size_bits: 8192, expected_elements: 1000,
284 false_positive_rate: 0.01, }
286 }
287}
288
289#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
291pub struct AdvancedMetricsIndex {
292 pub offset: u64,
294
295 pub size: u32,
297
298 pub metrics_bitmap: u32,
300
301 pub metric_sections: HashMap<String, MetricSectionIndex>,
303}
304
305#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
307pub struct MetricSectionIndex {
308 pub relative_offset: u32,
310
311 pub size: u32,
313
314 pub entry_count: u32,
316}
317
318#[derive(Debug, Clone)]
320pub struct RecordMetadata {
321 pub ptr: usize,
323
324 pub size: usize,
326
327 pub timestamp: u64,
329
330 pub thread_id: Option<String>,
332
333 pub type_name: Option<String>,
335}
336
337impl RecordMetadata {
338 pub fn new(ptr: usize, size: usize, timestamp: u64) -> Self {
340 Self {
341 ptr,
342 size,
343 timestamp,
344 thread_id: None,
345 type_name: None,
346 }
347 }
348
349 pub fn with_thread_id(mut self, thread_id: String) -> Self {
351 self.thread_id = Some(thread_id);
352 self
353 }
354
355 pub fn with_type_name(mut self, type_name: String) -> Self {
357 self.type_name = Some(type_name);
358 self
359 }
360}
361
362#[cfg(test)]
363mod tests {
364 use super::*;
365 use std::path::Path;
366
367 #[test]
368 fn test_binary_index_creation() {
369 let file_path = PathBuf::from("test.memscope");
370 let file_hash = 0x123456789ABCDEF0;
371 let file_size = 1024;
372 let header = FileHeader::new_legacy(100);
373
374 let index = BinaryIndex::new(file_path.clone(), file_hash, file_size, header.clone());
375
376 assert_eq!(index.version, BinaryIndex::INDEX_VERSION);
377 assert_eq!(index.file_hash, file_hash);
378 assert_eq!(index.file_path, file_path);
379 assert_eq!(index.header, header);
380 assert_eq!(index.file_size, file_size);
381 assert_eq!(index.record_count(), 0);
382 }
383
384 #[test]
385 fn test_index_validation() {
386 let file_path = PathBuf::from("test.memscope");
387 let file_hash = 0x123456789ABCDEF0;
388 let header = FileHeader::new_legacy(100);
389
390 let index = BinaryIndex::new(file_path.clone(), file_hash, 1024, header);
391
392 assert!(index.is_valid_for_file(&file_path, file_hash));
394
395 assert!(!index.is_valid_for_file(&file_path, 0xDEADBEEF));
397
398 assert!(!index.is_valid_for_file(Path::new("other.memscope"), file_hash));
400 }
401
402 #[test]
403 fn test_compact_allocation_index() {
404 let mut index = CompactAllocationIndex::new(1000);
405
406 assert!(index.add_record(1000, 100).is_ok());
408 assert!(index.add_record(1200, 200).is_ok());
409 assert!(index.add_record(1500, 150).is_ok());
410
411 assert_eq!(index.count, 3);
412 assert_eq!(index.relative_offsets, vec![0, 200, 500]);
413 assert_eq!(index.record_sizes, vec![100, 200, 150]);
414
415 assert!(index.add_record(500, 100).is_err()); }
418
419 #[test]
420 fn test_record_offset_calculation() {
421 let file_path = PathBuf::from("test.memscope");
422 let header = FileHeader::new_legacy(3);
423 let mut index = BinaryIndex::new(file_path, 0x123, 1024, header);
424
425 index.allocations.records_start_offset = 1000;
427 index
428 .allocations
429 .add_record(1000, 100)
430 .expect("Test operation failed");
431 index
432 .allocations
433 .add_record(1200, 200)
434 .expect("Test operation failed");
435 index
436 .allocations
437 .add_record(1500, 150)
438 .expect("Test operation failed");
439
440 assert_eq!(index.get_record_offset(0), Some(1000));
442 assert_eq!(index.get_record_offset(1), Some(1200));
443 assert_eq!(index.get_record_offset(2), Some(1500));
444 assert_eq!(index.get_record_offset(3), None); assert_eq!(index.get_record_size(0), Some(100));
448 assert_eq!(index.get_record_size(1), Some(200));
449 assert_eq!(index.get_record_size(2), Some(150));
450 assert_eq!(index.get_record_size(3), None); }
452
453 #[test]
454 fn test_quick_filter_data() {
455 let mut qfd = QuickFilterData::new(1000);
456
457 qfd.ptr_ranges.push((0x1000, 0x2000));
459 qfd.ptr_ranges.push((0x3000, 0x4000));
460 qfd.size_ranges.push((100, 500));
461 qfd.size_ranges.push((600, 1000));
462 qfd.timestamp_ranges.push((1000, 2000));
463 qfd.timestamp_ranges.push((3000, 4000));
464
465 assert!(qfd.ptr_might_be_in_batch(0, 0x1500));
467 assert!(!qfd.ptr_might_be_in_batch(0, 0x2500));
468 assert!(qfd.ptr_might_be_in_batch(1, 0x3500));
469
470 assert!(qfd.size_might_be_in_batch(0, 300));
471 assert!(!qfd.size_might_be_in_batch(0, 550));
472 assert!(qfd.size_might_be_in_batch(1, 800));
473
474 assert!(qfd.timestamp_might_be_in_batch(0, 1500));
475 assert!(!qfd.timestamp_might_be_in_batch(0, 2500));
476 assert!(qfd.timestamp_might_be_in_batch(1, 3500));
477
478 assert!(!qfd.ptr_might_be_in_batch(2, 0x1500));
480 }
481
482 #[test]
483 fn test_record_metadata() {
484 let metadata = RecordMetadata::new(0x1000, 1024, 1234567890)
485 .with_thread_id("main".to_string())
486 .with_type_name("Vec<u8>".to_string());
487
488 assert_eq!(metadata.ptr, 0x1000);
489 assert_eq!(metadata.size, 1024);
490 assert_eq!(metadata.timestamp, 1234567890);
491 assert_eq!(metadata.thread_id, Some("main".to_string()));
492 assert_eq!(metadata.type_name, Some("Vec<u8>".to_string()));
493 }
494
495 #[test]
496 fn test_memory_usage_calculation() {
497 let mut index = CompactAllocationIndex::new(1000);
498
499 for i in 0..100 {
501 index
502 .add_record(1000 + i * 100, 100)
503 .expect("Test operation failed");
504 }
505
506 let memory_usage = index.memory_usage();
507
508 let expected_min = std::mem::size_of::<CompactAllocationIndex>()
510 + 100 * std::mem::size_of::<u32>()
511 + 100 * std::mem::size_of::<u16>();
512
513 assert!(memory_usage >= expected_min);
514 }
515}