xet_client/cas_client/simulation/
random_xorb.rs

1//! Random XORB generation for testing with large files that are never fully materialized.
2//!
3//! This module provides `RandomXorb`, a structure that can generate XORB chunk data
4//! on-the-fly using deterministic random seeds, allowing testing with massive files
5//! without actually storing them in memory.
6
7use bytes::Bytes;
8use rand::prelude::*;
9use xet_core_structures::merklehash::{MerkleHash, compute_data_hash};
10use xet_core_structures::xorb_object::{
11    CompressionScheme, XORB_CHUNK_HEADER_LENGTH, XorbChunkHeader, XorbObject, XorbObjectInfoV1,
12};
13
14/// Information about a single chunk in a RandomXorb.
15#[derive(Clone, Debug)]
16pub struct RandomChunkInfo {
17    /// Random seed used to generate this chunk's data.
18    pub seed: u64,
19    /// Size of the uncompressed chunk data in bytes.
20    pub size: u32,
21    /// Cached hash of the chunk data.
22    pub hash: MerkleHash,
23}
24
25/// A XORB that generates its chunk data on-the-fly from random seeds.
26///
27/// This allows testing with massive files without actually storing the data in memory.
28/// Each chunk is defined by a seed and size, and the data is generated deterministically
29/// using `SmallRng` when needed.
30#[derive(Clone, Debug)]
31pub struct RandomXorb {
32    /// Information about each chunk.
33    chunks: Vec<RandomChunkInfo>,
34    /// Cached XorbObject header/footer.
35    xorb_object: XorbObject,
36}
37
38impl RandomXorb {
39    /// Creates a new RandomXorb from a list of (seed, size) pairs.
40    ///
41    /// The chunk data is generated deterministically from each seed, and
42    /// chunk hashes are computed and cached.
43    pub fn new(chunk_specs: &[(u64, u32)]) -> Self {
44        let chunks: Vec<RandomChunkInfo> = chunk_specs
45            .iter()
46            .map(|&(seed, size)| {
47                let data = Self::generate_chunk_data_from_seed(seed, size);
48                let hash = compute_data_hash(&data);
49                RandomChunkInfo { seed, size, hash }
50            })
51            .collect();
52
53        // Build the XorbObject header/footer
54        let xorb_obj = Self::build_xorb_object(&chunks);
55
56        Self {
57            chunks,
58            xorb_object: xorb_obj,
59        }
60    }
61
62    /// Creates a new RandomXorb from a seed, number of chunks, and chunk size.
63    ///
64    /// Each chunk gets a unique sub-seed derived from the main seed.
65    pub fn from_seed(seed: u64, num_chunks: u32, chunk_size: u32) -> Self {
66        use rand::prelude::*;
67
68        let mut rng = SmallRng::seed_from_u64(seed);
69        let chunk_specs: Vec<(u64, u32)> = (0..num_chunks)
70            .map(|_| {
71                let chunk_seed = rng.random::<u64>();
72                (chunk_seed, chunk_size)
73            })
74            .collect();
75
76        Self::new(&chunk_specs)
77    }
78
79    /// Builds a XorbObject from chunk information.
80    fn build_xorb_object(chunks: &[RandomChunkInfo]) -> XorbObject {
81        let num_chunks = chunks.len() as u32;
82
83        // Compute XORB hash from chunk hashes
84        let xorb_hash = if chunks.is_empty() {
85            MerkleHash::default()
86        } else {
87            let mut hash_data = Vec::with_capacity(chunks.len() * 32);
88            for chunk in chunks {
89                hash_data.extend_from_slice(chunk.hash.as_bytes());
90            }
91            compute_data_hash(&hash_data)
92        };
93
94        // Collect chunk hashes
95        let chunk_hashes: Vec<MerkleHash> = chunks.iter().map(|c| c.hash).collect();
96
97        // Compute chunk boundary offsets (physical layout with headers)
98        // Each chunk has: header (8 bytes) + data (chunk.size bytes)
99        let mut chunk_boundary_offsets = Vec::with_capacity(num_chunks as usize);
100        let mut cumulative_offset = 0u32;
101        for chunk in chunks {
102            cumulative_offset += XORB_CHUNK_HEADER_LENGTH as u32 + chunk.size;
103            chunk_boundary_offsets.push(cumulative_offset);
104        }
105
106        // Compute unpacked chunk offsets (uncompressed layout without headers)
107        let mut unpacked_chunk_offsets = Vec::with_capacity(num_chunks as usize);
108        let mut cumulative_unpacked = 0u32;
109        for chunk in chunks {
110            cumulative_unpacked += chunk.size;
111            unpacked_chunk_offsets.push(cumulative_unpacked);
112        }
113
114        // Start with default and override the fields we need
115        let mut info = XorbObjectInfoV1::default();
116        info.xorb_hash = xorb_hash;
117        info.chunk_hashes = chunk_hashes;
118        info.chunk_boundary_offsets = chunk_boundary_offsets;
119        info.unpacked_chunk_offsets = unpacked_chunk_offsets;
120        info.num_chunks = num_chunks;
121
122        // Fill in the offset fields
123        info.fill_in_boundary_offsets();
124
125        let info_length = info.serialized_length() as u32;
126
127        XorbObject { info, info_length }
128    }
129
130    /// Generates chunk data from a seed.
131    fn generate_chunk_data_from_seed(seed: u64, size: u32) -> Vec<u8> {
132        let mut rng = SmallRng::seed_from_u64(seed);
133        let mut data = vec![0u8; size as usize];
134        rng.fill_bytes(&mut data);
135        data
136    }
137
138    /// Returns the number of chunks in this XORB.
139    pub fn num_chunks(&self) -> u32 {
140        self.chunks.len() as u32
141    }
142
143    /// Returns the hash of the XORB.
144    pub fn xorb_hash(&self) -> MerkleHash {
145        self.xorb_object.info.xorb_hash
146    }
147
148    /// Returns the hash of a specific chunk.
149    pub fn chunk_hash(&self, idx: u32) -> Option<MerkleHash> {
150        self.chunks.get(idx as usize).map(|c| c.hash)
151    }
152
153    /// Returns the uncompressed size of a specific chunk.
154    pub fn chunk_size(&self, idx: u32) -> Option<u32> {
155        self.chunks.get(idx as usize).map(|c| c.size)
156    }
157
158    /// Returns the total uncompressed size of all chunks.
159    pub fn total_uncompressed_size(&self) -> u64 {
160        self.chunks.iter().map(|c| c.size as u64).sum()
161    }
162
163    /// Returns the total uncompressed size for a range of chunks [start, end).
164    pub fn chunk_range_size(&self, start: u32, end: u32) -> u64 {
165        (start..end).filter_map(|i| self.chunk_size(i)).map(|s| s as u64).sum()
166    }
167
168    /// Returns (hash, size) pairs for a range of chunks [start, end).
169    /// This is useful for computing file hashes.
170    pub fn chunk_hash_sizes(&self, start: u32, end: u32) -> Vec<(MerkleHash, u64)> {
171        (start..end)
172            .filter_map(|i| {
173                let hash = self.chunk_hash(i)?;
174                let size = self.chunk_size(i)? as u64;
175                Some((hash, size))
176            })
177            .collect()
178    }
179
180    /// Returns the chunk hashes for a range of chunks [start, end).
181    pub fn chunk_hashes_range(&self, start: u32, end: u32) -> Vec<MerkleHash> {
182        (start..end).filter_map(|i| self.chunk_hash(i)).collect()
183    }
184
185    /// Generates and returns the raw data for a specific chunk.
186    pub fn get_chunk_data(&self, idx: u32) -> Option<Bytes> {
187        self.chunks
188            .get(idx as usize)
189            .map(|chunk| Bytes::from(Self::generate_chunk_data_from_seed(chunk.seed, chunk.size)))
190    }
191
192    /// Generates and returns the raw data for a range of chunks [start, end).
193    pub fn get_chunk_range_data(&self, start: u32, end: u32) -> Option<Bytes> {
194        if start >= end || end > self.num_chunks() {
195            return None;
196        }
197
198        let mut data = Vec::new();
199        for idx in start..end {
200            let chunk = &self.chunks[idx as usize];
201            let chunk_data = Self::generate_chunk_data_from_seed(chunk.seed, chunk.size);
202            data.extend_from_slice(&chunk_data);
203        }
204        Some(Bytes::from(data))
205    }
206
207    /// Returns the XorbObject header/footer for this XORB.
208    ///
209    /// Uses no compression (CompressionScheme::None) for all chunks.
210    pub fn get_xorb_object(&self) -> XorbObject {
211        self.xorb_object.clone()
212    }
213
214    /// Returns the total serialized length of the XORB (chunks + footer).
215    pub fn serialized_length(&self) -> u64 {
216        let chunks_length: u64 = self
217            .chunks
218            .iter()
219            .map(|c| XORB_CHUNK_HEADER_LENGTH as u64 + c.size as u64)
220            .sum();
221
222        let footer_length = self.xorb_object.info.serialized_length() as u64 + 4; // +4 for info_length u32
223
224        chunks_length + footer_length
225    }
226
227    /// Returns the serialized bytes for a range within the XORB.
228    ///
229    /// This generates the bytes on-the-fly, including chunk headers and data.
230    /// The range is in terms of the serialized byte positions.
231    pub fn get_serialized_range(&self, start: u64, end: u64) -> Bytes {
232        let total_len = self.serialized_length();
233        let end = end.min(total_len);
234
235        if start >= end {
236            return Bytes::new();
237        }
238
239        // Calculate where the footer starts
240        let chunks_length: u64 = self
241            .chunks
242            .iter()
243            .map(|c| XORB_CHUNK_HEADER_LENGTH as u64 + c.size as u64)
244            .sum();
245
246        let mut result = Vec::with_capacity((end - start) as usize);
247
248        // Current position in the serialized stream
249        let mut pos = 0u64;
250
251        // Generate chunk data
252        for chunk in &self.chunks {
253            let chunk_serialized_len = XORB_CHUNK_HEADER_LENGTH as u64 + chunk.size as u64;
254            let chunk_end = pos + chunk_serialized_len;
255
256            if chunk_end > start && pos < end {
257                // This chunk overlaps with our range
258                let header = XorbChunkHeader::new(CompressionScheme::None, chunk.size, chunk.size);
259                let header_bytes = header_to_bytes(&header);
260                let chunk_data = Self::generate_chunk_data_from_seed(chunk.seed, chunk.size);
261
262                // Combine header and data
263                let mut serialized_chunk = Vec::with_capacity(chunk_serialized_len as usize);
264                serialized_chunk.extend_from_slice(&header_bytes);
265                serialized_chunk.extend_from_slice(&chunk_data);
266
267                // Extract the overlapping portion
268                let overlap_start = start.saturating_sub(pos) as usize;
269                let overlap_end = ((end - pos) as usize).min(serialized_chunk.len());
270
271                if overlap_start < overlap_end {
272                    result.extend_from_slice(&serialized_chunk[overlap_start..overlap_end]);
273                }
274            }
275
276            pos = chunk_end;
277            if pos >= end {
278                break;
279            }
280        }
281
282        // Generate footer if needed
283        if end > chunks_length && pos < end {
284            let mut footer_bytes = Vec::new();
285            self.xorb_object.info.serialize(&mut footer_bytes).unwrap();
286            footer_bytes.extend_from_slice(&self.xorb_object.info_length.to_le_bytes());
287
288            let footer_start_in_stream = chunks_length;
289            let overlap_start = start.saturating_sub(footer_start_in_stream) as usize;
290            let overlap_end = ((end - footer_start_in_stream) as usize).min(footer_bytes.len());
291
292            if overlap_start < overlap_end {
293                result.extend_from_slice(&footer_bytes[overlap_start..overlap_end]);
294            }
295        }
296
297        Bytes::from(result)
298    }
299
300    /// Returns the full serialized XORB.
301    ///
302    /// Note: This materializes the entire XORB, which may be large.
303    /// Prefer `get_serialized_range` for partial access.
304    pub fn get_full_serialized(&self) -> Bytes {
305        self.get_serialized_range(0, self.serialized_length())
306    }
307}
308
309/// Converts a XorbChunkHeader to bytes.
310fn header_to_bytes(header: &XorbChunkHeader) -> [u8; XORB_CHUNK_HEADER_LENGTH] {
311    let mut bytes = [0u8; XORB_CHUNK_HEADER_LENGTH];
312    bytes[0] = 0; // version
313    bytes[1..4].copy_from_slice(&header.get_compressed_length().to_le_bytes()[..3]);
314    bytes[4] = CompressionScheme::None as u8;
315    bytes[5..8].copy_from_slice(&header.get_uncompressed_length().to_le_bytes()[..3]);
316    bytes
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[test]
324    fn test_random_xorb_basic() {
325        let specs = vec![(42, 1024), (123, 2048), (456, 512)];
326        let xorb = RandomXorb::new(&specs);
327
328        assert_eq!(xorb.num_chunks(), 3);
329        assert!(xorb.chunk_hash(0).is_some());
330        assert!(xorb.chunk_hash(1).is_some());
331        assert!(xorb.chunk_hash(2).is_some());
332        assert!(xorb.chunk_hash(3).is_none());
333
334        // Verify deterministic generation
335        let data1 = xorb.get_chunk_data(0).unwrap();
336        let data2 = xorb.get_chunk_data(0).unwrap();
337        assert_eq!(data1, data2);
338        assert_eq!(data1.len(), 1024);
339    }
340
341    #[test]
342    fn test_random_xorb_object() {
343        let specs = vec![(1, 100), (2, 200)];
344        let xorb = RandomXorb::new(&specs);
345
346        let xorb_obj = xorb.get_xorb_object();
347        assert_eq!(xorb_obj.info.num_chunks, 2);
348        assert_eq!(xorb_obj.info.chunk_hashes.len(), 2);
349        assert_eq!(xorb_obj.info.chunk_boundary_offsets.len(), 2);
350        assert_eq!(xorb_obj.info.unpacked_chunk_offsets.len(), 2);
351
352        // First chunk: header (8) + data (100) = 108
353        assert_eq!(xorb_obj.info.chunk_boundary_offsets[0], 108);
354        // Second chunk: 108 + header (8) + data (200) = 316
355        assert_eq!(xorb_obj.info.chunk_boundary_offsets[1], 316);
356
357        // Unpacked offsets (no headers)
358        assert_eq!(xorb_obj.info.unpacked_chunk_offsets[0], 100);
359        assert_eq!(xorb_obj.info.unpacked_chunk_offsets[1], 300);
360    }
361
362    #[test]
363    fn test_random_xorb_chunk_range() {
364        let specs = vec![(1, 100), (2, 200), (3, 300)];
365        let xorb = RandomXorb::new(&specs);
366
367        let range_data = xorb.get_chunk_range_data(0, 2).unwrap();
368        assert_eq!(range_data.len(), 300); // 100 + 200
369
370        let chunk0 = xorb.get_chunk_data(0).unwrap();
371        let chunk1 = xorb.get_chunk_data(1).unwrap();
372        assert_eq!(&range_data[..100], &chunk0[..]);
373        assert_eq!(&range_data[100..], &chunk1[..]);
374    }
375
376    #[test]
377    fn test_random_xorb_serialized_length() {
378        let specs = vec![(1, 100)];
379        let xorb = RandomXorb::new(&specs);
380
381        let serialized = xorb.get_full_serialized();
382        assert_eq!(serialized.len() as u64, xorb.serialized_length());
383    }
384}
xet_client/cas_client/simulation/random_xorb.rs

xet_client/cas_client/simulation/
random_xorb.rs
