Skip to main content

readcon_core/
compression.rs

1//=============================================================================
2// Transparent compression support
3//=============================================================================
4
5use std::io::{self, Read, Write};
6use std::path::Path;
7
8/// zlib/gzip level used for CON payloads.
9///
10/// This is [`flate2::Compression::default`] (level 6). The historical
11/// campaign median of 0.120 s was this level applied to a frame whose
12/// coordinates had already been cleared; the first sample of that run,
13/// still holding the lattice, was 0.371 s.
14pub const GZIP_LEVEL: u32 = 6;
15
16/// zstd level used for CON payloads. Level 3 is the zstd CLI default.
17pub const ZSTD_LEVEL: i32 = 3;
18
19/// Detected compression format based on magic bytes.
20///
21/// `Zstd` is only constructed when the `zstd` Cargo feature is enabled.
22/// Builds without the feature treat `.zst` files as opaque bytes and
23/// return an error from [`read_file_contents`] indicating the feature
24/// is required.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum Compression {
27    None,
28    Gzip,
29    /// zstd data: a frame (`28 b5 2f fd`) or a skippable frame
30    /// (`5? 2a 4d 18`, RFC 8878). Build with `--features zstd`.
31    Zstd,
32}
33
34/// A Zstandard frame or a skippable frame, in the byte order on disk.
35///
36/// RFC 8878 gives a frame the magic `28 b5 2f fd` and a skippable frame
37/// one of `50 2a 4d 18` through `5f 2a 4d 18`. A file from `pzstd` can
38/// start with the skippable frame.
39fn is_zstd_magic(bytes: &[u8]) -> bool {
40    bytes.len() >= 4
41        && (bytes[..4] == [0x28, 0xb5, 0x2f, 0xfd]
42            || ((0x50..=0x5f).contains(&bytes[0])
43                && bytes[1] == 0x2a
44                && bytes[2] == 0x4d
45                && bytes[3] == 0x18))
46}
47
48/// Detect compression format from the first bytes of a file.
49///
50/// - `1f 8b` = gzip
51/// - `28 b5 2f fd` = zstd frame
52/// - `5? 2a 4d 18` = zstd skippable frame
53/// - Otherwise = uncompressed
54pub fn detect_compression(bytes: &[u8]) -> Compression {
55    if bytes.len() >= 2 && bytes[0] == 0x1f && bytes[1] == 0x8b {
56        return Compression::Gzip;
57    }
58    if is_zstd_magic(bytes) {
59        return Compression::Zstd;
60    }
61    Compression::None
62}
63
64/// Detect compression format from a file extension.
65///
66/// Returns `Compression::Gzip` for `.gz`, `Compression::Zstd` for
67/// `.zst`, `Compression::None` otherwise.
68pub fn detect_compression_from_extension(path: &Path) -> Compression {
69    match path.extension().and_then(|e| e.to_str()) {
70        Some("gz") => Compression::Gzip,
71        Some("zst") => Compression::Zstd,
72        _ => Compression::None,
73    }
74}
75
76/// Size in bytes from which [`read_file_contents`] maps an uncompressed file.
77/// A smaller one is read into a `String`.
78pub const MMAP_THRESHOLD: u64 = 64 * 1024;
79
80/// Reads file contents, decompressing if needed.
81///
82/// Detection strategy:
83/// 1. Read the first 4 bytes to check magic bytes.
84/// 2. If gzip or zstd: decompress entire file to a String.
85/// 3. If uncompressed and smaller than [`MMAP_THRESHOLD`]: read the rest into a String.
86/// 4. If uncompressed and at least [`MMAP_THRESHOLD`]: memory-mapped I/O.
87///
88/// The path is opened once and every byte comes from that handle, so a
89/// pipe works and a file replaced during the call is not mixed in.
90pub fn read_file_contents(path: &Path) -> Result<FileContents, Box<dyn std::error::Error>> {
91    read_file_contents_detected(path).map(|(contents, _)| contents)
92}
93
94/// [`read_file_contents`], also returning the compression it detected.
95pub fn read_file_contents_detected(
96    path: &Path,
97) -> Result<(FileContents, Compression), Box<dyn std::error::Error>> {
98    read_contents(path, true)
99}
100
101/// [`read_file_contents`] as one owned `String`. An uncompressed file is
102/// read into it at any size instead of being mapped, so the text is held
103/// once and later changes to the file do not reach it.
104pub fn read_file_text(path: &Path) -> Result<String, Box<dyn std::error::Error>> {
105    match read_contents(path, false)?.0 {
106        FileContents::Owned(text) => Ok(text),
107        FileContents::Mapped(_) => unreachable!("read_contents maps only when asked to"),
108    }
109}
110
111/// Bytes of one file, mapped when a plain file is large enough.
112///
113/// Compressed input is decoded into a buffer. A plain file under 64 KiB,
114/// or anything that is not a regular file, is read into a buffer. A larger
115/// plain file stays mapped for the parse and is not copied. The bytes are
116/// not checked as UTF-8; the parser reports a bad byte when it reaches it.
117pub enum ByteSource {
118    Buf(Vec<u8>),
119    Map(memmap2::Mmap),
120}
121
122impl ByteSource {
123    pub fn as_slice(&self) -> &[u8] {
124        match self {
125            ByteSource::Buf(bytes) => bytes,
126            ByteSource::Map(map) => map,
127        }
128    }
129}
130
131/// Load `path` for a parse. See [`ByteSource`].
132pub fn read_bytes_owned(path: &Path) -> Result<ByteSource, Box<dyn std::error::Error>> {
133    let mut file = std::fs::File::open(path)?;
134    let metadata = file.metadata()?;
135    let mut magic = [0u8; 4];
136    let n = (&mut file).take(4).read(&mut magic)?;
137    match detect_compression(&magic[..n]) {
138        Compression::None if metadata.is_file() && metadata.len() >= MMAP_THRESHOLD => {
139            let map = unsafe { memmap2::Mmap::map(&file)? };
140            Ok(ByteSource::Map(map))
141        }
142        Compression::None => {
143            let mut bytes = Vec::with_capacity(metadata.len().saturating_add(4) as usize);
144            bytes.extend_from_slice(&magic[..n]);
145            file.read_to_end(&mut bytes)?;
146            Ok(ByteSource::Buf(bytes))
147        }
148        Compression::Gzip | Compression::Zstd => {
149            drop(file);
150            Ok(ByteSource::Buf(read_file_bytes(path)?))
151        }
152    }
153}
154
155/// File bytes in one read when the file is not compressed.
156///
157/// Gzip and zstd still go through the decoder. Plain text skips the
158/// four-byte probe and the UTF-8 `String` allocation.
159pub fn read_file_bytes(path: &Path) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
160    let bytes = std::fs::read(path)?;
161    let magic = &bytes[..bytes.len().min(4)];
162    match detect_compression(magic) {
163        Compression::None => Ok(bytes),
164        Compression::Gzip => gunzip_all(&bytes).map_err(Into::into),
165        Compression::Zstd => {
166            #[cfg(feature = "zstd")]
167            {
168                zstd_decode_all(&bytes).map_err(Into::into)
169            }
170            #[cfg(not(feature = "zstd"))]
171            {
172                let _ = bytes;
173                Err(zstd_unsupported().into())
174            }
175        }
176    }
177}
178
179fn read_contents(
180    path: &Path,
181    map_large: bool,
182) -> Result<(FileContents, Compression), Box<dyn std::error::Error>> {
183    let file = std::fs::File::open(path)?;
184    let metadata = file.metadata()?;
185
186    // gzip needs 2 magic bytes and zstd needs 4. They stay in front of the
187    // rest of the stream, since a pipe cannot seek back.
188    let mut magic = Vec::with_capacity(4);
189    (&file).take(4).read_to_end(&mut magic)?;
190    let compression = detect_compression(&magic);
191    let mut stream = magic.as_slice().chain(&file);
192
193    let contents = match compression {
194        Compression::Gzip => {
195            // Every gzip member, in order (CON-compress.1, D16). A file
196            // this crate wrote stores member starts in the first extra
197            // field and those slices are inflated together. Any other
198            // gzip stream, including one member and a foreign multi-member
199            // file, stays on the serial decoder.
200            let mut bytes = Vec::new();
201            stream.read_to_end(&mut bytes)?;
202            let plain = gunzip_all(&bytes)?;
203            let contents = String::from_utf8(plain).map_err(|err| {
204                io::Error::new(io::ErrorKind::InvalidData, err.utf8_error().to_string())
205            })?;
206            FileContents::Owned(contents)
207        }
208        Compression::Zstd => {
209            #[cfg(feature = "zstd")]
210            {
211                let mut bytes = Vec::new();
212                stream.read_to_end(&mut bytes)?;
213                let plain = zstd_decode_all(&bytes)?;
214                let contents = String::from_utf8(plain).map_err(|err| {
215                    io::Error::new(io::ErrorKind::InvalidData, err.utf8_error().to_string())
216                })?;
217                FileContents::Owned(contents)
218            }
219            #[cfg(not(feature = "zstd"))]
220            {
221                return Err(zstd_unsupported().into());
222            }
223        }
224        Compression::None if metadata.len() < MMAP_THRESHOLD => {
225            let mut bytes = Vec::with_capacity(metadata.len() as usize);
226            stream.read_to_end(&mut bytes)?;
227            let contents = String::from_utf8(bytes).map_err(|_| {
228                io::Error::new(
229                    io::ErrorKind::InvalidData,
230                    "stream did not contain valid UTF-8",
231                )
232            })?;
233            FileContents::Owned(contents)
234        }
235        Compression::None if map_large => {
236            let mmap = unsafe { memmap2::Mmap::map(&file)? };
237            FileContents::Mapped(mmap)
238        }
239        Compression::None => {
240            let mut bytes = Vec::with_capacity(metadata.len() as usize);
241            stream.read_to_end(&mut bytes)?;
242            // The message `as_str` gives for the same bytes when mapped.
243            FileContents::Owned(String::from_utf8(bytes).map_err(|e| e.utf8_error())?)
244        }
245    };
246    Ok((contents, compression))
247}
248
249#[cfg(not(feature = "zstd"))]
250fn zstd_unsupported() -> io::Error {
251    io::Error::new(
252        io::ErrorKind::Unsupported,
253        "zstd-compressed input detected; rebuild readcon-core with --features zstd",
254    )
255}
256
257/// The bytes of `path` as one stream, decompressed while it is read.
258///
259/// Compression is detected by magic bytes, as in [`read_file_contents`],
260/// and every gzip member is read in order.
261/// The only memory used is the decoder's own buffers, so it does not grow
262/// with the file.
263pub fn open_decoded(path: &Path) -> Result<Box<dyn Read + Send>, Box<dyn std::error::Error>> {
264    let mut file = std::fs::File::open(path)?;
265    let mut magic = Vec::with_capacity(4);
266    (&mut file).take(4).read_to_end(&mut magic)?;
267    let compression = detect_compression(&magic);
268    let stream = io::Cursor::new(magic).chain(file);
269    Ok(match compression {
270        Compression::Gzip => Box::new(flate2::read::MultiGzDecoder::new(stream)),
271        #[cfg(feature = "zstd")]
272        Compression::Zstd => Box::new(zstd::stream::read::Decoder::new(stream)?),
273        #[cfg(not(feature = "zstd"))]
274        Compression::Zstd => return Err(zstd_unsupported().into()),
275        Compression::None => Box::new(stream),
276    })
277}
278
279/// Holds file contents either as an owned String or a memory-mapped region.
280pub enum FileContents {
281    Owned(String),
282    Mapped(memmap2::Mmap),
283}
284
285impl FileContents {
286    pub fn as_str(&self) -> Result<&str, std::str::Utf8Error> {
287        match self {
288            FileContents::Owned(s) => Ok(s.as_str()),
289            FileContents::Mapped(m) => std::str::from_utf8(m),
290        }
291    }
292
293    pub fn as_bytes(&self) -> &[u8] {
294        match self {
295            FileContents::Owned(s) => s.as_bytes(),
296            FileContents::Mapped(m) => m,
297        }
298    }
299
300    /// Like [`Self::as_str`], but validates mapped files in newline-aligned
301    /// slices (in parallel with the `parallel` feature).
302    pub fn as_str_sliced(&self) -> Result<&str, std::str::Utf8Error> {
303        match self {
304            FileContents::Owned(s) => Ok(s.as_str()),
305            FileContents::Mapped(m) => crate::frame_index::validate_utf8_sliced(m),
306        }
307    }
308}
309
310/// Mapped bytes between two releases of the pages already read.
311const RELEASE_STEP: usize = 32 << 20;
312
313/// The text of one file, read from the front.
314///
315/// With `map`, an uncompressed file of 64 KiB or more stays mapped, and
316/// [`Self::release_before`] advises `MADV_DONTNEED` for the pages before the
317/// read position. On Linux that drops the file pages, so resident memory
318/// stays near [`RELEASE_STEP`] instead of the file size. macOS gets the same
319/// call, but its `MADV_DONTNEED` does not discard pages. Outside Unix the
320/// release is a no-op. On those systems the whole file can become resident.
321/// A mapped file is not a snapshot and must not change while it is read; a
322/// truncation can end the process with `SIGBUS`.
323pub struct ForwardText {
324    contents: FileContents,
325    released: usize,
326}
327
328impl ForwardText {
329    /// The text of `path`. Without `map` it is [`read_file_text`].
330    pub fn open(path: &Path, map: bool) -> Result<Self, Box<dyn std::error::Error>> {
331        let contents = if map {
332            read_file_contents(path)?
333        } else {
334            FileContents::Owned(read_file_text(path)?)
335        };
336        // One sequential check. Discarding the pages here would fault
337        // them back in during the parse, and a parallel check would
338        // start a worker pool for a single slice.
339        if let FileContents::Mapped(m) = &contents {
340            std::str::from_utf8(m)?;
341        }
342        Ok(Self {
343            contents,
344            released: 0,
345        })
346    }
347
348    pub fn as_str(&self) -> &str {
349        match &self.contents {
350            FileContents::Owned(s) => s,
351            // SAFETY: `open` checked every byte, in slices that end after a
352            // newline. A mapped file must not change while mapped, as for
353            // every mapped reader in this crate.
354            FileContents::Mapped(m) => unsafe { std::str::from_utf8_unchecked(m) },
355        }
356    }
357
358    /// Release the mapped pages before byte `pos` once [`RELEASE_STEP`]
359    /// bytes have been read since the last release. A later read of them
360    /// maps them in again. A `pos` past the end counts as the end.
361    pub fn release_before(&mut self, pos: usize) -> io::Result<()> {
362        if let FileContents::Mapped(m) = &self.contents
363            && pos >= self.released + RELEASE_STEP
364        {
365            let pos = pos.min(m.len());
366            release(m, self.released, pos)?;
367            self.released = pos;
368        }
369        Ok(())
370    }
371}
372
373/// Drop resident pages of `map` on `workers` threads, then the caller
374/// unmaps. Each thread advises a disjoint byte range. The page size is
375/// `sysconf(SC_PAGESIZE)`, not a host-tuned cutoff.
376pub fn advise_dontneed_parallel(map: &memmap2::Mmap, workers: usize) {
377    let len = map.len();
378    if len == 0 {
379        return;
380    }
381    let workers = workers.max(1);
382    if workers == 1 {
383        let _ = release(map, 0, len);
384        return;
385    }
386    let step = len.div_ceil(workers);
387    let advise_one = |worker: usize| {
388        let start = worker.saturating_mul(step);
389        if start >= len {
390            return;
391        }
392        let end = (start + step).min(len);
393        let _ = release(map, start, end);
394    };
395    // Inside the read's pool this stays on those workers. Outside it,
396    // a private scope avoids the global pool.
397    #[cfg(feature = "parallel")]
398    if rayon::current_thread_index().is_some() {
399        use rayon::prelude::*;
400        (0..workers).into_par_iter().for_each(advise_one);
401        return;
402    }
403    std::thread::scope(|scope| {
404        for worker in 0..workers {
405            scope.spawn(move || advise_one(worker));
406        }
407    });
408}
409
410/// Advise `MADV_DONTNEED` for the whole 64 KiB blocks of the map that
411/// `start..end` reaches, cut at the map end.
412fn release(map: &memmap2::Mmap, start: usize, end: usize) -> io::Result<()> {
413    // A multiple of every Linux page size. The map starts on a page, so the
414    // range stays inside it, and the page that contains `end` stays mapped.
415    const BLOCK: usize = 64 * 1024;
416    let start = start & !(BLOCK - 1);
417    let end = end.min(map.len()) & !(BLOCK - 1);
418    if end <= start {
419        return Ok(());
420    }
421    // SAFETY: `start..end` lies inside the map, which is shared and read
422    // only, so dropped pages read back from the file. Miri has no
423    // `MADV_DONTNEED`; skipping the hint leaves the bytes in place.
424    #[cfg(all(unix, not(miri)))]
425    unsafe {
426        map.unchecked_advise_range(memmap2::UncheckedAdvice::DontNeed, start, end - start)?;
427    }
428    #[cfg(not(all(unix, not(miri))))]
429    let _ = (map, start, end);
430    Ok(())
431}
432
433/// Creates a gzip-compressed writer wrapping a file at the given path.
434pub fn gzip_writer(path: &Path) -> io::Result<flate2::write::GzEncoder<std::fs::File>> {
435    let file = std::fs::File::create(path)?;
436    Ok(flate2::write::GzEncoder::new(
437        file,
438        flate2::Compression::new(GZIP_LEVEL),
439    ))
440}
441
442/// How many independent compression members a payload of `nbytes` uses.
443///
444/// Below 1 MiB the file stays one member, so small fixtures keep a single
445/// gzip stream. Larger payloads use about 2 MiB of plain text per member,
446/// and at least four members per core, so a read can keep every worker
447/// busy instead of waiting on one fat member. Capped at 64.
448pub fn compression_workers(nbytes: usize) -> usize {
449    const MIN_CHUNKED: usize = 1 << 20;
450    const PLAIN_PER_MEMBER: usize = 2 << 20;
451    if nbytes < MIN_CHUNKED {
452        return 1;
453    }
454    let threads = std::thread::available_parallelism()
455        .map(|n| n.get())
456        .unwrap_or(1);
457    let by_size = (nbytes / PLAIN_PER_MEMBER).max(1);
458    by_size.max(threads.saturating_mul(4)).clamp(1, 64)
459}
460
461/// Deflate workers for a write that did not ask for a thread count.
462///
463/// The member count stays wide so a read can split the file. Running one
464/// thread per member oversubscribes an 8-core host once that count hits
465/// 32. The cap is `available_parallelism`.
466pub fn default_deflate_threads(members: usize) -> usize {
467    let cores = std::thread::available_parallelism()
468        .map(|count| count.get())
469        .unwrap_or(1)
470        .max(1);
471    members.max(1).min(cores)
472}
473
474fn split_chunks(data: &[u8], workers: usize) -> Vec<&[u8]> {
475    if data.is_empty() || workers <= 1 {
476        return vec![data];
477    }
478    let workers = workers.max(1);
479    let mut cuts = Vec::with_capacity(workers + 1);
480    cuts.push(0usize);
481    for index in 1..workers {
482        let mut at = index * data.len() / workers;
483        if at > 0 && at < data.len() && data[at - 1] != b'\n' {
484            let forward = memchr::memchr(b'\n', &data[at..]).map(|rel| at + rel + 1);
485            let back = memchr::memrchr(b'\n', &data[..at]).map(|rel| rel + 1);
486            let previous = *cuts.last().unwrap_or(&0);
487            at = match (forward, back) {
488                (Some(next), Some(prev))
489                    if next < data.len() && (prev <= previous || next - at <= at - prev) =>
490                {
491                    next
492                }
493                (Some(next), _) if next < data.len() => next,
494                (_, Some(prev)) if prev > previous => prev,
495                _ => at,
496            };
497        }
498        if cuts.last().copied() != Some(at) && at < data.len() {
499            cuts.push(at);
500        }
501    }
502    if cuts.last().copied() != Some(data.len()) {
503        cuts.push(data.len());
504    }
505    cuts.windows(2)
506        .filter(|span| span[1] > span[0])
507        .map(|span| &data[span[0]..span[1]])
508        .collect()
509}
510
511/// `(first coordinate row, row count)` for each plain slice.
512///
513/// `None` when the text is not one coordinate block after a header, so
514/// the reader keeps the full-frame parser. Old indexes omit this table
515/// and still decode.
516fn coord_member_rows(data: &[u8], parts: &[&[u8]]) -> Option<Vec<(u32, u32)>> {
517    if parts.len() < 2 || data.is_empty() {
518        return None;
519    }
520    let mut line_at = Vec::new();
521    let mut pos = 0usize;
522    while let Some(rel) = memchr::memchr(b'\n', &data[pos..]) {
523        line_at.push(pos);
524        pos += rel + 1;
525        if line_at.len() > 64 && line_at.len() > 9 {
526            break;
527        }
528    }
529    if line_at.len() < 9 {
530        return None;
531    }
532    let line = |index: usize| -> Option<&[u8]> {
533        let start = *line_at.get(index)?;
534        let end = line_at.get(index + 1).copied().unwrap_or(data.len());
535        let end = end.saturating_sub(usize::from(end > start && data[end - 1] == b'\n'));
536        Some(&data[start..end])
537    };
538    let types_text = std::str::from_utf8(line(6)?).ok()?;
539    let natm_types = crate::parser::parse_count_line(types_text, 1).ok()?[0];
540    // Species headers sit between coordinate groups. A member can be one
541    // contiguous coordinate run only when the frame has a single species;
542    // anything else keeps the full-frame parser.
543    if natm_types != 1 {
544        return None;
545    }
546    let counts_text = std::str::from_utf8(line(7)?).ok()?;
547    let counts = crate::parser::parse_count_line(counts_text, natm_types).ok()?;
548    let natoms: usize = counts.iter().sum();
549    if natoms == 0 {
550        return None;
551    }
552    let header_lines = 9 + natm_types * 2;
553    // A per-newline callback over a million rows is slower than the
554    // deflate it is trying to index. Count bytes in each slice instead.
555    let mut counts = Vec::with_capacity(parts.len());
556    for part in parts {
557        counts.push(count_newlines(part) as u32);
558    }
559    if counts[0] < header_lines as u32 {
560        return None;
561    }
562    counts[0] -= header_lines as u32;
563    let total: u32 = counts.iter().copied().sum();
564    if total as usize != natoms {
565        return None;
566    }
567    if let Some(last_nl) = memchr::memrchr(b'\n', data) {
568        if data[last_nl + 1..]
569            .iter()
570            .any(|byte| !byte.is_ascii_whitespace())
571        {
572            return None;
573        }
574    }
575    let mut row = 0u32;
576    let mut table = Vec::with_capacity(parts.len());
577    for count in counts {
578        table.push((row, count));
579        row = row.saturating_add(count);
580    }
581    if row as usize != natoms {
582        return None;
583    }
584    Some(table)
585}
586
587fn count_newlines(data: &[u8]) -> usize {
588    #[cfg(target_arch = "x86_64")]
589    {
590        static AVX2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
591        if *AVX2.get_or_init(|| is_x86_feature_detected!("avx2")) {
592            return unsafe { count_newlines_avx2(data) };
593        }
594    }
595    data.iter().filter(|byte| **byte == b'\n').count()
596}
597
598#[cfg(target_arch = "x86_64")]
599#[target_feature(enable = "avx2")]
600unsafe fn count_newlines_avx2(data: &[u8]) -> usize {
601    use std::arch::x86_64::*;
602    unsafe {
603        let mut n = 0usize;
604        let mut i = 0usize;
605        let nl = _mm256_set1_epi8(b'\n' as i8);
606        while i + 128 <= data.len() {
607            let mut mask_n = 0usize;
608            for lane in 0..4 {
609                let v = _mm256_loadu_si256(data.as_ptr().add(i + lane * 32) as *const _);
610                mask_n += _mm256_movemask_epi8(_mm256_cmpeq_epi8(v, nl)).count_ones() as usize;
611            }
612            n += mask_n;
613            i += 128;
614        }
615        while i + 32 <= data.len() {
616            let v = _mm256_loadu_si256(data.as_ptr().add(i) as *const _);
617            n += _mm256_movemask_epi8(_mm256_cmpeq_epi8(v, nl)).count_ones() as usize;
618            i += 32;
619        }
620        while i < data.len() {
621            n += usize::from(*data.get_unchecked(i) == b'\n');
622            i += 1;
623        }
624        n
625    }
626}
627
628fn gzip_member(data: &[u8]) -> io::Result<Vec<u8>> {
629    let mut enc = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::new(GZIP_LEVEL));
630    enc.write_all(data)?;
631    enc.finish()
632}
633
634/// Writes `data` as gzip members at [`GZIP_LEVEL`].
635///
636/// One member when the payload is small. Otherwise each core compresses a
637/// slice. A wide file stores the member starts in the first member's extra
638/// field so the reader can inflate those slices together. A one-member file
639/// is an ordinary gzip stream.
640pub fn write_gzip_members(path: &Path, data: &[u8]) -> io::Result<usize> {
641    write_gzip_members_n(path, data, compression_workers(data.len()))
642}
643
644/// [`write_gzip_members`] with an explicit member count.
645///
646/// `members == 1` is one gzip stream and does not use another core.
647/// The member count is also the worker count.
648pub fn write_gzip_members_n(path: &Path, data: &[u8], members: usize) -> io::Result<usize> {
649    write_gzip_members_limited(path, data, members, members.max(1))
650}
651
652/// [`write_gzip_members_n`] with a separate worker count.
653///
654/// Eight members on two workers compress four slices per worker. The
655/// bytes match a one-worker write of the same slices.
656/// Format each atom slice and deflate it on that same worker.
657///
658/// The plain text is the binding header on the first slice plus coordinate
659/// lines. One member stays on the streaming encoder so its bytes do not
660/// change. `threads == 1` is the caller's problem: this function always
661/// uses `threads` workers.
662pub fn write_fused_gzip_frame(
663    path: &Path,
664    frame: &crate::types::ConFrame,
665    members: usize,
666    threads: usize,
667) -> io::Result<usize> {
668    write_fused_frame(path, frame, members, threads, true)
669}
670
671/// [`write_fused_gzip_frame`] for zstd frames.
672#[cfg(feature = "zstd")]
673pub fn write_fused_zstd_frame(
674    path: &Path,
675    frame: &crate::types::ConFrame,
676    members: usize,
677    threads: usize,
678) -> io::Result<usize> {
679    write_fused_frame(path, frame, members, threads, false)
680}
681
682fn write_fused_frame(
683    path: &Path,
684    frame: &crate::types::ConFrame,
685    members: usize,
686    threads: usize,
687    gzip: bool,
688) -> io::Result<usize> {
689    let n = frame.atom_data.len();
690    let members = members.max(2).min(n.max(2));
691    let mut ranges = Vec::with_capacity(members);
692    let mut rows = Vec::with_capacity(members);
693    let mut row = 0u32;
694    let mut start = 0usize;
695    for index in 0..members {
696        let end = if index + 1 == members {
697            n
698        } else {
699            (index + 1) * n / members
700        };
701        if end <= start {
702            continue;
703        }
704        let count = (end - start) as u32;
705        rows.push((row, count));
706        row += count;
707        ranges.push((start, end));
708        start = end;
709    }
710    let header = crate::writer::binding_header_bytes(frame)
711        .map_err(|err| io::Error::other(err.to_string()))?;
712    let slots: Vec<std::sync::Mutex<Option<io::Result<(Vec<u8>, usize)>>>> = (0..ranges.len())
713        .map(|_| std::sync::Mutex::new(None))
714        .collect();
715    let threads = threads.max(1).min(ranges.len());
716    let ranges_ref = &ranges;
717    std::thread::scope(|scope| {
718        for (begin, end) in worker_ranges(ranges.len(), threads) {
719            let slots = &slots;
720            let header = &header;
721            scope.spawn(move || {
722                let mut plain = Vec::new();
723                for index in begin..end {
724                    let (atom_from, atom_to) = ranges_ref[index];
725                    plain.clear();
726                    if atom_from == 0 {
727                        plain.extend_from_slice(&header);
728                    }
729                    plain.extend(crate::writer::format_coord_range(frame, atom_from, atom_to));
730                    let plain_len = plain.len();
731                    let encoded = if gzip {
732                        gzip_member(&plain)
733                    } else {
734                        #[cfg(feature = "zstd")]
735                        {
736                            zstd_compress_part(&plain)
737                        }
738                        #[cfg(not(feature = "zstd"))]
739                        {
740                            let _ = &plain;
741                            Err(zstd_unsupported())
742                        }
743                    };
744                    *slots[index].lock().unwrap() = Some(encoded.map(|bytes| (bytes, plain_len)));
745                }
746            });
747        }
748    });
749    let mut encoded = Vec::with_capacity(ranges.len());
750    let mut plain_lens = Vec::with_capacity(ranges.len());
751    for slot in slots {
752        match slot.into_inner().unwrap_or(None) {
753            Some(Ok((bytes, plain_len))) => {
754                encoded.push(bytes);
755                plain_lens.push(plain_len);
756            }
757            Some(Err(err)) => return Err(err),
758            None => return Err(io::Error::other("compress")),
759        }
760    }
761    if gzip {
762        if encoded.len() > 1 {
763            stamp_gzip_index(&mut encoded, &plain_lens, Some(&rows));
764        }
765        write_members(path, &encoded)?;
766    } else {
767        #[cfg(feature = "zstd")]
768        {
769            let mut file = std::fs::File::create(path)?;
770            if encoded.len() > 1 {
771                file.write_all(&zstd_index_frame(&encoded, &plain_lens, Some(&rows)))?;
772            }
773            for part in &encoded {
774                file.write_all(part)?;
775            }
776            file.flush()?;
777        }
778    }
779    Ok(encoded.len())
780}
781
782pub fn write_gzip_members_limited(
783    path: &Path,
784    data: &[u8],
785    members: usize,
786    threads: usize,
787) -> io::Result<usize> {
788    let trace = std::env::var_os("READCON_WRITE_TRACE").is_some();
789    let t0 = std::time::Instant::now();
790    let parts = split_chunks(data, members.max(1));
791    let plain: Vec<usize> = parts.iter().map(|part| part.len()).collect();
792    let t_split = t0.elapsed().as_secs_f64();
793    let t1 = std::time::Instant::now();
794    let mut encoded = encode_parts(&parts, threads, gzip_member)?;
795    let t_encode = t1.elapsed().as_secs_f64();
796    let t2 = std::time::Instant::now();
797    if encoded.len() > 1 {
798        let rows = coord_member_rows(data, &parts);
799        stamp_gzip_index(&mut encoded, &plain, rows.as_deref());
800    }
801    let t_index = t2.elapsed().as_secs_f64();
802    let t3 = std::time::Instant::now();
803    let count = encoded.len();
804    write_members(path, &encoded)?;
805    if trace {
806        eprintln!(
807            "gzip-write split={t_split:.4} encode={t_encode:.4} index={t_index:.4} io={:.4} members={count} threads={threads}",
808            t3.elapsed().as_secs_f64()
809        );
810    }
811    Ok(count)
812}
813
814/// `RC` extra-field payload: `RCZ1`, member count, compressed starts,
815/// then each member's plain length.
816fn row_table(data: &[u8], count: usize) -> Option<Vec<(u32, u32)>> {
817    let at = 6 + 16 * count;
818    if data.len() < at + 8 * count {
819        return None;
820    }
821    let mut rows = Vec::with_capacity(count);
822    for index in 0..count {
823        let off = at + 8 * index;
824        let base = u32::from_le_bytes(data[off..off + 4].try_into().ok()?);
825        let n = u32::from_le_bytes(data[off + 4..off + 8].try_into().ok()?);
826        rows.push((base, n));
827    }
828    Some(rows)
829}
830
831fn stamp_gzip_index(members: &mut [Vec<u8>], plain: &[usize], rows: Option<&[(u32, u32)]>) {
832    let count = members.len();
833    let row_bytes = usize::from(rows.is_some()) * 8 * count;
834    let extra_len = 4 + 2 + 16 * count + row_bytes;
835    let inserted = 2 + 4 + extra_len;
836    let mut starts = Vec::with_capacity(count);
837    let mut offset = inserted;
838    for member in members.iter() {
839        starts.push(offset as u64);
840        offset += member.len();
841    }
842    starts[0] = 0;
843    let mut payload = Vec::with_capacity(extra_len);
844    payload.extend_from_slice(b"RCZ1");
845    payload.extend_from_slice(&(count as u16).to_le_bytes());
846    for start in &starts {
847        payload.extend_from_slice(&start.to_le_bytes());
848    }
849    for len in plain {
850        payload.extend_from_slice(&(*len as u64).to_le_bytes());
851    }
852    if let Some(rows) = rows {
853        for (base, n) in rows {
854            payload.extend_from_slice(&base.to_le_bytes());
855            payload.extend_from_slice(&n.to_le_bytes());
856        }
857    }
858    let mut extra = Vec::with_capacity(inserted);
859    extra.extend_from_slice(&((4 + payload.len()) as u16).to_le_bytes());
860    extra.extend_from_slice(b"RC");
861    extra.extend_from_slice(&(payload.len() as u16).to_le_bytes());
862    extra.extend_from_slice(&payload);
863    let first = &mut members[0];
864    debug_assert_eq!(&first[0..2], b"\x1f\x8b");
865    first[3] |= 0x04;
866    first.splice(10..10, extra);
867}
868
869struct MemberPlan {
870    starts: Vec<usize>,
871    plain: Option<Vec<usize>>,
872    /// `(first coordinate row, row count)` per member, when the payload is
873    /// one coordinate block after a single header.
874    rows: Option<Vec<(u32, u32)>>,
875}
876
877fn gzip_member_starts(bytes: &[u8]) -> Option<MemberPlan> {
878    if bytes.len() < 16 || bytes[0] != 0x1f || bytes[1] != 0x8b || bytes[3] & 0x04 == 0 {
879        return None;
880    }
881    let xlen = u16::from_le_bytes([bytes[10], bytes[11]]) as usize;
882    let extra = bytes.get(12..12 + xlen)?;
883    let mut cursor = 0;
884    while cursor + 4 <= extra.len() {
885        let len = u16::from_le_bytes([extra[cursor + 2], extra[cursor + 3]]) as usize;
886        let data = extra.get(cursor + 4..cursor + 4 + len)?;
887        if &extra[cursor..cursor + 2] == b"RC" && data.starts_with(b"RCZ1") && data.len() >= 6 {
888            let count = u16::from_le_bytes([data[4], data[5]]) as usize;
889            if data.len() < 6 + 8 * count || count == 0 {
890                return None;
891            }
892            let mut starts = Vec::with_capacity(count);
893            for index in 0..count {
894                let at = 6 + 8 * index;
895                let start = u64::from_le_bytes(data[at..at + 8].try_into().ok()?) as usize;
896                if start > bytes.len() || (index > 0 && start <= starts[index - 1]) {
897                    return None;
898                }
899                starts.push(start);
900            }
901            let plain = if data.len() >= 6 + 16 * count {
902                let mut lens = Vec::with_capacity(count);
903                for index in 0..count {
904                    let at = 6 + 8 * count + 8 * index;
905                    lens.push(u64::from_le_bytes(data[at..at + 8].try_into().ok()?) as usize);
906                }
907                Some(lens)
908            } else {
909                None
910            };
911            let rows = row_table(data, count);
912            return Some(MemberPlan {
913                starts,
914                plain,
915                rows,
916            });
917        }
918        cursor += 4 + len;
919    }
920    None
921}
922
923/// `0` uses one worker per member. `1` stays on the calling thread.
924pub(crate) fn decode_thread_budget() -> usize {
925    DECODE_THREADS.with(|flag| flag.get())
926}
927
928std::thread_local! {
929    static DECODE_THREADS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
930}
931
932/// Run `body` with indexed members inflated on at most `threads` workers.
933///
934/// `threads < 2` is the calling thread. The previous budget is restored
935/// when `body` returns, including on unwind.
936pub fn with_decode_threads<T>(threads: usize, body: impl FnOnce() -> T) -> T {
937    let previous = DECODE_THREADS.with(|flag| flag.replace(threads));
938    struct Reset(usize);
939    impl Drop for Reset {
940        fn drop(&mut self) {
941            DECODE_THREADS.with(|flag| flag.set(self.0));
942        }
943    }
944    let _reset = Reset(previous);
945    body()
946}
947
948/// Run `body` with multi-member inflation forced onto one thread.
949pub fn with_serial_decode<T>(body: impl FnOnce() -> T) -> T {
950    with_decode_threads(1, body)
951}
952
953fn gunzip_member(bytes: &[u8]) -> io::Result<Vec<u8>> {
954    let mut decoder = flate2::read::GzDecoder::new(bytes);
955    let mut out = Vec::new();
956    decoder.read_to_end(&mut out)?;
957    Ok(out)
958}
959
960fn gunzip_all(bytes: &[u8]) -> io::Result<Vec<u8>> {
961    let mut out = Vec::new();
962    inflate_gzip(bytes, &mut out)?;
963    Ok(out)
964}
965
966fn worker_limit_from_budget(members: usize) -> usize {
967    let members = members.max(1);
968    match decode_thread_budget() {
969        0 => members,
970        budget => budget.min(members).max(1),
971    }
972}
973
974fn worker_ranges(jobs: usize, threads: usize) -> Vec<(usize, usize)> {
975    let threads = threads.max(1).min(jobs.max(1));
976    let mut ranges = Vec::with_capacity(threads);
977    let mut start = 0usize;
978    for worker in 0..threads {
979        let end = if worker + 1 == threads {
980            jobs
981        } else {
982            (worker + 1) * jobs / threads
983        };
984        if end > start {
985            ranges.push((start, end));
986            start = end;
987        }
988    }
989    ranges
990}
991
992fn par_ranges(jobs: usize, threads: usize, work: impl Fn(usize, usize) + Sync + Send) {
993    let ranges = worker_ranges(jobs, threads);
994    if ranges.len() < 2 {
995        for (begin, end) in ranges {
996            work(begin, end);
997        }
998        return;
999    }
1000    #[cfg(feature = "parallel")]
1001    {
1002        crate::parser::on_pool(ranges.len(), || {
1003            use rayon::prelude::*;
1004            ranges
1005                .par_iter()
1006                .for_each(|(begin, end)| work(*begin, *end));
1007        });
1008    }
1009    #[cfg(not(feature = "parallel"))]
1010    {
1011        std::thread::scope(|scope| {
1012            for (worker, (begin, end)) in ranges.into_iter().enumerate() {
1013                let work = &work;
1014                scope.spawn(move || {
1015                    crate::parser::pin_worker(worker);
1016                    work(begin, end);
1017                });
1018            }
1019        });
1020    }
1021}
1022
1023fn gzip_into(src: &[u8], dst: &mut [u8]) -> io::Result<()> {
1024    let mut decoder = flate2::read::GzDecoder::new(src);
1025    decoder.read_exact(dst)?;
1026    let mut extra = [0u8; 1];
1027    if decoder.read(&mut extra)? != 0 {
1028        return Err(io::Error::new(
1029            io::ErrorKind::InvalidData,
1030            "gzip member is longer than the index",
1031        ));
1032    }
1033    Ok(())
1034}
1035
1036fn inflate_gzip(bytes: &[u8], out: &mut Vec<u8>) -> io::Result<()> {
1037    if decode_thread_budget() != 1 {
1038        if let Some(plan) = gzip_member_starts(bytes) {
1039            if plan.starts.len() > 1 {
1040                let threads = worker_limit_from_budget(plan.starts.len());
1041                if let Some(plain) = plan.plain.as_deref() {
1042                    if plain.len() == plan.starts.len() {
1043                        return inflate_sized(bytes, &plan.starts, plain, threads, gzip_into, out);
1044                    }
1045                }
1046                let parts = inflate_unknown(bytes, &plan.starts, threads, gunzip_member)?;
1047                out.clear();
1048                for part in parts {
1049                    out.extend_from_slice(&part);
1050                }
1051                return Ok(());
1052            }
1053        }
1054    }
1055    let mut decoder = flate2::read::MultiGzDecoder::new(bytes);
1056    out.clear();
1057    decoder.read_to_end(out)?;
1058    Ok(())
1059}
1060
1061/// Largest plain length an index may claim for a file of `file_len` bytes.
1062///
1063/// The index is attacker-controlled. A flipped high bit must fail the
1064/// decode instead of reserving that many bytes.
1065fn plain_length_budget(file_len: usize) -> usize {
1066    file_len.saturating_mul(8192).saturating_add(1024 * 1024)
1067}
1068
1069pub(crate) fn trusted_plain_total(file_len: usize, plain: &[usize]) -> Option<usize> {
1070    let cap = plain_length_budget(file_len);
1071    let mut total = 0usize;
1072    for len in plain {
1073        if *len > cap {
1074            return None;
1075        }
1076        total = total.checked_add(*len).filter(|sum| *sum <= cap)?;
1077    }
1078    Some(total)
1079}
1080
1081/// Atom count implied by a row table, when every span fits in `plain_total`.
1082#[cfg(feature = "parallel")]
1083pub(crate) fn trusted_row_total(plain_total: usize, rows: &[(u32, u32)]) -> Option<usize> {
1084    let cap = plain_total / 8 + 16;
1085    let mut total = 0usize;
1086    let mut spans = Vec::with_capacity(rows.len());
1087    for (base, count) in rows {
1088        let base = *base as usize;
1089        let count = *count as usize;
1090        total = total.checked_add(count).filter(|sum| *sum <= cap)?;
1091        let end = base.checked_add(count)?;
1092        spans.push((base, end));
1093    }
1094    // The spans must cover every atom once. A gap leaves an uninitialized
1095    // atom, and dropping that atom reads a dangling pointer.
1096    spans.sort_unstable();
1097    let mut cursor = 0usize;
1098    for (base, end) in spans {
1099        if base != cursor || end > total {
1100            return None;
1101        }
1102        cursor = end;
1103    }
1104    if cursor != total {
1105        return None;
1106    }
1107    Some(total)
1108}
1109
1110fn inflate_sized(
1111    bytes: &[u8],
1112    starts: &[usize],
1113    plain: &[usize],
1114    threads: usize,
1115    decode: impl Fn(&[u8], &mut [u8]) -> io::Result<()> + Sync + Send,
1116    out: &mut Vec<u8>,
1117) -> io::Result<()> {
1118    let Some(total) = trusted_plain_total(bytes.len(), plain) else {
1119        return Err(io::Error::new(
1120            io::ErrorKind::InvalidData,
1121            "member plain length exceeds the compressed file",
1122        ));
1123    };
1124    // `reserve(n)` grows until `capacity >= len + n`, not `capacity >= old + n`.
1125    // A reused buffer whose length was cleared still has its old capacity, so
1126    // reserving only the shortfall against that capacity leaves `set_len` past
1127    // the end.
1128    out.clear();
1129    if out.capacity() < total {
1130        let extra = total - out.len();
1131        out.reserve(extra);
1132    }
1133    // `u8` has no destructor. `set_len(0)` on the error path drops nothing
1134    // unread. Every successful return overwrites `0..total`.
1135    unsafe { out.set_len(total) };
1136    let mut offsets = Vec::with_capacity(plain.len());
1137    let mut cursor = 0usize;
1138    for len in plain {
1139        offsets.push(cursor);
1140        cursor += *len;
1141    }
1142    let mut ends = Vec::with_capacity(starts.len());
1143    ends.extend_from_slice(&starts[1..]);
1144    ends.push(bytes.len());
1145    let error = std::sync::Mutex::new(None);
1146    let base = out.as_mut_ptr() as usize;
1147    par_ranges(starts.len(), threads, |begin, end| {
1148        for index in begin..end {
1149            if ends[index] < starts[index] || ends[index] > bytes.len() {
1150                *error.lock().unwrap() = Some(io::Error::new(
1151                    io::ErrorKind::InvalidData,
1152                    "member range is outside the file",
1153                ));
1154                return;
1155            }
1156            let src = &bytes[starts[index]..ends[index]];
1157            let dst = unsafe {
1158                std::slice::from_raw_parts_mut((base as *mut u8).add(offsets[index]), plain[index])
1159            };
1160            if let Err(err) = decode(src, dst) {
1161                *error.lock().unwrap() = Some(err);
1162                return;
1163            }
1164        }
1165    });
1166    if let Some(err) = error.into_inner().unwrap_or(None) {
1167        unsafe { out.set_len(0) };
1168        return Err(err);
1169    }
1170    Ok(())
1171}
1172
1173fn inflate_unknown(
1174    bytes: &[u8],
1175    starts: &[usize],
1176    threads: usize,
1177    decode: impl Fn(&[u8]) -> io::Result<Vec<u8>> + Sync + Send,
1178) -> io::Result<Vec<Vec<u8>>> {
1179    let mut ends = Vec::with_capacity(starts.len());
1180    ends.extend_from_slice(&starts[1..]);
1181    ends.push(bytes.len());
1182    let slots: Vec<std::sync::Mutex<Option<io::Result<Vec<u8>>>>> = (0..starts.len())
1183        .map(|_| std::sync::Mutex::new(None))
1184        .collect();
1185    par_ranges(starts.len(), threads, |begin, end| {
1186        for index in begin..end {
1187            let part = if ends[index] < starts[index] || ends[index] > bytes.len() {
1188                Err(io::Error::new(
1189                    io::ErrorKind::InvalidData,
1190                    "member range is outside the file",
1191                ))
1192            } else {
1193                decode(&bytes[starts[index]..ends[index]])
1194            };
1195            *slots[index].lock().unwrap() = Some(part);
1196        }
1197    });
1198    let mut out = Vec::with_capacity(starts.len());
1199    for slot in slots {
1200        match slot.into_inner().unwrap_or(None) {
1201            Some(Ok(part)) => out.push(part),
1202            Some(Err(err)) => return Err(err),
1203            None => return Err(io::Error::other("inflate")),
1204        }
1205    }
1206    Ok(out)
1207}
1208
1209fn encode_parts<F>(parts: &[&[u8]], threads: usize, encode: F) -> io::Result<Vec<Vec<u8>>>
1210where
1211    F: Fn(&[u8]) -> io::Result<Vec<u8>> + Sync,
1212{
1213    if parts.len() < 2 || threads <= 1 {
1214        let mut out = Vec::with_capacity(parts.len().max(1));
1215        if parts.is_empty() {
1216            out.push(encode(b"")?);
1217            return Ok(out);
1218        }
1219        for part in parts {
1220            out.push(encode(part)?);
1221        }
1222        return Ok(out);
1223    }
1224    let slots: Vec<std::sync::Mutex<Option<io::Result<Vec<u8>>>>> = (0..parts.len())
1225        .map(|_| std::sync::Mutex::new(None))
1226        .collect();
1227    // One OS thread per worker, as the original member writer did. The
1228    // pinned pool is the parser's; sharing it with deflate left the
1229    // eight-thread gzip write slower than four.
1230    let ranges = worker_ranges(parts.len(), threads);
1231    std::thread::scope(|scope| {
1232        for (_worker, (begin, end)) in ranges.into_iter().enumerate() {
1233            let slots = &slots;
1234            let parts = &parts;
1235            let encode = &encode;
1236            scope.spawn(move || {
1237                for index in begin..end {
1238                    *slots[index].lock().unwrap() = Some(encode(parts[index]));
1239                }
1240            });
1241        }
1242    });
1243    let mut out = Vec::with_capacity(parts.len());
1244    for slot in slots {
1245        match slot.into_inner().unwrap_or(None) {
1246            Some(Ok(part)) => out.push(part),
1247            Some(Err(err)) => return Err(err),
1248            None => return Err(io::Error::other("compress")),
1249        }
1250    }
1251    Ok(out)
1252}
1253
1254fn write_members(path: &Path, members: &[Vec<u8>]) -> io::Result<()> {
1255    let mut file = std::fs::File::create(path)?;
1256    for member in members {
1257        file.write_all(member)?;
1258    }
1259    file.flush()
1260}
1261
1262/// Creates a zstd-compressed writer wrapping a file at the given path.
1263///
1264/// Returns a finished writer ready for streaming line writes; the caller
1265/// is responsible for dropping the writer to flush the final frame.
1266/// Available only when the `zstd` Cargo feature is enabled.
1267#[cfg(feature = "zstd")]
1268pub fn zstd_writer(
1269    path: &Path,
1270) -> io::Result<zstd::stream::write::AutoFinishEncoder<'static, std::fs::File>> {
1271    let file = std::fs::File::create(path)?;
1272    // Level 3 is the zstd CLI default; balances ratio against speed.
1273    let encoder = zstd::stream::write::Encoder::new(file, ZSTD_LEVEL)?;
1274    Ok(encoder.auto_finish())
1275}
1276
1277/// Writes `data` as concatenated zstd frames at [`ZSTD_LEVEL`].
1278///
1279/// One frame when the payload is small. A wide file leads with a skippable
1280/// frame that lists the data-frame starts, which the serial decoder ignores.
1281#[cfg(feature = "zstd")]
1282pub fn write_zstd_members(path: &Path, data: &[u8]) -> io::Result<usize> {
1283    write_zstd_members_n(path, data, compression_workers(data.len()))
1284}
1285
1286/// [`write_zstd_members`] with an explicit frame count.
1287#[cfg(feature = "zstd")]
1288pub fn write_zstd_members_n(path: &Path, data: &[u8], members: usize) -> io::Result<usize> {
1289    write_zstd_members_limited(path, data, members, members.max(1))
1290}
1291
1292/// [`write_zstd_members_n`] with a separate worker count.
1293#[cfg(feature = "zstd")]
1294pub fn write_zstd_members_limited(
1295    path: &Path,
1296    data: &[u8],
1297    members: usize,
1298    threads: usize,
1299) -> io::Result<usize> {
1300    let trace = std::env::var_os("READCON_WRITE_TRACE").is_some();
1301    let t0 = std::time::Instant::now();
1302    let parts = split_chunks(data, members.max(1));
1303    let plain: Vec<usize> = parts.iter().map(|part| part.len()).collect();
1304    let t_split = t0.elapsed().as_secs_f64();
1305    let t1 = std::time::Instant::now();
1306    // One frame keeps the `bulk::compress` bytes. Several frames reuse a
1307    // context; the plaintext is what callers compare.
1308    let encoded = if parts.len() < 2 {
1309        encode_parts(&parts, 1, |part| {
1310            zstd::bulk::compress(part, ZSTD_LEVEL).map_err(|err| io::Error::other(err.to_string()))
1311        })?
1312    } else {
1313        encode_parts(&parts, threads, zstd_compress_part)?
1314    };
1315    let t_encode = t1.elapsed().as_secs_f64();
1316    let t2 = std::time::Instant::now();
1317    let count = encoded.len();
1318    let mut file = std::fs::File::create(path)?;
1319    if count > 1 {
1320        let rows = coord_member_rows(data, &parts);
1321        file.write_all(&zstd_index_frame(&encoded, &plain, rows.as_deref()))?;
1322    }
1323    let t_index = t2.elapsed().as_secs_f64();
1324    let t3 = std::time::Instant::now();
1325    for frame in &encoded {
1326        file.write_all(frame)?;
1327    }
1328    file.flush()?;
1329    if trace {
1330        eprintln!(
1331            "zstd-write split={t_split:.4} encode={t_encode:.4} index={t_index:.4} io={:.4} members={count} threads={threads}",
1332            t3.elapsed().as_secs_f64()
1333        );
1334    }
1335    Ok(count)
1336}
1337
1338#[cfg(feature = "zstd")]
1339fn zstd_compress_part(part: &[u8]) -> io::Result<Vec<u8>> {
1340    // One member stays on `bulk::compress`, which is the single-frame
1341    // byte image. Several members reuse a context on the worker.
1342    if part.is_empty() {
1343        return zstd::bulk::compress(part, ZSTD_LEVEL)
1344            .map_err(|err| io::Error::other(err.to_string()));
1345    }
1346    std::thread_local! {
1347        static COMPRESSOR: std::cell::RefCell<Option<zstd::bulk::Compressor<'static>>> =
1348            const { std::cell::RefCell::new(None) };
1349    }
1350    COMPRESSOR.with(|cell| {
1351        let mut slot = cell.borrow_mut();
1352        if slot.is_none() {
1353            *slot = Some(
1354                zstd::bulk::Compressor::new(ZSTD_LEVEL)
1355                    .map_err(|err| io::Error::other(err.to_string()))?,
1356            );
1357        }
1358        slot.as_mut()
1359            .expect("compressor")
1360            .compress(part)
1361            .map_err(|err| io::Error::other(err.to_string()))
1362    })
1363}
1364
1365#[cfg(feature = "zstd")]
1366fn zstd_index_frame(frames: &[Vec<u8>], plain: &[usize], rows: Option<&[(u32, u32)]>) -> Vec<u8> {
1367    let row_bytes = usize::from(rows.is_some()) * 8 * frames.len();
1368    let payload_len = 4 + 2 + 16 * frames.len() + row_bytes;
1369    let mut starts = Vec::with_capacity(frames.len());
1370    let mut offset = 8 + payload_len;
1371    for frame in frames {
1372        starts.push(offset as u64);
1373        offset += frame.len();
1374    }
1375    let mut payload = Vec::with_capacity(payload_len);
1376    payload.extend_from_slice(b"RCZ1");
1377    payload.extend_from_slice(&(frames.len() as u16).to_le_bytes());
1378    for start in starts {
1379        payload.extend_from_slice(&start.to_le_bytes());
1380    }
1381    for len in plain {
1382        payload.extend_from_slice(&(*len as u64).to_le_bytes());
1383    }
1384    if let Some(rows) = rows {
1385        for (base, n) in rows {
1386            payload.extend_from_slice(&base.to_le_bytes());
1387            payload.extend_from_slice(&n.to_le_bytes());
1388        }
1389    }
1390    let mut bytes = Vec::with_capacity(8 + payload.len());
1391    bytes.extend_from_slice(&0x184D_2A50u32.to_le_bytes());
1392    bytes.extend_from_slice(&(payload.len() as u32).to_le_bytes());
1393    bytes.extend_from_slice(&payload);
1394    bytes
1395}
1396
1397#[cfg(feature = "zstd")]
1398fn zstd_frame_starts(bytes: &[u8]) -> Option<MemberPlan> {
1399    if bytes.len() < 12 || bytes[0..4] != [0x50, 0x2a, 0x4d, 0x18] {
1400        return None;
1401    }
1402    let size = u32::from_le_bytes(bytes[4..8].try_into().ok()?) as usize;
1403    let payload = bytes.get(8..8 + size)?;
1404    if !payload.starts_with(b"RCZ1") || payload.len() < 6 {
1405        return None;
1406    }
1407    let count = u16::from_le_bytes([payload[4], payload[5]]) as usize;
1408    if payload.len() < 6 + 8 * count || count == 0 {
1409        return None;
1410    }
1411    let mut starts = Vec::with_capacity(count);
1412    for index in 0..count {
1413        let at = 6 + 8 * index;
1414        let start = u64::from_le_bytes(payload[at..at + 8].try_into().ok()?) as usize;
1415        if start > bytes.len() || (index > 0 && start <= starts[index - 1]) {
1416            return None;
1417        }
1418        starts.push(start);
1419    }
1420    let plain = if payload.len() >= 6 + 16 * count {
1421        let mut lens = Vec::with_capacity(count);
1422        for index in 0..count {
1423            let at = 6 + 8 * count + 8 * index;
1424            lens.push(u64::from_le_bytes(payload[at..at + 8].try_into().ok()?) as usize);
1425        }
1426        Some(lens)
1427    } else {
1428        None
1429    };
1430    let rows = row_table(payload, count);
1431    Some(MemberPlan {
1432        starts,
1433        plain,
1434        rows,
1435    })
1436}
1437
1438#[cfg(feature = "zstd")]
1439fn zstd_into(src: &[u8], dst: &mut [u8]) -> io::Result<()> {
1440    std::thread_local! {
1441        static DECODER: std::cell::RefCell<Option<zstd::bulk::Decompressor<'static>>> =
1442            const { std::cell::RefCell::new(None) };
1443    }
1444    DECODER.with(|cell| {
1445        let mut slot = cell.borrow_mut();
1446        if slot.is_none() {
1447            *slot = Some(
1448                zstd::bulk::Decompressor::new().map_err(|err| io::Error::other(err.to_string()))?,
1449            );
1450        }
1451        let n = slot
1452            .as_mut()
1453            .expect("decoder")
1454            .decompress_to_buffer(src, dst)
1455            .map_err(|err| io::Error::other(err.to_string()))?;
1456        if n != dst.len() {
1457            return Err(io::Error::new(
1458                io::ErrorKind::InvalidData,
1459                "zstd frame length does not match the index",
1460            ));
1461        }
1462        Ok(())
1463    })
1464}
1465
1466#[cfg(feature = "zstd")]
1467fn zstd_decode_all(bytes: &[u8]) -> io::Result<Vec<u8>> {
1468    let mut out = Vec::new();
1469    inflate_zstd(bytes, &mut out)?;
1470    Ok(out)
1471}
1472
1473#[cfg(feature = "zstd")]
1474fn inflate_zstd(bytes: &[u8], out: &mut Vec<u8>) -> io::Result<()> {
1475    if decode_thread_budget() != 1 {
1476        if let Some(plan) = zstd_frame_starts(bytes) {
1477            if plan.starts.len() > 1 {
1478                let threads = worker_limit_from_budget(plan.starts.len());
1479                if let Some(plain) = plan.plain.as_deref() {
1480                    if plain.len() == plan.starts.len() {
1481                        return inflate_sized(bytes, &plan.starts, plain, threads, zstd_into, out);
1482                    }
1483                }
1484                let parts = inflate_unknown(bytes, &plan.starts, threads, |part| {
1485                    zstd::stream::decode_all(part).map_err(|err| io::Error::other(err.to_string()))
1486                })?;
1487                out.clear();
1488                for part in parts {
1489                    out.extend_from_slice(&part);
1490                }
1491                return Ok(());
1492            }
1493        }
1494    }
1495    let mut decoder = zstd::stream::read::Decoder::new(bytes)?;
1496    out.clear();
1497    decoder.read_to_end(out)?;
1498    Ok(())
1499}
1500
1501/// Inflate a gzip or zstd buffer with `threads` workers and no CON parse.
1502pub fn codec_inflate(bytes: &[u8], threads: usize) -> io::Result<Vec<u8>> {
1503    with_decode_threads(threads.max(1), || {
1504        if bytes.len() >= 2 && bytes[0] == 0x1f && bytes[1] == 0x8b {
1505            gunzip_all(bytes)
1506        } else if is_zstd_magic(bytes) {
1507            #[cfg(feature = "zstd")]
1508            {
1509                zstd_decode_all(bytes)
1510            }
1511            #[cfg(not(feature = "zstd"))]
1512            {
1513                Err(zstd_unsupported())
1514            }
1515        } else {
1516            Ok(bytes.to_vec())
1517        }
1518    })
1519}
1520
1521/// Deflate `plain` into `members` slices on `threads` workers.
1522///
1523/// The returned bytes are the concatenated members, including the index
1524/// a reader needs. The timed work is the codec, not a CON parse.
1525pub fn codec_deflate(
1526    plain: &[u8],
1527    members: usize,
1528    threads: usize,
1529    gzip: bool,
1530) -> io::Result<Vec<u8>> {
1531    if gzip {
1532        let parts = split_chunks(plain, members.max(1));
1533        let lens: Vec<usize> = parts.iter().map(|part| part.len()).collect();
1534        let mut encoded = encode_parts(&parts, threads, gzip_member)?;
1535        if encoded.len() > 1 {
1536            let rows = coord_member_rows(plain, &parts);
1537            stamp_gzip_index(&mut encoded, &lens, rows.as_deref());
1538        }
1539        let mut out = Vec::new();
1540        for part in encoded {
1541            out.extend(part);
1542        }
1543        return Ok(out);
1544    }
1545    #[cfg(feature = "zstd")]
1546    {
1547        let parts = split_chunks(plain, members.max(1));
1548        let lens: Vec<usize> = parts.iter().map(|part| part.len()).collect();
1549        let encoded = encode_parts(&parts, threads, zstd_compress_part)?;
1550        let mut out = Vec::new();
1551        if encoded.len() > 1 {
1552            let rows = coord_member_rows(plain, &parts);
1553            out.extend(zstd_index_frame(&encoded, &lens, rows.as_deref()));
1554        }
1555        for part in encoded {
1556            out.extend(part);
1557        }
1558        return Ok(out);
1559    }
1560    #[cfg(not(feature = "zstd"))]
1561    {
1562        let _ = (plain, members, threads);
1563        Err(zstd_unsupported())
1564    }
1565}
1566
1567/// Indexed members whose plain lengths and coordinate-row spans are known.
1568pub(crate) struct IndexedMembers {
1569    pub starts: Vec<usize>,
1570    pub plain: Vec<usize>,
1571    pub rows: Vec<(u32, u32)>,
1572    pub gzip: bool,
1573}
1574
1575pub(crate) fn indexed_members(bytes: &[u8]) -> Option<IndexedMembers> {
1576    if bytes.len() >= 2 && bytes[0] == 0x1f && bytes[1] == 0x8b {
1577        let plan = gzip_member_starts(bytes)?;
1578        let plain = plan.plain?;
1579        let rows = plan.rows?;
1580        if plain.len() != plan.starts.len() || rows.len() != plan.starts.len() || plain.len() < 2 {
1581            return None;
1582        }
1583        return Some(IndexedMembers {
1584            starts: plan.starts,
1585            plain,
1586            rows,
1587            gzip: true,
1588        });
1589    }
1590    #[cfg(feature = "zstd")]
1591    if is_zstd_magic(bytes) {
1592        let plan = zstd_frame_starts(bytes)?;
1593        let plain = plan.plain?;
1594        let rows = plan.rows?;
1595        if plain.len() != plan.starts.len() || rows.len() != plan.starts.len() || plain.len() < 2 {
1596            return None;
1597        }
1598        return Some(IndexedMembers {
1599            starts: plan.starts,
1600            plain,
1601            rows,
1602            gzip: false,
1603        });
1604    }
1605    None
1606}
1607
1608pub(crate) fn inflate_member(gzip: bool, src: &[u8], dst: &mut [u8]) -> io::Result<()> {
1609    if gzip {
1610        gzip_into(src, dst)
1611    } else {
1612        #[cfg(feature = "zstd")]
1613        {
1614            zstd_into(src, dst)
1615        }
1616        #[cfg(not(feature = "zstd"))]
1617        {
1618            let _ = (src, dst);
1619            Err(zstd_unsupported())
1620        }
1621    }
1622}
1623
1624std::thread_local! {
1625    static PLAIN_BUF: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
1626}
1627
1628/// Inflate `file` into a reused buffer and hand that slice to `body`.
1629///
1630/// The buffer stays on this thread. The next call keeps its capacity.
1631pub(crate) fn with_inflated_plain<T>(
1632    file: &[u8],
1633    body: impl FnOnce(&[u8]) -> Result<T, Box<dyn std::error::Error>>,
1634) -> Result<T, Box<dyn std::error::Error>> {
1635    PLAIN_BUF.with(|cell| {
1636        let mut buf = cell.borrow_mut();
1637        if file.len() >= 2 && file[0] == 0x1f && file[1] == 0x8b {
1638            inflate_gzip(file, &mut buf)
1639                .map_err(|err| Box::new(err) as Box<dyn std::error::Error>)?;
1640        } else if is_zstd_magic(file) {
1641            #[cfg(feature = "zstd")]
1642            {
1643                inflate_zstd(file, &mut buf)
1644                    .map_err(|err| Box::new(err) as Box<dyn std::error::Error>)?;
1645            }
1646            #[cfg(not(feature = "zstd"))]
1647            {
1648                return Err(Box::new(zstd_unsupported()) as Box<dyn std::error::Error>);
1649            }
1650        } else {
1651            buf.clear();
1652            buf.extend_from_slice(file);
1653        }
1654        let result = body(buf.as_slice())?;
1655        buf.clear();
1656        Ok(result)
1657    })
1658}
1659
1660#[cfg(test)]
1661mod tests {
1662    use super::*;
1663    use std::io::Write;
1664
1665    fn text_of_at_least(len: usize) -> Vec<u8> {
1666        let frame = include_bytes!("../resources/test/tiny_cuh2.con");
1667        frame.repeat(len / frame.len() + 2)
1668    }
1669
1670    fn large_text() -> Vec<u8> {
1671        text_of_at_least(MMAP_THRESHOLD as usize)
1672    }
1673
1674    /// Independent members inflate back to the same bytes. One member stays
1675    /// a single frame; several members carry the index a parallel read uses.
1676    #[test]
1677    fn codec_deflate_round_trips_one_member_and_several() {
1678        let plain = b"aaaa\nbbbb\ncccc\ndddd\neeee\nffff\n".repeat(80);
1679        for members in [1usize, 4] {
1680            let gzip = codec_deflate(&plain, members, 2, true).unwrap();
1681            assert!(gzip.starts_with(&[0x1f, 0x8b]));
1682            assert_eq!(codec_inflate(&gzip, 2).unwrap(), plain);
1683            #[cfg(feature = "zstd")]
1684            {
1685                let zstd_bytes = codec_deflate(&plain, members, 2, false).unwrap();
1686                assert_eq!(codec_inflate(&zstd_bytes, 2).unwrap(), plain);
1687            }
1688        }
1689        let raw = b"not a compressed buffer";
1690        assert_eq!(codec_inflate(raw, 1).unwrap(), raw);
1691    }
1692
1693    fn mapped_text(path: &Path) -> Result<String, String> {
1694        let contents = read_file_contents(path).map_err(|e| e.to_string())?;
1695        contents
1696            .as_str()
1697            .map(str::to_owned)
1698            .map_err(|e| e.to_string())
1699    }
1700
1701    #[test]
1702    fn compressed_text_that_is_not_utf8_is_rejected() {
1703        let dir = tempfile::tempdir().unwrap();
1704        let raw = b"Cu \xff";
1705        let gz = dir.path().join("bad.con.gz");
1706        {
1707            let mut enc = flate2::write::GzEncoder::new(
1708                std::fs::File::create(&gz).unwrap(),
1709                flate2::Compression::default(),
1710            );
1711            enc.write_all(raw).unwrap();
1712            enc.finish().unwrap();
1713        }
1714        let err = read_file_text(&gz).unwrap_err().to_string();
1715        assert!(err.contains("invalid utf-8"), "{err}");
1716        #[cfg(feature = "zstd")]
1717        {
1718            let zst = dir.path().join("bad.con.zst");
1719            std::fs::write(&zst, zstd::stream::encode_all(&raw[..], 3).unwrap()).unwrap();
1720            let err = read_file_text(&zst).unwrap_err().to_string();
1721            assert!(err.contains("invalid utf-8"), "{err}");
1722        }
1723    }
1724
1725    /// File-backed maps are not available under Miri.
1726    #[test]
1727    #[cfg(not(miri))]
1728    fn dontneed_outside_a_rayon_pool_keeps_the_bytes() {
1729        let dir = tempfile::tempdir().unwrap();
1730        let path = dir.path().join("wide.con");
1731        let bytes = large_text();
1732        std::fs::write(&path, &bytes).unwrap();
1733        let file = std::fs::File::open(&path).unwrap();
1734        let map = unsafe { memmap2::Mmap::map(&file).unwrap() };
1735        advise_dontneed_parallel(&map, 4);
1736        assert_eq!(&map[..bytes.len()], bytes.as_slice());
1737    }
1738
1739    #[test]
1740    fn file_text_matches_the_mapped_read() {
1741        let dir = tempfile::tempdir().unwrap();
1742        let small = dir.path().join("small.con");
1743        std::fs::write(&small, include_bytes!("../resources/test/tiny_cuh2.con")).unwrap();
1744        let large = dir.path().join("large.con");
1745        std::fs::write(&large, large_text()).unwrap();
1746        assert!(matches!(
1747            read_file_contents(&large).unwrap(),
1748            FileContents::Mapped(_)
1749        ));
1750        let gz = dir.path().join("large.con.gz");
1751        let mut enc = gzip_writer(&gz).unwrap();
1752        enc.write_all(&large_text()).unwrap();
1753        enc.finish().unwrap();
1754        for path in [&small, &large, &gz] {
1755            let text = read_file_text(path).unwrap();
1756            assert_eq!(Ok(text), mapped_text(path), "{}", path.display());
1757        }
1758    }
1759
1760    #[test]
1761    fn file_text_reports_invalid_utf8_as_the_mapped_read() {
1762        let dir = tempfile::tempdir().unwrap();
1763        for (name, len) in [("small.con", 100), ("large.con", large_text().len())] {
1764            let mut bytes = large_text()[..len].to_vec();
1765            bytes[len - 40] = 0xff;
1766            let path = dir.path().join(name);
1767            std::fs::write(&path, &bytes).unwrap();
1768            let err = read_file_text(&path).unwrap_err().to_string();
1769            assert_eq!(Err(err), mapped_text(&path), "{name}");
1770        }
1771    }
1772
1773    #[test]
1774    fn forward_text_reads_back_released_pages() {
1775        let dir = tempfile::tempdir().unwrap();
1776        let bytes = text_of_at_least(RELEASE_STEP + (200 << 10));
1777        let path = dir.path().join("two-steps.con");
1778        std::fs::write(&path, &bytes).unwrap();
1779        let mut mapped = ForwardText::open(&path, true).unwrap();
1780        assert!(matches!(mapped.contents, FileContents::Mapped(_)));
1781        assert_eq!(mapped.as_str().as_bytes(), bytes.as_slice());
1782        for pos in (0..=bytes.len()).step_by(1 << 20) {
1783            mapped.release_before(pos).unwrap();
1784        }
1785        assert_eq!(mapped.released, RELEASE_STEP);
1786        assert_eq!(mapped.as_str().as_bytes(), bytes.as_slice());
1787        let owned = ForwardText::open(&path, false).unwrap();
1788        assert!(matches!(owned.contents, FileContents::Owned(_)));
1789        assert_eq!(owned.as_str(), mapped.as_str());
1790    }
1791
1792    /// The guard is private anonymous memory right after the file map, which
1793    /// `MADV_DONTNEED` would zero.
1794    ///
1795    /// The kernel puts a new map at the top of the highest free gap that fits
1796    /// it. The test frees the space below the guard after it fills every gap
1797    /// above, so the file map must end at the guard. A map made by another
1798    /// test thread in between would take that space, so the test body runs
1799    /// alone in a child process of this test binary.
1800    #[cfg(target_os = "linux")]
1801    #[test]
1802    fn release_before_past_the_end_leaves_the_next_mapping_alone() {
1803        const NAME: &str =
1804            "compression::tests::release_before_past_the_end_leaves_the_next_mapping_alone";
1805        const CHILD: &str = "READCON_RELEASE_TEST_CHILD";
1806        if std::env::var_os(CHILD).is_none() {
1807            let out = std::process::Command::new(std::env::current_exe().unwrap())
1808                .args([NAME, "--exact", "--test-threads=1"])
1809                .env(CHILD, "1")
1810                .output()
1811                .unwrap();
1812            let stdout = String::from_utf8_lossy(&out.stdout);
1813            assert!(
1814                out.status.success() && stdout.contains("test result: ok. 1 passed;"),
1815                "{stdout}{}",
1816                String::from_utf8_lossy(&out.stderr)
1817            );
1818            return;
1819        }
1820        const GUARD: usize = 1 << 20;
1821        let dir = tempfile::tempdir().unwrap();
1822        let bytes = text_of_at_least(3 * MMAP_THRESHOLD as usize);
1823        let path = dir.path().join("short.con");
1824        std::fs::write(&path, &bytes).unwrap();
1825        let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as usize;
1826        let map_len = bytes.len().next_multiple_of(page);
1827        let anonymous = |len: usize| {
1828            let map = unsafe {
1829                libc::mmap(
1830                    std::ptr::null_mut(),
1831                    len,
1832                    libc::PROT_READ | libc::PROT_WRITE,
1833                    libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
1834                    -1,
1835                    0,
1836                )
1837            };
1838            assert_ne!(map, libc::MAP_FAILED);
1839            map
1840        };
1841        let region = anonymous(map_len + GUARD);
1842        let guard =
1843            unsafe { std::slice::from_raw_parts_mut(region.cast::<u8>().add(map_len), GUARD) };
1844        guard.fill(0xab);
1845        // Each filler takes the highest gap that fits the file map. The first
1846        // one below the region shows that no such gap is left above it.
1847        let mut fillers = Vec::new();
1848        loop {
1849            let filler = anonymous(map_len);
1850            if filler < region {
1851                unsafe { libc::munmap(filler, map_len) };
1852                break;
1853            }
1854            fillers.push(filler);
1855        }
1856        unsafe { libc::munmap(region, map_len) };
1857        let mut text = ForwardText::open(&path, true).unwrap();
1858        assert!(
1859            std::ptr::eq(text.as_str().as_ptr(), region.cast::<u8>()),
1860            "the file map does not end at the guard"
1861        );
1862        for pos in [bytes.len() + RELEASE_STEP, 3 * RELEASE_STEP, usize::MAX] {
1863            text.release_before(pos).unwrap();
1864        }
1865        assert_eq!(text.released, bytes.len());
1866        assert!(guard.iter().all(|&b| b == 0xab));
1867        assert_eq!(text.as_str().as_bytes(), bytes.as_slice());
1868        unsafe { libc::munmap(guard.as_mut_ptr().cast(), GUARD) };
1869        for filler in fillers {
1870            unsafe { libc::munmap(filler, map_len) };
1871        }
1872    }
1873
1874    #[test]
1875    fn forward_text_reports_invalid_utf8_as_the_mapped_read() {
1876        let dir = tempfile::tempdir().unwrap();
1877        let mut bytes = text_of_at_least(RELEASE_STEP + (200 << 10));
1878        let at = bytes.len() - 100;
1879        bytes[at] = 0xc3;
1880        let path = dir.path().join("bad.con");
1881        std::fs::write(&path, &bytes).unwrap();
1882        for map in [true, false] {
1883            let err = ForwardText::open(&path, map).err().unwrap().to_string();
1884            assert_eq!(Err(err), mapped_text(&path), "map {map}");
1885        }
1886    }
1887
1888    #[test]
1889    fn skippable_frame_magic_is_zstd() {
1890        assert_eq!(
1891            detect_compression(&[0x50, 0x2a, 0x4d, 0x18]),
1892            Compression::Zstd
1893        );
1894        assert_eq!(
1895            detect_compression(&[0x5f, 0x2a, 0x4d, 0x18]),
1896            Compression::Zstd
1897        );
1898        assert_eq!(
1899            detect_compression(&[0x4f, 0x2a, 0x4d, 0x18]),
1900            Compression::None
1901        );
1902        assert_eq!(detect_compression(&[0x50, 0x2a, 0x4d]), Compression::None);
1903        assert_eq!(
1904            detect_compression(&[0x28, 0xb5, 0x2f, 0xfd]),
1905            Compression::Zstd
1906        );
1907    }
1908
1909    fn skippable_prefix(nibble: u8, payload: &[u8]) -> Vec<u8> {
1910        let mut bytes = Vec::new();
1911        bytes.extend_from_slice(&(0x184D_2A50u32 + u32::from(nibble)).to_le_bytes());
1912        bytes.extend_from_slice(&(payload.len() as u32).to_le_bytes());
1913        bytes.extend_from_slice(payload);
1914        bytes
1915    }
1916
1917    #[cfg(not(feature = "zstd"))]
1918    #[test]
1919    fn a_leading_skippable_frame_needs_the_zstd_feature() {
1920        let dir = tempfile::tempdir().unwrap();
1921        let path = dir.path().join("skip-first.zst");
1922        std::fs::write(&path, skippable_prefix(0, b"index")).unwrap();
1923        let err = read_file_text(&path).unwrap_err().to_string();
1924        assert!(err.contains("zstd-compressed input detected"), "{err}");
1925        match open_decoded(&path) {
1926            Err(err) => {
1927                assert!(
1928                    err.to_string().contains("zstd-compressed input detected"),
1929                    "{err}"
1930                );
1931            }
1932            Ok(_) => panic!("a skippable frame was treated as uncompressed"),
1933        }
1934    }
1935
1936    #[cfg(feature = "zstd")]
1937    #[test]
1938    fn a_leading_skippable_frame_decodes_with_the_frames_after_it() {
1939        let text = include_str!("../resources/test/tiny_cuh2.con");
1940        let frame = zstd::stream::encode_all(text.as_bytes(), 1).unwrap();
1941        let dir = tempfile::tempdir().unwrap();
1942        let first = dir.path().join("skip-first.con.zst");
1943        let mut bytes = skippable_prefix(0, b"index");
1944        bytes.extend_from_slice(&frame);
1945        std::fs::write(&first, &bytes).unwrap();
1946        assert_eq!(read_file_text(&first).unwrap(), text);
1947        let frames = crate::iterators::read_all_frames(&first).unwrap();
1948        assert_eq!(frames.len(), 1);
1949
1950        let mut between = frame.clone();
1951        between.extend(skippable_prefix(0, b"index"));
1952        between.extend_from_slice(&frame);
1953        let mid = dir.path().join("skip-between.con.zst");
1954        std::fs::write(&mid, &between).unwrap();
1955        assert_eq!(read_file_text(&mid).unwrap(), format!("{text}{text}"));
1956    }
1957
1958    fn gzip_member_count(bytes: &[u8]) -> usize {
1959        let mut rest = bytes;
1960        let mut count = 0;
1961        while !rest.is_empty() {
1962            let mut decoder = flate2::bufread::GzDecoder::new(rest);
1963            let mut out = Vec::new();
1964            std::io::Read::read_to_end(&mut decoder, &mut out).unwrap();
1965            let next = decoder.into_inner();
1966            assert!(
1967                next.len() < rest.len() || out.is_empty(),
1968                "gzip member did not consume input"
1969            );
1970            if next.len() == rest.len() {
1971                break;
1972            }
1973            rest = next;
1974            count += 1;
1975        }
1976        count
1977    }
1978
1979    #[test]
1980    fn indexed_gzip_members_are_cut_on_newlines() {
1981        let mut data = Vec::new();
1982        for i in 0..50_000 {
1983            writeln!(data, "Cu {i}.0 2.0 3.0 0 {i}").unwrap();
1984        }
1985        assert!(data.len() > (1 << 20));
1986        let dir = tempfile::tempdir().unwrap();
1987        let path = dir.path().join("wide.con.gz");
1988        let members = write_gzip_members_n(&path, &data, 8).unwrap();
1989        assert!(members > 1, "a wide payload is more than one gzip member");
1990        let bytes = std::fs::read(&path).unwrap();
1991        let starts = gzip_member_starts(&bytes).expect("index").starts;
1992        assert_eq!(starts.len(), members);
1993        let mut plain = Vec::new();
1994        for (index, start) in starts.iter().copied().enumerate() {
1995            let end = starts.get(index + 1).copied().unwrap_or(bytes.len());
1996            let part = gunzip_member(&bytes[start..end]).unwrap();
1997            assert!(
1998                part.ends_with(b"\n"),
1999                "member {index} ends mid-line ({} bytes)",
2000                part.len()
2001            );
2002            plain.extend_from_slice(&part);
2003        }
2004        assert_eq!(plain, data);
2005    }
2006
2007    #[test]
2008    fn a_short_indexed_plain_length_is_rejected() {
2009        let data = b"Cu 1.0 2.0 3.0 0 1\nH 4.0 5.0 6.0 0 2\n".repeat(8);
2010        let dir = tempfile::tempdir().unwrap();
2011        let path = dir.path().join("short-index.con.gz");
2012        assert!(write_gzip_members_n(&path, &data, 2).unwrap() > 1);
2013        let mut bytes = std::fs::read(&path).unwrap();
2014        assert_eq!(gunzip_all(&bytes).unwrap(), data);
2015        let at = bytes
2016            .windows(4)
2017            .position(|window| window == b"RCZ1")
2018            .expect("index");
2019        let count = u16::from_le_bytes([bytes[at + 4], bytes[at + 5]]) as usize;
2020        let len_at = at + 6 + 8 * count;
2021        let len = u64::from_le_bytes(bytes[len_at..len_at + 8].try_into().unwrap());
2022        assert!(len > 1);
2023        bytes[len_at..len_at + 8].copy_from_slice(&(len - 1).to_le_bytes());
2024        assert!(gunzip_all(&bytes).is_err());
2025    }
2026
2027    #[test]
2028    fn one_member_gzip_matches_the_streaming_encoder() {
2029        let data = b"Cu 1.0 2.0 3.0 0 1\n".repeat(32);
2030        let direct = gzip_member(&data).unwrap();
2031        let dir = tempfile::tempdir().unwrap();
2032        let path = dir.path().join("one.con.gz");
2033        assert_eq!(write_gzip_members_n(&path, &data, 1).unwrap(), 1);
2034        assert_eq!(std::fs::read(&path).unwrap(), direct);
2035        let auto = dir.path().join("auto.con.gz");
2036        assert_eq!(write_gzip_members(&auto, &data).unwrap(), 1);
2037        assert_eq!(std::fs::read(&auto).unwrap(), direct);
2038    }
2039
2040    #[test]
2041    fn default_deflate_threads_stay_within_the_machine() {
2042        let cores = std::thread::available_parallelism()
2043            .map(|count| count.get())
2044            .unwrap_or(1)
2045            .max(1);
2046        assert_eq!(default_deflate_threads(1), 1);
2047        assert_eq!(default_deflate_threads(cores.saturating_mul(4)), cores);
2048        assert!(default_deflate_threads(64) <= cores);
2049    }
2050
2051    #[test]
2052    fn gzip_worker_count_does_not_change_the_bytes() {
2053        let mut data = Vec::new();
2054        for i in 0..20_000 {
2055            writeln!(data, "Cu {i}.25 1.5 2.5 0 {i}").unwrap();
2056        }
2057        let dir = tempfile::tempdir().unwrap();
2058        let one = dir.path().join("t1.con.gz");
2059        let eight = dir.path().join("t8.con.gz");
2060        write_gzip_members_limited(&one, &data, 8, 1).unwrap();
2061        write_gzip_members_limited(&eight, &data, 8, 8).unwrap();
2062        assert_eq!(std::fs::read(&one).unwrap(), std::fs::read(&eight).unwrap());
2063        assert_eq!(read_file_bytes(&one).unwrap(), data);
2064    }
2065
2066    #[test]
2067    fn chunked_gzip_roundtrips_at_level_6_with_one_member_per_slice() {
2068        // A single gzip member of this payload is the old streaming writer.
2069        // Four members is the chunked level-6 writer.
2070        let unit = b"Cu 1.0 2.0 3.0 0 1\n";
2071        let data = unit.repeat(80_000);
2072        assert!(data.len() > (1 << 20));
2073        let dir = tempfile::tempdir().unwrap();
2074        let path = dir.path().join("wide.con.gz");
2075        let members = write_gzip_members(&path, &data).unwrap();
2076        assert!(
2077            members > 1,
2078            "a wide payload must be more than one gzip member"
2079        );
2080        let bytes = std::fs::read(&path).unwrap();
2081        assert_eq!(gzip_member_count(&bytes), members);
2082        assert_eq!(bytes[8], 0, "XFL 0 is gzip level 6, not the fast level");
2083        assert_eq!(read_file_bytes(&path).unwrap(), data);
2084    }
2085
2086    #[cfg(feature = "zstd")]
2087    #[test]
2088    fn chunked_zstd_roundtrips_at_level_3() {
2089        let data = b"0.0 1.0 2.0 0 0\n".repeat(80_000);
2090        let dir = tempfile::tempdir().unwrap();
2091        let path = dir.path().join("wide.con.zst");
2092        let members = write_zstd_members(&path, &data).unwrap();
2093        assert!(members > 1);
2094        assert_eq!(read_file_bytes(&path).unwrap(), data);
2095        assert_eq!(
2096            with_serial_decode(|| read_file_bytes(&path).unwrap()),
2097            data,
2098            "serial decode of an indexed file matches the parallel decode"
2099        );
2100    }
2101
2102    #[test]
2103    fn a_single_gzip_stream_without_an_index_still_reads() {
2104        let text = include_str!("../resources/test/tiny_cuh2.con");
2105        let dir = tempfile::tempdir().unwrap();
2106        let path = dir.path().join("one.con.gz");
2107        {
2108            let mut writer = gzip_writer(&path).unwrap();
2109            writer.write_all(text.as_bytes()).unwrap();
2110            writer.finish().unwrap();
2111        }
2112        let bytes = std::fs::read(&path).unwrap();
2113        assert!(gzip_member_starts(&bytes).is_none());
2114        assert_eq!(read_file_text(&path).unwrap(), text);
2115        assert_eq!(read_file_bytes(&path).unwrap(), text.as_bytes());
2116    }
2117
2118    #[test]
2119    fn a_foreign_multi_member_gzip_without_our_index_still_reads() {
2120        let mut bytes = gzip_member(b"alpha\n").unwrap();
2121        bytes.extend(gzip_member(b"beta\n").unwrap());
2122        assert!(gzip_member_starts(&bytes).is_none());
2123        let dir = tempfile::tempdir().unwrap();
2124        let path = dir.path().join("foreign.con.gz");
2125        std::fs::write(&path, &bytes).unwrap();
2126        assert_eq!(read_file_bytes(&path).unwrap(), b"alpha\nbeta\n");
2127    }
2128
2129    #[cfg(feature = "zstd")]
2130    #[test]
2131    fn a_single_zstd_frame_without_an_index_still_reads() {
2132        let text = include_str!("../resources/test/tiny_cuh2.con");
2133        let dir = tempfile::tempdir().unwrap();
2134        let path = dir.path().join("one.con.zst");
2135        std::fs::write(&path, zstd::stream::encode_all(text.as_bytes(), 3).unwrap()).unwrap();
2136        assert!(zstd_frame_starts(&std::fs::read(&path).unwrap()).is_none());
2137        assert_eq!(read_file_text(&path).unwrap(), text);
2138    }
2139
2140    fn cli_stdout(program: &str, args: &[&str]) -> std::process::Output {
2141        std::process::Command::new(program)
2142            .args(args)
2143            .output()
2144            .unwrap_or_else(|err| panic!("{program} failed to start: {err}"))
2145    }
2146
2147    #[test]
2148    fn gzip_cli_decodes_one_member_and_an_indexed_file() {
2149        let text = include_str!("../resources/test/tiny_cuh2.con");
2150        let wide = b"Cu 1.0 2.0 3.0 0 1\n".repeat(80_000);
2151        let dir = tempfile::tempdir().unwrap();
2152        let one = dir.path().join("one.con.gz");
2153        {
2154            let mut writer = gzip_writer(&one).unwrap();
2155            writer.write_all(text.as_bytes()).unwrap();
2156            writer.finish().unwrap();
2157        }
2158        let many = dir.path().join("wide.con.gz");
2159        write_gzip_members(&many, &wide).unwrap();
2160        for (path, expect) in [(&one, text.as_bytes()), (&many, wide.as_slice())] {
2161            let out = cli_stdout("gzip", &["-dc", path.to_str().unwrap()]);
2162            assert!(
2163                out.status.success(),
2164                "{}",
2165                String::from_utf8_lossy(&out.stderr)
2166            );
2167            assert_eq!(out.stdout, expect);
2168        }
2169    }
2170
2171    #[cfg(feature = "zstd")]
2172    #[test]
2173    fn zstd_cli_decodes_one_frame_and_an_indexed_file() {
2174        let text = include_str!("../resources/test/tiny_cuh2.con");
2175        let wide = b"0.0 1.0 2.0 0 0\n".repeat(80_000);
2176        let dir = tempfile::tempdir().unwrap();
2177        let one = dir.path().join("one.con.zst");
2178        std::fs::write(&one, zstd::stream::encode_all(text.as_bytes(), 3).unwrap()).unwrap();
2179        let many = dir.path().join("wide.con.zst");
2180        write_zstd_members(&many, &wide).unwrap();
2181        for (path, expect) in [(&one, text.as_bytes()), (&many, wide.as_slice())] {
2182            let out = cli_stdout("zstd", &["-dc", path.to_str().unwrap()]);
2183            assert!(
2184                out.status.success(),
2185                "{}",
2186                String::from_utf8_lossy(&out.stderr)
2187            );
2188            assert_eq!(out.stdout, expect);
2189        }
2190    }
2191
2192    fn copper_frame(n: usize) -> crate::types::ConFrame {
2193        let mut builder =
2194            crate::types::ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
2195        for i in 0..n {
2196            builder.add_atom(
2197                "Cu",
2198                i as f64,
2199                0.25,
2200                0.5,
2201                [false, false, false],
2202                i as u64,
2203                63.546,
2204            );
2205        }
2206        builder.build().unwrap()
2207    }
2208
2209    fn coords(frame: &crate::types::ConFrame) -> Vec<f64> {
2210        frame
2211            .atom_data
2212            .iter()
2213            .flat_map(|atom| [atom.x(), atom.y(), atom.z()])
2214            .collect()
2215    }
2216
2217    #[test]
2218    fn fused_gzip_round_trips_an_odd_atom_count_on_one_and_many_threads() {
2219        let frame = copper_frame(5);
2220        let dir = tempfile::tempdir().unwrap();
2221        let path = dir.path().join("odd.con.gz");
2222        let members = write_fused_gzip_frame(&path, &frame, 3, 2).unwrap();
2223        assert!(members >= 2);
2224        let one = crate::iterators::read_all_frames_phased(&path, Some(1)).unwrap();
2225        let many = crate::iterators::read_all_frames_phased(&path, Some(4)).unwrap();
2226        assert_eq!(one.frames.len(), 1);
2227        assert_eq!(coords(&one.frames[0]), coords(&frame));
2228        assert_eq!(coords(&many.frames[0]), coords(&frame));
2229    }
2230
2231    #[test]
2232    fn fused_gzip_round_trips_two_species_on_many_threads() {
2233        let mut builder =
2234            crate::types::ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
2235        for i in 0..4 {
2236            builder.add_atom("Cu", i as f64, 0.0, 0.0, [false; 3], i as u64, 63.546);
2237        }
2238        for i in 0..3 {
2239            builder.add_atom("H", i as f64, 1.0, 0.0, [false; 3], 4 + i as u64, 1.008);
2240        }
2241        let frame = builder.build().unwrap();
2242        let dir = tempfile::tempdir().unwrap();
2243        let path = dir.path().join("alloy.con.gz");
2244        write_fused_gzip_frame(&path, &frame, 2, 2).unwrap();
2245        let one = crate::iterators::read_all_frames_phased(&path, Some(1)).unwrap();
2246        let many = crate::iterators::read_all_frames_phased(&path, Some(4)).unwrap();
2247        assert_eq!(coords(&one.frames[0]), coords(&frame));
2248        assert_eq!(coords(&many.frames[0]), coords(&frame));
2249        assert_eq!(
2250            many.frames[0].atom_data[0].symbol.as_ref(),
2251            frame.atom_data[0].symbol.as_ref()
2252        );
2253        assert_eq!(
2254            many.frames[0].atom_data[4].symbol.as_ref(),
2255            frame.atom_data[4].symbol.as_ref()
2256        );
2257    }
2258
2259    #[test]
2260    fn fused_worker_count_does_not_change_the_bytes() {
2261        let frame = copper_frame(4_096);
2262        let dir = tempfile::tempdir().unwrap();
2263        let one = dir.path().join("one.con.gz");
2264        let eight = dir.path().join("eight.con.gz");
2265        assert!(write_fused_gzip_frame(&one, &frame, 8, 1).unwrap() > 1);
2266        assert!(write_fused_gzip_frame(&eight, &frame, 8, 8).unwrap() > 1);
2267        assert_eq!(std::fs::read(&one).unwrap(), std::fs::read(&eight).unwrap());
2268        let loaded = crate::iterators::read_all_frames_phased(&eight, Some(8)).unwrap();
2269        assert_eq!(coords(&loaded.frames[0]), coords(&frame));
2270    }
2271
2272    #[cfg(feature = "zstd")]
2273    #[test]
2274    fn fused_zstd_worker_count_does_not_change_the_bytes() {
2275        let frame = copper_frame(4_096);
2276        let dir = tempfile::tempdir().unwrap();
2277        let one = dir.path().join("one.con.zst");
2278        let eight = dir.path().join("eight.con.zst");
2279        assert!(write_fused_zstd_frame(&one, &frame, 8, 1).unwrap() > 1);
2280        assert!(write_fused_zstd_frame(&eight, &frame, 8, 8).unwrap() > 1);
2281        assert_eq!(std::fs::read(&one).unwrap(), std::fs::read(&eight).unwrap());
2282        let loaded = crate::iterators::read_all_frames_phased(&eight, Some(4)).unwrap();
2283        assert_eq!(coords(&loaded.frames[0]), coords(&frame));
2284    }
2285}