nectar_primitives/file/
tree.rs

1//! Tree structure calculations for parallel file operations.
2
3use super::constants::{LEVEL_LIMIT, REF_SIZE};
4use crate::bmt::DEFAULT_BODY_SIZE;
5
6/// Tree structure for a file of known size.
7///
8/// Pre-computes chunk counts and spans for efficient parallel operations.
9/// The branching factor defaults to `BODY_SIZE / 32` (plain mode, 128 branches).
10/// Use [`with_ref_size`](Self::with_ref_size) for encrypted mode (`BODY_SIZE / 64` = 64 branches).
11#[derive(Debug, Clone, Copy)]
12pub struct TreeParams<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
13    size: u64,
14    depth: usize,
15    data_chunks: u64,
16    branches: usize,
17}
18
19impl<const BODY_SIZE: usize> TreeParams<BODY_SIZE> {
20    /// Create tree parameters for a file of given size (plain mode, `REF_SIZE = 32`).
21    pub fn new(size: u64) -> Self {
22        Self::with_ref_size(size, REF_SIZE)
23    }
24
25    /// Create tree parameters with a custom reference size (e.g. 64 for encrypted mode).
26    pub fn with_ref_size(size: u64, ref_size: usize) -> Self {
27        let branches = BODY_SIZE / ref_size;
28        let (depth, data_chunks) = if size == 0 {
29            (1, 1) // Empty file still produces one chunk
30        } else {
31            let data_chunks = size.div_ceil(BODY_SIZE as u64);
32            let depth = Self::calculate_depth_with(data_chunks, branches);
33            (depth, data_chunks)
34        };
35
36        Self {
37            size,
38            depth,
39            data_chunks,
40            branches,
41        }
42    }
43
44    /// File size in bytes.
45    pub const fn size(&self) -> u64 {
46        self.size
47    }
48
49    /// Tree depth (1 = single chunk, 2 = one intermediate level, etc.).
50    pub const fn depth(&self) -> usize {
51        self.depth
52    }
53
54    /// Number of data chunks (level 0).
55    pub const fn data_chunks(&self) -> u64 {
56        self.data_chunks
57    }
58
59    /// Total chunks across all levels.
60    pub const fn total_chunks(&self) -> u64 {
61        let mut total = self.data_chunks;
62        let mut chunks_at_level = self.data_chunks;
63
64        while chunks_at_level > 1 {
65            chunks_at_level = chunks_at_level.div_ceil(self.branches as u64);
66            total += chunks_at_level;
67        }
68
69        total
70    }
71
72    /// Chunks at a specific level.
73    pub fn chunks_at_level(&self, level: usize) -> u64 {
74        if level >= self.depth {
75            return 0;
76        }
77        if level == 0 {
78            return self.data_chunks;
79        }
80
81        let mut chunks = self.data_chunks;
82        for _ in 0..level {
83            chunks = chunks.div_ceil(self.branches as u64);
84        }
85        chunks
86    }
87
88    /// Span for a chunk at given level and index.
89    pub fn span_at(&self, level: usize, index: u64) -> u64 {
90        let spans = super::constants::compute_spans_inline(self.branches);
91        let max_span = spans[level] * BODY_SIZE as u64;
92        let chunks_at_level = self.chunks_at_level(level);
93
94        if index + 1 == chunks_at_level {
95            // Last chunk may have smaller span
96            let preceding = index * max_span;
97            let level_total = self.level_span(level);
98            level_total.saturating_sub(preceding)
99        } else {
100            max_span
101        }
102    }
103
104    /// Total bytes covered by chunks at this level.
105    const fn level_span(&self, _level: usize) -> u64 {
106        self.size
107    }
108
109    /// Byte offset for start of chunk at level 0.
110    pub const fn chunk_offset(&self, chunk_index: u64) -> u64 {
111        chunk_index * BODY_SIZE as u64
112    }
113
114    /// Byte range covered by a data chunk.
115    pub fn chunk_range(&self, chunk_index: u64) -> (u64, u64) {
116        let start = self.chunk_offset(chunk_index);
117        let end = (start + BODY_SIZE as u64).min(self.size);
118        (start, end)
119    }
120
121    /// Calculate required data chunks for a byte range.
122    pub fn chunks_for_range(&self, offset: u64, len: u64) -> ChunkRange {
123        if len == 0 || offset >= self.size {
124            return ChunkRange { start: 0, end: 0 };
125        }
126
127        let end_offset = (offset + len).min(self.size);
128        let start_chunk = offset / BODY_SIZE as u64;
129        let end_chunk = end_offset.div_ceil(BODY_SIZE as u64);
130
131        ChunkRange {
132            start: start_chunk,
133            end: end_chunk.min(self.data_chunks),
134        }
135    }
136
137    fn calculate_depth_with(data_chunks: u64, branches: usize) -> usize {
138        if data_chunks <= 1 {
139            return 1;
140        }
141
142        let mut depth = 1;
143        let mut chunks = data_chunks;
144        while chunks > 1 {
145            chunks = chunks.div_ceil(branches as u64);
146            depth += 1;
147        }
148        depth.min(LEVEL_LIMIT)
149    }
150}
151
152/// Range of chunk indices.
153#[derive(Debug, Clone, Copy, PartialEq, Eq)]
154pub struct ChunkRange {
155    /// First chunk index (inclusive).
156    pub start: u64,
157    /// Last chunk index (exclusive).
158    pub end: u64,
159}
160
161impl ChunkRange {
162    /// Number of chunks in range.
163    pub const fn len(&self) -> u64 {
164        self.end.saturating_sub(self.start)
165    }
166
167    /// Whether the range is empty.
168    pub const fn is_empty(&self) -> bool {
169        self.start >= self.end
170    }
171
172    /// Iterate over chunk indices.
173    pub fn iter(&self) -> impl Iterator<Item = u64> {
174        self.start..self.end
175    }
176}
177
178/// Assemble decoded chunk bodies into a contiguous output buffer.
179///
180/// Given a set of chunk bodies corresponding to `chunk_range`, copies the
181/// relevant byte slices into a single output buffer accounting for partial
182/// reads at the start/end of the range.
183pub(crate) fn assemble_range<const BODY_SIZE: usize>(
184    tree: &TreeParams<BODY_SIZE>,
185    offset: u64,
186    actual_len: usize,
187    chunk_range: &ChunkRange,
188    bodies: &[bytes::Bytes],
189) -> Vec<u8> {
190    let mut result = vec![0u8; actual_len];
191    let mut result_offset = 0;
192
193    for (i, chunk_idx) in chunk_range.iter().enumerate() {
194        let (chunk_start, chunk_end) = tree.chunk_range(chunk_idx);
195        let chunk_data_len = (chunk_end - chunk_start) as usize;
196
197        let read_start = if chunk_start < offset {
198            (offset - chunk_start) as usize
199        } else {
200            0
201        };
202
203        let read_end = if chunk_end > offset + actual_len as u64 {
204            chunk_data_len - ((chunk_end - offset - actual_len as u64) as usize)
205        } else {
206            chunk_data_len
207        };
208
209        let bytes_to_copy = read_end - read_start;
210        let body = &bodies[i];
211
212        result[result_offset..result_offset + bytes_to_copy]
213            .copy_from_slice(&body[read_start..read_end]);
214        result_offset += bytes_to_copy;
215    }
216
217    result
218}
219
220/// Calculate subspan size for children of a node with given span (plain mode).
221#[cfg(test)]
222fn subspan_size<const BODY_SIZE: usize>(span: u64) -> u64 {
223    let spans = super::constants::compute_spans_inline(BODY_SIZE / super::constants::REF_SIZE);
224    super::constants::subspan_for_spans::<BODY_SIZE>(span, &spans)
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn test_tree_params_empty() {
233        let tree = TreeParams::<DEFAULT_BODY_SIZE>::new(0);
234        assert_eq!(tree.size(), 0);
235        assert_eq!(tree.depth(), 1);
236        assert_eq!(tree.data_chunks(), 1);
237        assert_eq!(tree.total_chunks(), 1);
238    }
239
240    #[test]
241    fn test_tree_params_single_chunk() {
242        let tree = TreeParams::<DEFAULT_BODY_SIZE>::new(100);
243        assert_eq!(tree.depth(), 1);
244        assert_eq!(tree.data_chunks(), 1);
245        assert_eq!(tree.total_chunks(), 1);
246        assert_eq!(tree.chunks_at_level(0), 1);
247        assert_eq!(tree.chunks_at_level(1), 0);
248    }
249
250    #[test]
251    fn test_tree_params_two_chunks() {
252        let tree = TreeParams::<DEFAULT_BODY_SIZE>::new(DEFAULT_BODY_SIZE as u64 + 1);
253        assert_eq!(tree.depth(), 2);
254        assert_eq!(tree.data_chunks(), 2);
255        assert_eq!(tree.total_chunks(), 3); // 2 data + 1 intermediate
256        assert_eq!(tree.chunks_at_level(0), 2);
257        assert_eq!(tree.chunks_at_level(1), 1);
258    }
259
260    #[test]
261    fn test_tree_params_128_chunks() {
262        let size = DEFAULT_BODY_SIZE as u64 * 128;
263        let tree = TreeParams::<DEFAULT_BODY_SIZE>::new(size);
264        assert_eq!(tree.depth(), 2);
265        assert_eq!(tree.data_chunks(), 128);
266        assert_eq!(tree.chunks_at_level(0), 128);
267        assert_eq!(tree.chunks_at_level(1), 1);
268    }
269
270    #[test]
271    fn test_tree_params_129_chunks() {
272        let size = DEFAULT_BODY_SIZE as u64 * 129;
273        let tree = TreeParams::<DEFAULT_BODY_SIZE>::new(size);
274        assert_eq!(tree.depth(), 3);
275        assert_eq!(tree.data_chunks(), 129);
276        assert_eq!(tree.chunks_at_level(0), 129);
277        assert_eq!(tree.chunks_at_level(1), 2);
278        assert_eq!(tree.chunks_at_level(2), 1);
279    }
280
281    #[test]
282    fn test_chunk_range() {
283        let tree = TreeParams::<DEFAULT_BODY_SIZE>::new(10000);
284
285        // First chunk
286        let (start, end) = tree.chunk_range(0);
287        assert_eq!(start, 0);
288        assert_eq!(end, 4096);
289
290        // Second chunk
291        let (start, end) = tree.chunk_range(1);
292        assert_eq!(start, 4096);
293        assert_eq!(end, 8192);
294
295        // Last partial chunk
296        let (start, end) = tree.chunk_range(2);
297        assert_eq!(start, 8192);
298        assert_eq!(end, 10000);
299    }
300
301    #[test]
302    fn test_chunks_for_range() {
303        let tree = TreeParams::<DEFAULT_BODY_SIZE>::new(10000);
304
305        // Read from start
306        let range = tree.chunks_for_range(0, 100);
307        assert_eq!(range.start, 0);
308        assert_eq!(range.end, 1);
309        assert_eq!(range.len(), 1);
310
311        // Read spanning chunks
312        let range = tree.chunks_for_range(4000, 200);
313        assert_eq!(range.start, 0);
314        assert_eq!(range.end, 2);
315        assert_eq!(range.len(), 2);
316
317        // Read at end
318        let range = tree.chunks_for_range(9000, 1000);
319        assert_eq!(range.start, 2);
320        assert_eq!(range.end, 3);
321    }
322
323    #[test]
324    fn test_span_at() {
325        let size = DEFAULT_BODY_SIZE as u64 * 3;
326        let tree = TreeParams::<DEFAULT_BODY_SIZE>::new(size);
327
328        // Full chunks have full span
329        assert_eq!(tree.span_at(0, 0), DEFAULT_BODY_SIZE as u64);
330        assert_eq!(tree.span_at(0, 1), DEFAULT_BODY_SIZE as u64);
331        assert_eq!(tree.span_at(0, 2), DEFAULT_BODY_SIZE as u64);
332    }
333
334    #[test]
335    fn test_subspan_size() {
336        // Single chunk span -> subspan is chunk size
337        assert_eq!(subspan_size::<DEFAULT_BODY_SIZE>(4096), 4096);
338
339        // Two chunk span -> subspan is chunk size
340        assert_eq!(subspan_size::<DEFAULT_BODY_SIZE>(8000), 4096);
341
342        // Large span -> subspan is previous level
343        let large_span = 128 * DEFAULT_BODY_SIZE as u64 + 1;
344        assert_eq!(subspan_size::<DEFAULT_BODY_SIZE>(large_span), 128 * 4096);
345    }
346
347    #[test]
348    fn test_encrypted_tree_params() {
349        use super::super::constants::ENCRYPTED_REF_SIZE;
350
351        // 64 data chunks fills one encrypted intermediate exactly
352        let size = DEFAULT_BODY_SIZE as u64 * 64;
353        let tree = TreeParams::<DEFAULT_BODY_SIZE>::with_ref_size(size, ENCRYPTED_REF_SIZE);
354        assert_eq!(tree.depth(), 2); // 64 chunks / 64 branches = 1 intermediate
355        assert_eq!(tree.data_chunks(), 64);
356        assert_eq!(tree.chunks_at_level(0), 64);
357        assert_eq!(tree.chunks_at_level(1), 1);
358        assert_eq!(tree.total_chunks(), 65);
359
360        // 65 data chunks needs a third level
361        let size2 = DEFAULT_BODY_SIZE as u64 * 65;
362        let tree2 = TreeParams::<DEFAULT_BODY_SIZE>::with_ref_size(size2, ENCRYPTED_REF_SIZE);
363        assert_eq!(tree2.depth(), 3);
364        assert_eq!(tree2.chunks_at_level(0), 65);
365        assert_eq!(tree2.chunks_at_level(1), 2);
366        assert_eq!(tree2.chunks_at_level(2), 1);
367    }
368}
nectar_primitives/file/tree.rs

nectar_primitives/file/
tree.rs
