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