1use rayon::prelude::*;
10
11use crate::central_dir;
12use crate::entry::{self, ZipEntry, ZipError};
13
14pub fn decompress_parallel(data: &[u8]) -> Result<Vec<ZipEntry>, ZipError> {
20 decompress_parallel_filter(data, "")
21}
22
23pub 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
121pub 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
149pub 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}