1#![cfg_attr(not(feature = "std"), no_std)]
2#![cfg_attr(feature = "nightly", feature(optimize_attribute))]
3
4#[cfg(feature = "alloc")]
5extern crate alloc;
6
7pub(crate) mod block_decoder;
8#[cfg(feature = "std")]
9pub mod context;
10pub(crate) mod exec;
11pub(crate) mod literals;
12pub(crate) mod primitives;
13pub(crate) mod ring_buffer;
14pub(crate) mod sequences;
15#[cfg(feature = "std")]
16pub mod streaming;
17
18#[allow(dead_code)]
19pub(crate) mod simd_decode;
20
21#[cfg(feature = "alloc")]
22use alloc::boxed::Box;
23#[cfg(feature = "alloc")]
24use alloc::vec::Vec;
25
26use crate::exec::decode_execute_sequences;
27use crate::literals::decode_literals_ws;
28use crate::sequences::{SequenceDecodeTables, parse_sequence_count, parse_sequence_tables_ws};
29use zrip_core::block::{BlockType, parse_block_header};
30use zrip_core::error::DecompressError;
31use zrip_core::frame::MAX_WINDOW_SIZE;
32use zrip_core::frame::header::parse_frame_header;
33use zrip_core::huffman::HuffmanDecodeEntry;
34#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
35use zrip_core::simd::CpuTier;
36use zrip_core::xxhash::Xxh64State;
37
38pub(crate) struct BlockDecodeWorkspace {
39 pub literal_buf: Vec<u8>,
40 pub huf_table: Vec<HuffmanDecodeEntry>,
41 pub huf_table_log: u8,
42 pub huf_valid: bool,
43 pub huf_all_weights: Vec<u8>,
44 pub huf_rank_count: Vec<u32>,
45 pub huf_rank_start: Vec<u32>,
46 pub fse_dist: Vec<i16>,
47 pub fse_symbol_next: Vec<u16>,
48 pub fse_build_buf: Vec<zrip_core::fse::FseDecodeEntry>,
49}
50
51impl BlockDecodeWorkspace {
52 pub(crate) fn new() -> Self {
53 Self {
54 literal_buf: Vec::new(),
55 huf_table: Vec::new(),
56 huf_table_log: 0,
57 huf_valid: false,
58 huf_all_weights: Vec::new(),
59 huf_rank_count: Vec::new(),
60 huf_rank_start: Vec::new(),
61 fse_dist: Vec::new(),
62 fse_symbol_next: Vec::new(),
63 fse_build_buf: Vec::new(),
64 }
65 }
66}
67
68pub(crate) fn skip_skippable_frame(data: &[u8]) -> Option<usize> {
69 if data.len() < 8 {
70 return None;
71 }
72 let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
73 if (magic & 0xFFFFFFF0) != 0x184D2A50 {
74 return None;
75 }
76 let frame_size = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize;
77 let total = 8 + frame_size;
78 if total > data.len() {
79 return None;
80 }
81 Some(total)
82}
83
84pub fn decompress(input: &[u8]) -> Result<Vec<u8>, DecompressError> {
85 decompress_with_dict(input, None)
86}
87
88pub fn decompress_into(input: &[u8], output: &mut Vec<u8>) -> Result<usize, DecompressError> {
89 let max_output = zrip_core::DEFAULT_DECOMPRESS_LIMIT;
90 let mut ws = Box::new(BlockDecodeWorkspace::new());
91 let start = output.len();
92 let mut offset = 0;
93 while offset < input.len() {
94 let remaining = &input[offset..];
95 if let Some(skip_len) = skip_skippable_frame(remaining) {
96 offset += skip_len;
97 continue;
98 }
99 let consumed = decompress_frame(remaining, output, max_output, None, &mut ws)?;
100 offset += consumed;
101 }
102 Ok(output.len() - start)
103}
104
105pub fn decompress_with_dict(
106 input: &[u8],
107 dict: Option<&zrip_core::dict::Dictionary>,
108) -> Result<Vec<u8>, DecompressError> {
109 let max_output = zrip_core::DEFAULT_DECOMPRESS_LIMIT;
110 let mut output = Vec::new();
111 let mut ws = Box::new(BlockDecodeWorkspace::new());
112 let mut offset = 0;
113
114 while offset < input.len() {
115 let remaining = &input[offset..];
116 if let Some(skip_len) = skip_skippable_frame(remaining) {
117 offset += skip_len;
118 continue;
119 }
120 let consumed = decompress_frame(remaining, &mut output, max_output, dict, &mut ws)?;
121 offset += consumed;
122 }
123
124 Ok(output)
125}
126
127pub(crate) fn decompress_frame(
128 input: &[u8],
129 output: &mut Vec<u8>,
130 max_output: usize,
131 dict: Option<&zrip_core::dict::Dictionary>,
132 ws: &mut BlockDecodeWorkspace,
133) -> Result<usize, DecompressError> {
134 let header = parse_frame_header(input)?;
135
136 if header.window_size > MAX_WINDOW_SIZE {
137 return Err(DecompressError::WindowTooLarge {
138 requested: header.window_size,
139 max: MAX_WINDOW_SIZE,
140 });
141 }
142
143 if let Some(frame_dict_id) = header.dict_id {
144 match dict {
145 Some(d) if d.id() == frame_dict_id => {}
146 Some(d) => {
147 return Err(DecompressError::DictMismatch {
148 expected: frame_dict_id,
149 got: d.id(),
150 });
151 }
152 None => return Err(DecompressError::DictRequired),
153 }
154 }
155
156 if let Some(fcs) = header.frame_content_size {
157 if max_output < usize::MAX && fcs as usize > max_output {
158 return Err(DecompressError::OutputTooSmall);
159 }
160 let hint = (fcs as usize).min(MAX_WINDOW_SIZE as usize);
161 output.reserve(hint + 32);
162 }
163
164 let mut offset = header.header_size;
165 let output_start = output.len();
166
167 let dict_history: &[u8] = if let Some(d) = dict { d.content() } else { &[] };
168
169 let mut seq_tables = if let Some(d) = dict {
170 let mut st = SequenceDecodeTables::new_default();
171 if let Some((t, l)) = d.of_table() {
172 st.of_table = zrip_core::fse::promote_of_table(t);
173 st.of_accuracy = l;
174 }
175 if let Some((t, l)) = d.ml_table() {
176 st.ml_table = zrip_core::fse::promote_ml_table(t);
177 st.ml_accuracy = l;
178 }
179 if let Some((t, l)) = d.ll_table() {
180 st.ll_table = zrip_core::fse::promote_ll_table(t);
181 st.ll_accuracy = l;
182 }
183 st
184 } else {
185 SequenceDecodeTables::new_default()
186 };
187 let mut rep_offsets: [u32; 3] = if let Some(d) = dict {
188 *d.rep_offsets()
189 } else {
190 [1, 4, 8]
191 };
192 ws.huf_valid = false;
193 if let Some(d) = dict {
194 if let Some((t, l)) = d.huf_table() {
195 ws.huf_table.clear();
196 ws.huf_table.extend_from_slice(t);
197 ws.huf_table_log = l;
198 ws.huf_valid = true;
199 }
200 }
201
202 let mut hasher = if header.content_checksum {
203 Some(Xxh64State::new(0))
204 } else {
205 None
206 };
207
208 loop {
209 if offset + 3 > input.len() {
210 return Err(DecompressError::InputExhausted);
211 }
212 let block_header = parse_block_header(&input[offset..])?;
213 offset += 3;
214
215 let block_size = block_header.block_size as usize;
216
217 if block_size > zrip_core::frame::MAX_BLOCK_SIZE {
218 match block_header.block_type {
219 BlockType::Raw | BlockType::Rle => {
220 return Err(DecompressError::CorruptSequences);
221 }
222 BlockType::Compressed => {}
223 }
224 }
225
226 match block_header.block_type {
227 BlockType::Raw => {
228 if offset + block_size > input.len() {
229 return Err(DecompressError::InputExhausted);
230 }
231 if output.len() - output_start + block_size > max_output {
232 return Err(DecompressError::OutputTooSmall);
233 }
234 output.extend_from_slice(&input[offset..offset + block_size]);
235 offset += block_size;
236 }
237 BlockType::Rle => {
238 if offset >= input.len() {
239 return Err(DecompressError::InputExhausted);
240 }
241 if output.len() - output_start + block_size > max_output {
242 return Err(DecompressError::OutputTooSmall);
243 }
244 let byte = input[offset];
245 output.resize(output.len() + block_size, byte);
246 offset += 1;
247 }
248 BlockType::Compressed => {
249 if offset + block_size > input.len() {
250 return Err(DecompressError::InputExhausted);
251 }
252 let block_data = &input[offset..offset + block_size];
253 decode_compressed_block(
254 block_data,
255 output,
256 output_start,
257 max_output,
258 &mut seq_tables,
259 &mut rep_offsets,
260 ws,
261 dict_history,
262 )?;
263 offset += block_size;
264 }
265 }
266
267 if block_header.last_block {
268 break;
269 }
270 }
271
272 if let Some(ref mut hasher) = hasher {
273 hasher.update(&output[output_start..]);
274 let hash = hasher.finish();
275 let expected_checksum = (hash & 0xFFFFFFFF) as u32;
276
277 if offset + 4 > input.len() {
278 return Err(DecompressError::InputExhausted);
279 }
280 let stored_checksum = u32::from_le_bytes([
281 input[offset],
282 input[offset + 1],
283 input[offset + 2],
284 input[offset + 3],
285 ]);
286 offset += 4;
287
288 if expected_checksum != stored_checksum {
289 return Err(DecompressError::ChecksumMismatch {
290 expected: stored_checksum,
291 got: expected_checksum,
292 });
293 }
294 }
295
296 if let Some(fcs) = header.frame_content_size {
297 if (output.len() - output_start) as u64 != fcs {
298 return Err(DecompressError::CorruptSequences);
299 }
300 }
301
302 Ok(offset)
303}
304
305#[allow(clippy::too_many_arguments)]
306fn decode_compressed_block(
307 data: &[u8],
308 output: &mut Vec<u8>,
309 output_start: usize,
310 max_output: usize,
311 seq_tables: &mut SequenceDecodeTables,
312 rep_offsets: &mut [u32; 3],
313 ws: &mut BlockDecodeWorkspace,
314 dict_history: &[u8],
315) -> Result<(), DecompressError> {
316 let lit_consumed = decode_literals_ws(data, ws)?;
317
318 let remaining = &data[lit_consumed..];
319
320 if remaining.is_empty() {
321 if output.len() - output_start + ws.literal_buf.len() > max_output {
322 return Err(DecompressError::OutputTooSmall);
323 }
324 output.extend_from_slice(&ws.literal_buf);
325 return Ok(());
326 }
327
328 let (num_sequences, seq_count_size) = parse_sequence_count(remaining)?;
329
330 if num_sequences == 0 {
331 if output.len() - output_start + ws.literal_buf.len() > max_output {
332 return Err(DecompressError::OutputTooSmall);
333 }
334 output.extend_from_slice(&ws.literal_buf);
335 return Ok(());
336 }
337
338 let table_data = &remaining[seq_count_size..];
339 let tables_consumed = parse_sequence_tables_ws(table_data, seq_tables, ws)?;
340
341 let seq_data = &table_data[tables_consumed..];
342
343 let before = output.len();
344
345 #[cfg(target_arch = "x86_64")]
346 {
347 if zrip_core::simd::cpu_tier() >= CpuTier::Avx2 {
348 decode_execute_block_avx2(
349 seq_data,
350 num_sequences,
351 seq_tables,
352 rep_offsets,
353 &ws.literal_buf,
354 output,
355 dict_history,
356 )?;
357 if output.len() - before > zrip_core::frame::MAX_BLOCK_SIZE {
358 return Err(DecompressError::CorruptSequences);
359 }
360 return Ok(());
361 }
362 }
363 #[cfg(target_arch = "aarch64")]
364 {
365 if zrip_core::simd::cpu_tier() >= CpuTier::Neon {
366 decode_execute_block_neon(
367 seq_data,
368 num_sequences,
369 seq_tables,
370 rep_offsets,
371 &ws.literal_buf,
372 output,
373 dict_history,
374 )?;
375 if output.len() - before > zrip_core::frame::MAX_BLOCK_SIZE {
376 return Err(DecompressError::CorruptSequences);
377 }
378 return Ok(());
379 }
380 }
381
382 decode_execute_sequences(
383 seq_data,
384 num_sequences,
385 seq_tables,
386 rep_offsets,
387 &ws.literal_buf,
388 output,
389 dict_history,
390 )?;
391 if output.len() - before > zrip_core::frame::MAX_BLOCK_SIZE {
392 return Err(DecompressError::CorruptSequences);
393 }
394
395 Ok(())
396}
397
398#[cfg(target_arch = "x86_64")]
399fn decode_execute_block_avx2(
400 seq_data: &[u8],
401 num_sequences: u32,
402 tables: &mut SequenceDecodeTables,
403 rep_offsets: &mut [u32; 3],
404 literals: &[u8],
405 output: &mut Vec<u8>,
406 history: &[u8],
407) -> Result<(), DecompressError> {
408 crate::simd_decode::x86_64::decode::decode_execute_avx2_safe(
409 seq_data,
410 num_sequences,
411 tables,
412 rep_offsets,
413 literals,
414 output,
415 history,
416 )
417}
418
419#[cfg(target_arch = "aarch64")]
420fn decode_execute_block_neon(
421 seq_data: &[u8],
422 num_sequences: u32,
423 tables: &mut SequenceDecodeTables,
424 rep_offsets: &mut [u32; 3],
425 literals: &[u8],
426 output: &mut Vec<u8>,
427 history: &[u8],
428) -> Result<(), DecompressError> {
429 crate::simd_decode::aarch64::decode::decode_execute_neon_safe(
430 seq_data,
431 num_sequences,
432 tables,
433 rep_offsets,
434 literals,
435 output,
436 history,
437 )
438}