nectar_primitives/file/
splitter.rs1use std::fmt;
8use std::io::{self, Write};
9use std::marker::PhantomData;
10
11use crate::bmt::DEFAULT_BODY_SIZE;
12use crate::chunk::AnyChunk;
13
14use super::error::{FileError, Result};
15use super::mode::{PlainMode, SplitMode};
16use super::splitter_parallel::GenericParallelSplitter;
17
18#[cfg(feature = "encryption")]
19use super::mode::EncryptedMode;
20
21pub struct GenericSplitter<M: SplitMode, const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
26 span_length: u64,
27 buffer: Vec<u8>,
28 _mode: PhantomData<M>,
29}
30
31pub type Splitter<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> =
33 GenericSplitter<PlainMode, BODY_SIZE>;
34
35#[cfg(feature = "encryption")]
37pub type EncryptedSplitter<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> =
38 GenericSplitter<EncryptedMode, BODY_SIZE>;
39
40impl<M, const BODY_SIZE: usize> fmt::Debug for GenericSplitter<M, BODY_SIZE>
41where
42 M: SplitMode,
43{
44 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45 f.debug_struct("GenericSplitter")
46 .field("span_length", &self.span_length)
47 .field("length", &self.buffer.len())
48 .finish_non_exhaustive()
49 }
50}
51
52impl<M, const BODY_SIZE: usize> GenericSplitter<M, BODY_SIZE>
53where
54 M: SplitMode,
55{
56 pub fn new(span_length: u64) -> Self {
58 const { super::constants::assert_valid_body_size::<BODY_SIZE>() };
59
60 Self {
61 span_length,
62 buffer: Vec::with_capacity(span_length.min(BODY_SIZE as u64 * 2) as usize),
63 _mode: PhantomData,
64 }
65 }
66
67 pub const fn len(&self) -> u64 {
69 self.buffer.len() as u64
70 }
71
72 pub const fn is_empty(&self) -> bool {
74 self.buffer.is_empty()
75 }
76
77 pub const fn span_length(&self) -> u64 {
79 self.span_length
80 }
81}
82
83impl<M, const BODY_SIZE: usize> GenericSplitter<M, BODY_SIZE>
84where
85 M: SplitMode + Send + Sync,
86{
87 pub fn finish(self) -> Result<(M::RootRef, Vec<AnyChunk<BODY_SIZE>>)> {
89 if self.buffer.len() as u64 != self.span_length {
90 return Err(FileError::SpanMismatch {
91 expected: self.span_length,
92 actual: self.buffer.len() as u64,
93 });
94 }
95
96 if self.buffer.is_empty() {
97 let (chunk, root) = M::empty_chunk::<BODY_SIZE>()?;
98 return Ok((root, vec![chunk.into()]));
99 }
100
101 GenericParallelSplitter::<M, BODY_SIZE>::split_to_vec(&self.buffer)
102 }
103}
104
105impl<M, const BODY_SIZE: usize> Write for GenericSplitter<M, BODY_SIZE>
106where
107 M: SplitMode,
108{
109 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
110 let remaining = self.span_length.saturating_sub(self.buffer.len() as u64) as usize;
111 let to_write = buf.len().min(remaining);
112 if to_write == 0 && !buf.is_empty() {
113 return Err(io::Error::other(
114 FileError::WritePastSpan {
115 span: self.span_length,
116 written: self.span_length + 1,
117 }
118 .to_string(),
119 ));
120 }
121 self.buffer.extend_from_slice(&buf[..to_write]);
122 Ok(to_write)
123 }
124
125 fn flush(&mut self) -> io::Result<()> {
126 Ok(())
127 }
128}
129
130#[cfg(test)]
131mod tests {
132 use super::*;
133 use crate::store::MemoryStore;
134
135 fn split_and_store(
136 data: &[u8],
137 ) -> (crate::chunk::ChunkAddress, MemoryStore<DEFAULT_BODY_SIZE>) {
138 let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
139 splitter.write_all(data).unwrap();
140 let (root, chunks) = splitter.finish().unwrap();
141 (root, MemoryStore::from_chunks(chunks))
142 }
143
144 generate_plain_splitter_tests!(split_and_store);
145
146 #[test]
147 fn test_splitter_incremental_writes() {
148 let mut data = vec![0u8; DEFAULT_BODY_SIZE * 2 + 100];
149 rand::RngExt::fill(&mut rand::rng(), &mut data);
150 let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
151
152 for chunk in data.chunks(100) {
153 splitter.write_all(chunk).unwrap();
154 }
155 let (root, chunks) = splitter.finish().unwrap();
156 let store = MemoryStore::from_chunks(chunks);
157
158 assert_eq!(store.len(), 4);
159 assert!(!root.is_zero());
160 }
161
162 #[test]
163 fn test_splitter_deterministic() {
164 let data = vec![0x56; DEFAULT_BODY_SIZE * 3];
165
166 let root1 = {
167 let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
168 splitter.write_all(&data).unwrap();
169 splitter.finish().unwrap().0
170 };
171
172 let root2 = {
173 let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
174 splitter.write_all(&data).unwrap();
175 splitter.finish().unwrap().0
176 };
177
178 assert_eq!(root1, root2);
179 }
180
181 #[test]
182 fn test_splitter_write_past_span() {
183 let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(10);
184
185 let result = splitter.write_all(b"this is more than 10 bytes");
186 assert!(result.is_err());
187 }
188
189 #[test]
190 fn test_splitter_span_mismatch() {
191 let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(100);
192
193 splitter.write_all(b"short").unwrap();
194 let result = splitter.finish();
195
196 assert!(matches!(result, Err(FileError::SpanMismatch { .. })));
197 }
198
199 #[cfg(feature = "encryption")]
200 mod encrypted {
201 use super::*;
202
203 fn encrypted_split_and_store(
204 data: &[u8],
205 ) -> (
206 crate::chunk::encryption::EncryptedChunkRef,
207 MemoryStore<DEFAULT_BODY_SIZE>,
208 ) {
209 let mut splitter = EncryptedSplitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
210 splitter.write_all(data).unwrap();
211 let (root_ref, chunks) = splitter.finish().unwrap();
212 (root_ref, MemoryStore::from_chunks(chunks))
213 }
214
215 generate_encrypted_splitter_tests!(encrypted_split_and_store);
216
217 #[test]
218 fn test_encrypted_splitter_write_past_span() {
219 let mut splitter = EncryptedSplitter::<DEFAULT_BODY_SIZE>::new(10);
220
221 let result = splitter.write_all(b"this is more than 10 bytes");
222 assert!(result.is_err());
223 }
224
225 #[test]
226 fn test_encrypted_splitter_span_mismatch() {
227 let mut splitter = EncryptedSplitter::<DEFAULT_BODY_SIZE>::new(100);
228
229 splitter.write_all(b"short").unwrap();
230 let result = splitter.finish();
231
232 assert!(matches!(result, Err(FileError::SpanMismatch { .. })));
233 }
234
235 #[test]
236 fn test_encrypted_differs_from_plaintext() {
237 let data = b"test data for encryption comparison";
238 let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
239 splitter.write_all(data).unwrap();
240 let (plain_root, _) = splitter.finish().unwrap();
241
242 let mut enc_splitter = EncryptedSplitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
243 enc_splitter.write_all(data).unwrap();
244 let (enc_root, _) = enc_splitter.finish().unwrap();
245
246 assert_ne!(enc_root.address(), &plain_root);
247 }
248 }
249}