Skip to main content

mmap_chunker_core/
lib.rs

1pub mod ffi;
2pub mod mmap;
3pub mod scanner;
4
5mod plan;
6
7pub use ffi::{
8    CChunkView, CEngineHandle, ABI_VERSION, CAP_CONFIGURABLE_DELIMITER, CAP_ERROR_STRINGS,
9    CAP_FIXED_SIZE_CHUNKING, CAP_MULTI_BYTE_DELIMITER, CAP_RECORD_PARTITIONING, CAP_ZERO_COPY,
10};
11pub use mmap::MmapFile;
12pub use scanner::ChunkCursor;
13pub use scanner::PatternChunkCursor;
14
15use std::io;
16use std::path::Path;
17
18use plan::ChunkPlan;
19
20/// Safe Rust interface for memory-mapped file chunking.
21///
22/// Wraps [`MmapFile`] with a chunk-layout state machine supporting
23/// three scan modes: delimited, fixed-size, and record-aligned
24/// partition planning.
25///
26/// # Safety
27///
28/// The constructor [`MmapChunker::open`] is `unsafe` because
29/// file-backed memory mappings can violate Rust's `&[u8]` immutability
30/// guarantee if the underlying file is mutated concurrently. Every
31/// other method on this type is safe after construction.
32///
33/// # Example
34///
35/// ```no_run
36/// use std::io;
37/// # fn main() -> io::Result<()> {
38/// let mut file = unsafe {
39///     mmap_chunker_core::MmapChunker::open("records.jsonl")?
40/// };
41/// let count = file.scan_delimited(64 * 1024, b'\n');
42/// for i in 0..count {
43///     if let Some(chunk) = file.get_chunk(i) {
44///         let _data: &[u8] = chunk;
45///     }
46/// }
47/// # Ok(())
48/// # }
49/// ```
50#[derive(Debug)]
51pub struct MmapChunker {
52    mmap: MmapFile,
53    plan: ChunkPlan,
54}
55
56impl MmapChunker {
57    /// Open and memory-map the file at `path` for read-only chunked
58    /// access.
59    ///
60    /// Accepts any type that converts to `Path` (`&str`, `&Path`,
61    /// `PathBuf`, `&OsStr`). On Windows the path is encoded as UTF-16
62    /// directly; on Unix the raw OS path bytes are used.
63    ///
64    /// # Safety
65    ///
66    /// The caller must ensure that the backing file is not modified,
67    /// truncated, deleted, or otherwise invalidated for the entire
68    /// lifetime of this `MmapChunker` and all `&[u8]` slices derived
69    /// from it (via [`get_chunk`](Self::get_chunk) or
70    /// [`as_bytes`](Self::as_bytes)).
71    ///
72    /// Concurrent file mutation by any process — including the calling
73    /// process — violates the immutability guarantee of `&[u8]` and is
74    /// Rust undefined behavior.
75    ///
76    /// On POSIX systems, another process may freely open the same file
77    /// for writing. Use external synchronization (file locks,
78    /// snapshots, or immutable files) to satisfy this contract. On
79    /// Windows, `FILE_SHARE_READ` prevents other processes from
80    /// opening the file for writing, but same-process mutation remains
81    /// possible.
82    pub unsafe fn open(path: impl AsRef<Path>) -> io::Result<Self> {
83        let mmap = MmapFile::open_path(path)?;
84        Ok(Self {
85            mmap,
86            plan: ChunkPlan::empty(),
87        })
88    }
89
90    /// Returns the number of chunks in the current layout.
91    ///
92    /// Returns 0 if no scan has been performed or if the file is empty.
93    #[inline]
94    pub fn chunk_count(&self) -> usize {
95        self.plan.len()
96    }
97
98    /// Scan the file with the given approximate chunk size and
99    /// single-byte delimiter.
100    ///
101    /// Chunk boundaries are placed at or after each `chunk_size`
102    /// interval, snapped to the next occurrence of `delimiter`.
103    /// The last chunk extends to EOF.
104    ///
105    /// Replaces any previous layout. Returns the number of chunks.
106    pub fn scan_delimited(&mut self, chunk_size: usize, delimiter: u8) -> usize {
107        let data = self.mmap.as_bytes();
108        if data.is_empty() {
109            self.plan = ChunkPlan::empty();
110            return 0;
111        }
112        let chunks = scanner::find_chunk_boundaries(data, chunk_size, delimiter);
113        self.plan = ChunkPlan::from_ranges(chunks);
114        self.plan.len()
115    }
116
117    /// Partition the file into sequential fixed-size chunks.
118    ///
119    /// Chunks are at exact `chunk_size` intervals with the last chunk
120    /// potentially shorter at EOF. No delimiter scanning.
121    ///
122    /// Replaces any previous layout. Returns the number of chunks.
123    pub fn scan_fixed(&mut self, chunk_size: usize) -> usize {
124        let file_len = self.mmap.len();
125        self.plan = ChunkPlan::fixed(file_len, chunk_size);
126        self.plan.len()
127    }
128
129    /// Plan record-aligned partition byte ranges for N-way parallel
130    /// consumers.
131    ///
132    /// Computes approximately balanced byte ranges where every
133    /// partition boundary falls on a record boundary (immediately
134    /// after `delimiter`), ensuring no record is split.
135    ///
136    /// Actual partition count may be less than `num_partitions` if
137    /// giant records span multiple ideal target positions.
138    ///
139    /// Replaces any previous layout. Returns the number of partitions.
140    pub fn partition_records(&mut self, num_partitions: usize, delimiter: u8) -> usize {
141        let data = self.mmap.as_bytes();
142        let file_len = data.len();
143        if file_len == 0 || num_partitions == 0 {
144            self.plan = ChunkPlan::empty();
145            return 0;
146        }
147        let partitions = scanner::find_partition_boundaries(data, num_partitions, delimiter);
148        self.plan = ChunkPlan::from_ranges(partitions);
149        self.plan.len()
150    }
151
152    /// Create a lazy streaming cursor for sequential chunk consumption.
153    ///
154    /// Returns a [`ChunkCursor`] that yields chunks one at a time using
155    /// the same delimiter-aware boundary semantics as
156    /// [`scan_delimited`](Self::scan_delimited), but without
157    /// pre-computing a `Vec` of all boundaries.
158    ///
159    /// O(1) state (~40 bytes on 64-bit) regardless of file size.
160    /// Ideal for low-memory streaming consumers where random access
161    /// via [`get_chunk`](Self::get_chunk) is not needed.
162    ///
163    /// # Example
164    ///
165    /// ```no_run
166    /// use mmap_chunker_core::MmapChunker;
167    ///
168    /// let file = unsafe { MmapChunker::open("records.jsonl")? };
169    /// for chunk in file.delimited_cursor(64 * 1024, b'\n') {
170    ///     let _data: &[u8] = chunk;
171    /// }
172    /// # Ok::<(), std::io::Error>(())
173    /// ```
174    #[inline]
175    pub fn delimited_cursor(&self, chunk_size: usize, delimiter: u8) -> ChunkCursor<'_> {
176        ChunkCursor::new(self.as_bytes(), chunk_size, delimiter)
177    }
178
179    /// Scan with a multi-byte delimiter (e.g., `b"\r\n"` for CRLF).
180    ///
181    /// Same semantics as [`scan_delimited`](Self::scan_delimited) but
182    /// the delimiter can be multiple bytes. Chunk boundaries are placed
183    /// immediately after the complete delimiter.
184    ///
185    /// When `delimiter.len() == 1`, this produces identical results to
186    /// the single-byte path. Delegates to the SWAR fast path internally.
187    ///
188    /// # Panics
189    ///
190    /// Panics if `delimiter` is empty.
191    pub fn scan_delimited_pattern(&mut self, chunk_size: usize, delimiter: &[u8]) -> usize {
192        let data = self.mmap.as_bytes();
193        if data.is_empty() {
194            self.plan = ChunkPlan::empty();
195            return 0;
196        }
197        let chunks = scanner::find_chunk_boundaries_pattern(data, chunk_size, delimiter);
198        self.plan = ChunkPlan::from_ranges(chunks);
199        self.plan.len()
200    }
201
202    /// Create a lazy streaming cursor with a multi-byte delimiter.
203    ///
204    /// Returns a [`PatternChunkCursor`] — same O(1) memory semantics
205    /// as [`delimited_cursor`](Self::delimited_cursor), but for
206    /// multi-byte delimiters like `b"\r\n"`.
207    ///
208    /// # Panics
209    ///
210    /// Panics if `delimiter` is empty.
211    ///
212    /// # Example
213    ///
214    /// ```no_run
215    /// use mmap_chunker_core::MmapChunker;
216    ///
217    /// let file = unsafe { MmapChunker::open("records.jsonl")? };
218    /// for chunk in file.delimited_cursor_pattern(64 * 1024, b"\r\n") {
219    ///     let _data: &[u8] = chunk;
220    /// }
221    /// # Ok::<(), std::io::Error>(())
222    /// ```
223    #[inline]
224    pub fn delimited_cursor_pattern<'a>(
225        &'a self,
226        chunk_size: usize,
227        delimiter: &'a [u8],
228    ) -> PatternChunkCursor<'a, 'a> {
229        PatternChunkCursor::new(self.as_bytes(), chunk_size, delimiter)
230    }
231
232    /// Retrieve a zero-copy chunk by index.
233    ///
234    /// Returns `Some(&[u8])` pointing directly into the mapped file,
235    /// or `None` if the index is out of bounds or no scan has been
236    /// performed.
237    ///
238    /// The returned slice is valid for the lifetime of `self`.
239    pub fn get_chunk(&self, index: usize) -> Option<&[u8]> {
240        let data = self.mmap.as_bytes();
241        let (start, end) = self.plan.range_at(index, data.len())?;
242        Some(&data[start..end])
243    }
244
245    /// Returns the mapped file contents as a byte slice.
246    #[inline]
247    pub fn as_bytes(&self) -> &[u8] {
248        self.mmap.as_bytes()
249    }
250
251    /// Returns the file size in bytes.
252    #[inline]
253    pub fn len(&self) -> usize {
254        self.mmap.len()
255    }
256
257    /// Returns `true` if the file is empty.
258    #[inline]
259    pub fn is_empty(&self) -> bool {
260        self.mmap.is_empty()
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    fn temp_file(name: &str, content: &[u8]) -> std::path::PathBuf {
269        let dir = std::env::temp_dir().join(format!("mmap_chunker_core_mc_{name}"));
270        let _ = std::fs::remove_dir_all(&dir);
271        std::fs::create_dir_all(&dir).unwrap();
272        let file_path = dir.join("data.txt");
273        std::fs::write(&file_path, content).unwrap();
274        file_path
275    }
276
277    fn cleanup(path: &std::path::Path) {
278        if let Some(parent) = path.parent() {
279            let _ = std::fs::remove_dir_all(parent);
280        }
281    }
282
283    #[test]
284    fn test_chunker_open_nonexistent() {
285        unsafe {
286            let err = MmapChunker::open("definitely_does_not_exist_12345.dat").unwrap_err();
287            assert!(
288                err.kind() == std::io::ErrorKind::NotFound
289                    || err.kind() == std::io::ErrorKind::Other
290            );
291        }
292    }
293
294    #[test]
295    fn test_chunker_open_empty_file() {
296        let path = temp_file("empty", b"");
297
298        unsafe {
299            let file = MmapChunker::open(&path).unwrap();
300            assert!(file.is_empty());
301            assert_eq!(file.len(), 0);
302            assert_eq!(file.chunk_count(), 0);
303            assert_eq!(file.as_bytes(), b"");
304        }
305
306        cleanup(&path);
307    }
308
309    #[test]
310    fn test_chunker_scan_delimited_basic() {
311        let path = temp_file("delimited", b"aaa\nbbb\nccc\nddd\n");
312
313        unsafe {
314            let mut file = MmapChunker::open(&path).unwrap();
315            let count = file.scan_delimited(4, b'\n');
316            assert_eq!(count, 2);
317            assert_eq!(file.chunk_count(), 2);
318
319            assert_eq!(file.get_chunk(0), Some(b"aaa\nbbb\n" as &[u8]));
320            assert_eq!(file.get_chunk(1), Some(b"ccc\nddd\n" as &[u8]));
321            assert_eq!(file.get_chunk(2), None);
322        }
323
324        cleanup(&path);
325    }
326
327    #[test]
328    fn test_chunker_get_chunk_before_scan() {
329        let path = temp_file("prescan", b"some data\n");
330
331        unsafe {
332            let file = MmapChunker::open(&path).unwrap();
333            assert_eq!(file.chunk_count(), 0);
334            assert_eq!(file.get_chunk(0), None);
335        }
336
337        cleanup(&path);
338    }
339
340    #[test]
341    fn test_chunker_scan_fixed() {
342        let path = temp_file("fixed", b"AAAABBBBCCCCDDDD");
343
344        unsafe {
345            let mut file = MmapChunker::open(&path).unwrap();
346            let count = file.scan_fixed(4);
347            assert_eq!(count, 4);
348            assert_eq!(file.chunk_count(), 4);
349
350            assert_eq!(file.get_chunk(0), Some(b"AAAA" as &[u8]));
351            assert_eq!(file.get_chunk(1), Some(b"BBBB" as &[u8]));
352            assert_eq!(file.get_chunk(2), Some(b"CCCC" as &[u8]));
353            assert_eq!(file.get_chunk(3), Some(b"DDDD" as &[u8]));
354            assert_eq!(file.get_chunk(4), None);
355        }
356
357        cleanup(&path);
358    }
359
360    #[test]
361    fn test_chunker_scan_fixed_short_last() {
362        let path = temp_file("fixed_short", b"XXXXXXXXX");
363
364        unsafe {
365            let mut file = MmapChunker::open(&path).unwrap();
366            let count = file.scan_fixed(4);
367            assert_eq!(count, 3);
368            assert_eq!(file.get_chunk(0).map(|c| c.len()), Some(4));
369            assert_eq!(file.get_chunk(1).map(|c| c.len()), Some(4));
370            assert_eq!(file.get_chunk(2).map(|c| c.len()), Some(1));
371        }
372
373        cleanup(&path);
374    }
375
376    #[test]
377    fn test_chunker_partition_records() {
378        let path = temp_file("partition", b"record1\nrecord2\nrecord3\nrecord4\n");
379
380        unsafe {
381            let mut file = MmapChunker::open(&path).unwrap();
382            let count = file.partition_records(2, b'\n');
383            assert!(count == 2);
384
385            let mut total = 0usize;
386            for i in 0..count {
387                let chunk = file.get_chunk(i).unwrap();
388                total += chunk.len();
389                assert!(!chunk.is_empty());
390            }
391            assert_eq!(total, file.len());
392        }
393
394        cleanup(&path);
395    }
396
397    #[test]
398    fn test_chunker_as_bytes() {
399        let path = temp_file("as_bytes", b"hello world!");
400
401        unsafe {
402            let file = MmapChunker::open(&path).unwrap();
403            assert_eq!(file.as_bytes(), b"hello world!");
404            assert_eq!(file.len(), 12);
405            assert!(!file.is_empty());
406        }
407
408        cleanup(&path);
409    }
410
411    #[test]
412    fn test_chunker_mode_switching() {
413        let path = temp_file("mode_switch", b"aaa\nbbb\nccc\nddd\n");
414
415        unsafe {
416            let mut file = MmapChunker::open(&path).unwrap();
417
418            let dc = file.scan_delimited(4, b'\n');
419            assert!(dc > 0);
420
421            let fc = file.scan_fixed(4);
422            assert!(fc > 0);
423            assert_eq!(file.chunk_count(), fc);
424
425            let dc2 = file.scan_delimited(4, b'\n');
426            assert_eq!(dc2, dc);
427
428            let pc = file.partition_records(2, b'\n');
429            assert_eq!(pc, 2);
430            assert_eq!(file.chunk_count(), 2);
431        }
432
433        cleanup(&path);
434    }
435
436    #[test]
437    fn test_chunker_large_file() {
438        let path = temp_file("large", &vec![b'x'; 100_000]);
439
440        unsafe {
441            let mut file = MmapChunker::open(&path).unwrap();
442            assert_eq!(file.len(), 100_000);
443
444            let count = file.scan_fixed(4096);
445            assert!(count > 0);
446
447            let mut total = 0usize;
448            for i in 0..count {
449                let chunk = file.get_chunk(i).unwrap();
450                total += chunk.len();
451            }
452            assert_eq!(total, 100_000);
453        }
454
455        cleanup(&path);
456    }
457}