Skip to main content

lzip_parallel/
parallel.rs

1//! Parallel ZIP decompression for small files (< SMALL_FILE_THRESHOLD).
2//!
3//! Reads the Central Directory once, then decompresses all entries
4//! concurrently on the rayon thread pool.  Results are returned in
5//! Central Directory order.
6//!
7//! For large files use `reader::decompress_batched` instead.
8
9use crate::central_dir;
10use crate::entry::{self, ZipEntry, ZipError};
11
12/// Decompress all file entries in `data` in parallel.
13///
14/// `data` must be a complete ZIP file in memory.
15/// Directory entries are skipped.  All file entries are decoded
16/// concurrently and returned in Central Directory order.
17pub fn decompress_parallel(data: &[u8]) -> Result<Vec<ZipEntry>, ZipError> {
18    decompress_parallel_filter(data, "")
19}
20
21/// Decompress only entries whose name contains `needle`.
22///
23/// Empty needle matches all entries (same as `decompress_parallel`).
24/// Pre-filters at the central directory level — non-matching entries are never inflated.
25pub fn decompress_parallel_filter(data: &[u8], needle: &str) -> Result<Vec<ZipEntry>, ZipError> {
26    let locations = central_dir::read_central_directory(data)?;
27    let matched: Vec<_> = locations.iter()
28        .filter(|l| !l.is_directory && (needle.is_empty() || l.name.contains(needle)))
29        .collect();
30
31    let results: Vec<Result<ZipEntry, ZipError>> =
32        gatling::gatling_forkjoin::gatling_for_each_balanced(
33            matched.len(),
34            0,
35            1,
36            |i| matched[i].uncompressed_size,
37            |i| {
38                let loc = matched[i];
39                let bytes = entry::decompress_entry(data, loc)?;
40                Ok(ZipEntry { name: loc.name.clone(), data: bytes })
41            },
42        );
43
44    results.into_iter().collect()
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50    use std::io::Write;
51
52    fn make_jar(items: &[(&str, &[u8])]) -> Vec<u8> {
53        use std::io::Cursor;
54        use zip::write::SimpleFileOptions;
55        let mut buf = Vec::new();
56        let mut zw = zip::ZipWriter::new(Cursor::new(&mut buf));
57        let opts = SimpleFileOptions::default()
58            .compression_method(zip::CompressionMethod::Deflated);
59        for (name, data) in items {
60            zw.start_file(*name, opts).unwrap();
61            zw.write_all(data).unwrap();
62        }
63        zw.finish().unwrap();
64        buf
65    }
66
67    #[test]
68    fn parallel_single() {
69        let jar = make_jar(&[("Hello.class", b"cafebabe deadbeef")]);
70        let entries = decompress_parallel(&jar).unwrap();
71        assert_eq!(entries.len(), 1);
72        assert_eq!(entries[0].name, "Hello.class");
73        assert_eq!(entries[0].data, b"cafebabe deadbeef");
74    }
75
76    #[test]
77    fn parallel_multi() {
78        let jar = make_jar(&[
79            ("a.txt", b"aaa"),
80            ("b.txt", b"bbb"),
81            ("c.txt", b"ccc"),
82        ]);
83        let entries = decompress_parallel(&jar).unwrap();
84        assert_eq!(entries.len(), 3);
85        assert_eq!(entries[0].data, b"aaa");
86        assert_eq!(entries[1].data, b"bbb");
87        assert_eq!(entries[2].data, b"ccc");
88    }
89
90    #[test]
91    fn parallel_skips_directories() {
92        let jar = make_jar(&[
93            ("META-INF/", b""),
94            ("META-INF/MANIFEST.MF", b"Manifest-Version: 1.0\n"),
95        ]);
96        let entries = decompress_parallel(&jar).unwrap();
97        assert_eq!(entries.len(), 1);
98        assert_eq!(entries[0].name, "META-INF/MANIFEST.MF");
99    }
100
101    #[test]
102    fn filter_by_needle() {
103        let jar = make_jar(&[
104            ("com/example/App.class", b"bytecode"),
105            ("META-INF/maven/com.example/app/pom.xml", b"<project/>"),
106            ("META-INF/MANIFEST.MF", b"Manifest-Version: 1.0\n"),
107        ]);
108        let entries = decompress_parallel_filter(&jar, "pom.xml").unwrap();
109        assert_eq!(entries.len(), 1);
110        assert_eq!(entries[0].name, "META-INF/maven/com.example/app/pom.xml");
111        assert_eq!(entries[0].data, b"<project/>");
112    }
113
114    #[test]
115    fn filter_empty_matches_all() {
116        let jar = make_jar(&[("a.txt", b"a"), ("b.txt", b"b")]);
117        let entries = decompress_parallel_filter(&jar, "").unwrap();
118        assert_eq!(entries.len(), 2);
119    }
120}
121
122/// Decompress only entries whose name exactly matches one of `names`.
123///
124/// Empty slice returns nothing. Pre-filters at the central directory level —
125/// only matching entries are dispatched to rayon workers.
126pub fn decompress_parallel_filter_set(data: &[u8], names: &[&str]) -> Result<Vec<ZipEntry>, ZipError> {
127    use std::collections::HashSet;
128    if names.is_empty() {
129        return Ok(Vec::new());
130    }
131    let name_set: HashSet<&str> = names.iter().copied().collect();
132    let locations = central_dir::read_central_directory(data)?;
133    let matched: Vec<_> = locations.iter()
134        .filter(|l| !l.is_directory && name_set.contains(l.name.as_str()))
135        .collect();
136
137    let results: Vec<Result<ZipEntry, ZipError>> =
138        gatling::gatling_forkjoin::gatling_for_each_balanced(
139            matched.len(),
140            0,
141            1,
142            |i| matched[i].uncompressed_size,
143            |i| {
144                let loc = matched[i];
145                let bytes = entry::decompress_entry(data, loc)?;
146                Ok(ZipEntry { name: loc.name.clone(), data: bytes })
147            },
148        );
149
150    results.into_iter().collect()
151}
152
153/// Decompress only entries whose name ends with one of `suffixes`.
154///
155/// Empty slice returns nothing. Pre-filters at the central directory level —
156/// only matching entries are dispatched to rayon workers.
157pub fn decompress_parallel_filter_suffixes(data: &[u8], suffixes: &[&str]) -> Result<Vec<ZipEntry>, ZipError> {
158    if suffixes.is_empty() {
159        return Ok(Vec::new());
160    }
161    let locations = central_dir::read_central_directory(data)?;
162    let matched: Vec<_> = locations.iter()
163        .filter(|l| !l.is_directory && suffixes.iter().any(|s| l.name.ends_with(s)))
164        .collect();
165
166    let results: Vec<Result<ZipEntry, ZipError>> =
167        gatling::gatling_forkjoin::gatling_for_each_balanced(
168            matched.len(),
169            0,
170            1,
171            |i| matched[i].uncompressed_size,
172            |i| {
173                let loc = matched[i];
174                let bytes = entry::decompress_entry(data, loc)?;
175                Ok(ZipEntry { name: loc.name.clone(), data: bytes })
176            },
177        );
178
179    results.into_iter().collect()
180}
181
182#[cfg(all(test, feature = "gen"))]
183mod filter_set_tests {
184    use super::*;
185    use std::io::Write;
186
187    fn make_zip(items: &[(&str, &[u8])]) -> Vec<u8> {
188        use std::io::Cursor;
189        use zip::write::SimpleFileOptions;
190        let mut buf = Vec::new();
191        let mut zw = zip::ZipWriter::new(Cursor::new(&mut buf));
192        let opts = SimpleFileOptions::default()
193            .compression_method(zip::CompressionMethod::Deflated);
194        for (name, data) in items {
195            zw.start_file(*name, opts).unwrap();
196            zw.write_all(data).unwrap();
197        }
198        zw.finish().unwrap();
199        buf
200    }
201
202    #[test]
203    fn filter_set_exact_match() {
204        let zip = make_zip(&[
205            ("numpy/__init__.py", b"import numpy"),
206            ("numpy-1.26.0.dist-info/METADATA", b"Name: numpy"),
207            ("numpy-1.26.0.dist-info/RECORD", b"files here"),
208            ("numpy/core/setup.py", b"setup()"),
209        ]);
210        let entries = decompress_parallel_filter_set(
211            &zip,
212            &["numpy-1.26.0.dist-info/METADATA", "numpy-1.26.0.dist-info/RECORD"],
213        ).unwrap();
214        assert_eq!(entries.len(), 2);
215        assert!(entries.iter().any(|e| e.data == b"Name: numpy"));
216        assert!(entries.iter().any(|e| e.data == b"files here"));
217    }
218
219    #[test]
220    fn filter_set_empty_returns_nothing() {
221        let zip = make_zip(&[("a.txt", b"a")]);
222        let entries = decompress_parallel_filter_set(&zip, &[]).unwrap();
223        assert_eq!(entries.len(), 0);
224    }
225
226    #[test]
227    fn filter_set_no_match() {
228        let zip = make_zip(&[("a.txt", b"a"), ("b.txt", b"b")]);
229        let entries = decompress_parallel_filter_set(&zip, &["c.txt"]).unwrap();
230        assert_eq!(entries.len(), 0);
231    }
232
233    #[test]
234    fn filter_suffixes_py() {
235        let zip = make_zip(&[
236            ("lib/__init__.py", b"init"),
237            ("lib/core.py", b"core"),
238            ("lib/types.pyi", b"types"),
239            ("lib/data.json", b"{}"),
240        ]);
241        let entries = decompress_parallel_filter_suffixes(&zip, &[".py", ".pyi"]).unwrap();
242        assert_eq!(entries.len(), 3);
243        assert!(entries.iter().all(|e| e.name.ends_with(".py") || e.name.ends_with(".pyi")));
244    }
245
246    #[test]
247    fn filter_suffixes_empty_returns_nothing() {
248        let zip = make_zip(&[("a.py", b"a")]);
249        let entries = decompress_parallel_filter_suffixes(&zip, &[]).unwrap();
250        assert_eq!(entries.len(), 0);
251    }
252}