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