nectar_primitives/file/
fold.rs

1//! Closure-driven fold combinators over the async joiner.
2//!
3//! Thin adapters that let a consumer scan the whole file through a closure
4//! without buffering it: [`for_each_chunk`](GenericJoiner::for_each_chunk)
5//! visits each leaf out of order at full throughput, and
6//! [`try_for_each_window`](GenericJoiner::try_for_each_window) folds contiguous
7//! in-order byte runs with bounded memory. Both compose over the existing
8//! streaming methods, so neither re-walks the tree.
9
10use core::ops::ControlFlow;
11
12use bytes::Bytes;
13use futures::stream::StreamExt;
14
15use super::error::Result;
16use super::joiner::GenericJoiner;
17use super::mode::JoinMode;
18use crate::store::{ChunkGet, MaybeSend};
19
20impl<G, M, const BODY_SIZE: usize> GenericJoiner<G, M, BODY_SIZE>
21where
22    G: ChunkGet<BODY_SIZE> + 'static,
23    M: JoinMode + MaybeSend + Sync,
24{
25    /// Visit each leaf body as it lands, out of order, tagged with its absolute
26    /// offset. Peak memory is the in-flight width. `f` returns `Break` to stop
27    /// early, which drops the stream and cancels in-flight fetches.
28    pub async fn for_each_chunk<F>(self, mut f: F) -> Result<()>
29    where
30        F: FnMut(u64, &Bytes) -> ControlFlow<()>,
31    {
32        let s = self.into_offset_stream_chunked();
33        futures::pin_mut!(s);
34        while let Some(p) = s.next().await {
35            let (off, body) = p?;
36            if f(off, &body).is_break() {
37                break;
38            }
39        }
40        Ok(())
41    }
42
43    /// Fold contiguous in-order byte runs through `f` without buffering the
44    /// file. `window` bounds peak memory (see
45    /// [`into_windowed_reader`](Self::into_windowed_reader)). `f` returns
46    /// `Break(acc)` to stop early with the carried accumulator; otherwise the
47    /// final accumulator is returned at stream end.
48    pub async fn try_for_each_window<B, F>(self, window: usize, init: B, mut f: F) -> Result<B>
49    where
50        F: FnMut(B, &[u8]) -> ControlFlow<B, B>,
51    {
52        let mut reader = self.into_windowed_reader(window);
53        let s = reader.stream();
54        futures::pin_mut!(s);
55        let mut acc = init;
56        while let Some(run) = s.next().await {
57            let bytes = run?;
58            match f(acc, &bytes) {
59                ControlFlow::Continue(b) => acc = b,
60                ControlFlow::Break(b) => return Ok(b),
61            }
62        }
63        Ok(acc)
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use crate::bmt::DEFAULT_BODY_SIZE;
71    use crate::chunk::{AnyChunk, ChunkAddress};
72    use crate::file::Joiner;
73    use crate::file::split;
74    use futures::executor::block_on;
75    use std::collections::HashMap;
76
77    fn split_and_store(data: &[u8]) -> (ChunkAddress, HashMap<ChunkAddress, AnyChunk>) {
78        let (root, store) = split::<DEFAULT_BODY_SIZE>(data).unwrap();
79        (root, store.into_chunks())
80    }
81
82    /// `for_each_chunk` visits every leaf exactly once; reassembling by offset
83    /// equals `read_all`.
84    fn assert_for_each_chunk_matches(data: &[u8], width: usize) {
85        let (root, store) = split_and_store(data);
86        block_on(async {
87            let expected = Joiner::new(store.clone(), root)
88                .await
89                .unwrap()
90                .read_all()
91                .await
92                .unwrap();
93
94            let joiner = Joiner::new(store, root)
95                .await
96                .unwrap()
97                .with_concurrency(width);
98            let total = joiner.size();
99
100            let mut reassembled = vec![0u8; total as usize];
101            let mut covered = 0u64;
102            let mut seen = std::collections::HashSet::new();
103            joiner
104                .for_each_chunk(|off, body| {
105                    assert!(seen.insert(off), "offset {off} visited more than once");
106                    let start = off as usize;
107                    reassembled[start..start + body.len()].copy_from_slice(body);
108                    covered += body.len() as u64;
109                    ControlFlow::Continue(())
110                })
111                .await
112                .unwrap();
113
114            assert_eq!(covered, total, "every byte covered exactly once");
115            assert_eq!(reassembled, expected, "reassembly equals read_all");
116            assert_eq!(reassembled, data, "reassembly equals input");
117        });
118    }
119
120    #[test]
121    fn for_each_chunk_reassembles() {
122        let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 3 + 123)
123            .map(|i| (i % 256) as u8)
124            .collect();
125        assert_for_each_chunk_matches(b"hello world", 8);
126        assert_for_each_chunk_matches(&data, 8);
127        // width 1 (degenerate concurrent path) still visits every leaf.
128        assert_for_each_chunk_matches(&data, 1);
129    }
130
131    #[test]
132    fn for_each_chunk_break_stops_early() {
133        let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5)
134            .map(|i| (i % 256) as u8)
135            .collect();
136        let (root, store) = split_and_store(&data);
137        block_on(async {
138            let joiner = Joiner::new(store, root).await.unwrap();
139            let mut visited = 0usize;
140            let res = joiner
141                .for_each_chunk(|_off, _body| {
142                    visited += 1;
143                    if visited == 2 {
144                        ControlFlow::Break(())
145                    } else {
146                        ControlFlow::Continue(())
147                    }
148                })
149                .await;
150            assert!(res.is_ok(), "early break returns Ok");
151            assert_eq!(visited, 2, "stops on the breaking leaf");
152        });
153    }
154
155    /// Concatenating the in-order runs equals the file.
156    fn assert_window_concat_matches(data: &[u8], window: usize) {
157        let (root, store) = split_and_store(data);
158        block_on(async {
159            let joiner = Joiner::new(store, root).await.unwrap();
160            let out = joiner
161                .try_for_each_window(window, Vec::new(), |mut acc: Vec<u8>, run| {
162                    acc.extend_from_slice(run);
163                    ControlFlow::Continue(acc)
164                })
165                .await
166                .unwrap();
167            assert_eq!(out, data, "window concat equals file (window {window})");
168        });
169    }
170
171    #[test]
172    fn try_for_each_window_concat_equals_file() {
173        let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 321)
174            .map(|i| (i % 256) as u8)
175            .collect();
176        assert_window_concat_matches(&data, 16);
177        // window 1 is the tightest in-order path.
178        assert_window_concat_matches(&data, 1);
179    }
180
181    #[test]
182    fn try_for_each_window_byte_count_equals_size() {
183        let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 4 + 7)
184            .map(|i| (i % 256) as u8)
185            .collect();
186        let (root, store) = split_and_store(&data);
187        block_on(async {
188            let joiner = Joiner::new(store, root).await.unwrap();
189            let size = joiner.size();
190            let count = joiner
191                .try_for_each_window(8, 0u64, |acc, run| {
192                    ControlFlow::Continue(acc + run.len() as u64)
193                })
194                .await
195                .unwrap();
196            assert_eq!(count, size, "folded byte count equals size()");
197        });
198    }
199
200    #[test]
201    fn try_for_each_window_break_returns_accumulator() {
202        let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5)
203            .map(|i| (i % 256) as u8)
204            .collect();
205        let (root, store) = split_and_store(&data);
206        block_on(async {
207            let joiner = Joiner::new(store, root).await.unwrap();
208            // Stop once at least one body has been seen; the carried count comes
209            // back through `Break`.
210            let count = joiner
211                .try_for_each_window(8, 0u64, |acc, _run| {
212                    let acc = acc + 1;
213                    ControlFlow::Break(acc)
214                })
215                .await
216                .unwrap();
217            assert_eq!(count, 1, "break carries the accumulator and stops");
218        });
219    }
220
221    #[cfg(feature = "encryption")]
222    mod encrypted {
223        use super::*;
224        use crate::file::EncryptedJoiner;
225        use crate::file::split_encrypted;
226
227        fn enc_split_and_store(
228            data: &[u8],
229        ) -> (
230            crate::chunk::encryption::EncryptedChunkRef,
231            HashMap<ChunkAddress, AnyChunk>,
232        ) {
233            let (root_ref, store) = split_encrypted::<DEFAULT_BODY_SIZE>(data).unwrap();
234            (root_ref, store.into_chunks())
235        }
236
237        #[test]
238        fn for_each_chunk_reassembles_encrypted() {
239            let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 3 + 123)
240                .map(|i| (i % 256) as u8)
241                .collect();
242            let (root_ref, store) = enc_split_and_store(&data);
243            block_on(async {
244                let expected = EncryptedJoiner::new(store.clone(), root_ref.clone())
245                    .await
246                    .unwrap()
247                    .read_all()
248                    .await
249                    .unwrap();
250
251                let joiner = EncryptedJoiner::new(store, root_ref).await.unwrap();
252                let total = joiner.size();
253                let mut reassembled = vec![0u8; total as usize];
254                let mut seen = std::collections::HashSet::new();
255                joiner
256                    .for_each_chunk(|off, body| {
257                        assert!(seen.insert(off), "offset {off} visited more than once");
258                        let start = off as usize;
259                        reassembled[start..start + body.len()].copy_from_slice(body);
260                        ControlFlow::Continue(())
261                    })
262                    .await
263                    .unwrap();
264                assert_eq!(
265                    reassembled, expected,
266                    "encrypted reassembly equals read_all"
267                );
268            });
269        }
270
271        #[test]
272        fn try_for_each_window_concat_equals_file_encrypted() {
273            let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 321)
274                .map(|i| (i % 256) as u8)
275                .collect();
276            let (root_ref, store) = enc_split_and_store(&data);
277            block_on(async {
278                let joiner = EncryptedJoiner::new(store, root_ref).await.unwrap();
279                let out = joiner
280                    .try_for_each_window(16, Vec::new(), |mut acc: Vec<u8>, run| {
281                        acc.extend_from_slice(run);
282                        ControlFlow::Continue(acc)
283                    })
284                    .await
285                    .unwrap();
286                assert_eq!(out, data, "encrypted window concat equals file");
287            });
288        }
289    }
290}
nectar_primitives/file/fold.rs

nectar_primitives/file/
fold.rs
