readcon_core/
compression.rs

1//=============================================================================
2// Transparent compression support
3//=============================================================================
4
5use std::io::{self, Read};
6use std::path::Path;
7
8/// Detected compression format based on magic bytes.
9///
10/// `Zstd` is only constructed when the `zstd` Cargo feature is enabled.
11/// Builds without the feature treat `.zst` files as opaque bytes and
12/// return an error from [`read_file_contents`] indicating the feature
13/// is required.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum Compression {
16    None,
17    Gzip,
18    /// zstd frame, magic `28 B5 2F FD`. Build with `--features zstd`.
19    Zstd,
20}
21
22/// Detect compression format from the first bytes of a file.
23///
24/// - `1f 8b` = gzip
25/// - `28 b5 2f fd` = zstd
26/// - Otherwise = uncompressed
27pub fn detect_compression(bytes: &[u8]) -> Compression {
28    if bytes.len() >= 2 && bytes[0] == 0x1f && bytes[1] == 0x8b {
29        return Compression::Gzip;
30    }
31    if bytes.len() >= 4
32        && bytes[0] == 0x28
33        && bytes[1] == 0xb5
34        && bytes[2] == 0x2f
35        && bytes[3] == 0xfd
36    {
37        return Compression::Zstd;
38    }
39    Compression::None
40}
41
42/// Detect compression format from a file extension.
43///
44/// Returns `Compression::Gzip` for `.gz`, `Compression::Zstd` for
45/// `.zst`, `Compression::None` otherwise.
46pub fn detect_compression_from_extension(path: &Path) -> Compression {
47    match path.extension().and_then(|e| e.to_str()) {
48        Some("gz") => Compression::Gzip,
49        Some("zst") => Compression::Zstd,
50        _ => Compression::None,
51    }
52}
53
54/// Size threshold below which we use `read_to_string` instead of mmap.
55const MMAP_THRESHOLD: u64 = 64 * 1024;
56
57/// Reads file contents, decompressing if needed.
58///
59/// Detection strategy:
60/// 1. Read first 2 bytes to check magic bytes.
61/// 2. If gzip: decompress entire file to a String.
62/// 3. If uncompressed and < 64 KiB: `read_to_string`.
63/// 4. If uncompressed and >= 64 KiB: memory-mapped I/O.
64pub fn read_file_contents(path: &Path) -> Result<FileContents, Box<dyn std::error::Error>> {
65    let file = std::fs::File::open(path)?;
66    let metadata = file.metadata()?;
67
68    // Read first 4 bytes for magic detection (gzip needs 2, zstd needs 4)
69    let mut magic = [0u8; 4];
70    let bytes_read = {
71        let mut f = &file;
72        f.read(&mut magic)?
73    };
74
75    let compression = if bytes_read > 0 {
76        detect_compression(&magic[..bytes_read])
77    } else {
78        Compression::None
79    };
80
81    match compression {
82        Compression::Gzip => {
83            // Re-open and decompress the entire file
84            let file = std::fs::File::open(path)?;
85            let mut decoder = flate2::read::GzDecoder::new(file);
86            let mut contents = String::new();
87            decoder.read_to_string(&mut contents)?;
88            Ok(FileContents::Owned(contents))
89        }
90        Compression::Zstd => {
91            #[cfg(feature = "zstd")]
92            {
93                let file = std::fs::File::open(path)?;
94                let mut decoder = zstd::stream::read::Decoder::new(file)?;
95                let mut contents = String::new();
96                decoder.read_to_string(&mut contents)?;
97                Ok(FileContents::Owned(contents))
98            }
99            #[cfg(not(feature = "zstd"))]
100            {
101                Err(io::Error::new(
102                    io::ErrorKind::Unsupported,
103                    "zstd-compressed input detected; rebuild readcon-core with --features zstd",
104                )
105                .into())
106            }
107        }
108        Compression::None => {
109            if metadata.len() < MMAP_THRESHOLD {
110                let contents = std::fs::read_to_string(path)?;
111                Ok(FileContents::Owned(contents))
112            } else {
113                let file = std::fs::File::open(path)?;
114                let mmap = unsafe { memmap2::Mmap::map(&file)? };
115                Ok(FileContents::Mapped(mmap))
116            }
117        }
118    }
119}
120
121/// Holds file contents either as an owned String or a memory-mapped region.
122pub enum FileContents {
123    Owned(String),
124    Mapped(memmap2::Mmap),
125}
126
127impl FileContents {
128    pub fn as_str(&self) -> Result<&str, std::str::Utf8Error> {
129        match self {
130            FileContents::Owned(s) => Ok(s.as_str()),
131            FileContents::Mapped(m) => std::str::from_utf8(m),
132        }
133    }
134}
135
136/// Creates a gzip-compressed writer wrapping a file at the given path.
137pub fn gzip_writer(path: &Path) -> io::Result<flate2::write::GzEncoder<std::fs::File>> {
138    let file = std::fs::File::create(path)?;
139    Ok(flate2::write::GzEncoder::new(
140        file,
141        flate2::Compression::default(),
142    ))
143}
144
145/// Creates a zstd-compressed writer wrapping a file at the given path.
146///
147/// Returns a finished writer ready for streaming line writes; the caller
148/// is responsible for dropping the writer to flush the final frame.
149/// Available only when the `zstd` Cargo feature is enabled.
150#[cfg(feature = "zstd")]
151pub fn zstd_writer(
152    path: &Path,
153) -> io::Result<zstd::stream::write::AutoFinishEncoder<'static, std::fs::File>> {
154    let file = std::fs::File::create(path)?;
155    // Level 3 is the zstd CLI default; balances ratio against speed.
156    let encoder = zstd::stream::write::Encoder::new(file, 3)?;
157    Ok(encoder.auto_finish())
158}
readcon_core/compression.rs

readcon_core/
compression.rs
