1use std::io::{self, Read, Write};
6use std::path::Path;
7
8pub const GZIP_LEVEL: u32 = 6;
15
16pub const ZSTD_LEVEL: i32 = 3;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum Compression {
27 None,
28 Gzip,
29 Zstd,
32}
33
34fn 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
48pub 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
64pub 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
76pub const MMAP_THRESHOLD: u64 = 64 * 1024;
79
80pub fn read_file_contents(path: &Path) -> Result<FileContents, Box<dyn std::error::Error>> {
91 read_file_contents_detected(path).map(|(contents, _)| contents)
92}
93
94pub fn read_file_contents_detected(
96 path: &Path,
97) -> Result<(FileContents, Compression), Box<dyn std::error::Error>> {
98 read_contents(path, true)
99}
100
101pub 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
111pub 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
131pub 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
155pub 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 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 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 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
257pub 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
279pub 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 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
310const RELEASE_STEP: usize = 32 << 20;
312
313pub struct ForwardText {
324 contents: FileContents,
325 released: usize,
326}
327
328impl ForwardText {
329 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 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 FileContents::Mapped(m) => unsafe { std::str::from_utf8_unchecked(m) },
355 }
356 }
357
358 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
373pub 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 #[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
410fn release(map: &memmap2::Mmap, start: usize, end: usize) -> io::Result<()> {
413 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 #[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
433pub 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
442pub 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
461pub 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
511fn 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 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 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
634pub fn write_gzip_members(path: &Path, data: &[u8]) -> io::Result<usize> {
641 write_gzip_members_n(path, data, compression_workers(data.len()))
642}
643
644pub 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
652pub 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#[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
814fn 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 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
923pub(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
932pub 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
948pub 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
1061fn 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#[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 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 out.clear();
1129 if out.capacity() < total {
1130 let extra = total - out.len();
1131 out.reserve(extra);
1132 }
1133 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 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#[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 let encoder = zstd::stream::write::Encoder::new(file, ZSTD_LEVEL)?;
1274 Ok(encoder.auto_finish())
1275}
1276
1277#[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#[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#[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 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 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
1501pub 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
1521pub 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
1567pub(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
1628pub(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 #[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 #[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 #[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 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 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}