1pub(crate) mod block_encoder;
2#[cfg(feature = "std")]
3pub mod context;
4pub(crate) mod dfast;
5pub(crate) mod fast;
6pub(crate) mod literals;
7pub(crate) mod sequences;
8pub mod strategy;
9#[cfg(feature = "std")]
10pub mod streaming;
11pub(crate) mod unchecked;
12
13#[cfg(feature = "alloc")]
14use alloc::vec::Vec;
15
16use crate::encode::strategy::Strategy;
17use crate::error::CompressError;
18use crate::frame::{MAX_BLOCK_SIZE, ZSTD_MAGIC};
19use crate::xxhash::xxh64;
20
21pub(crate) fn block_looks_incompressible(data: &[u8]) -> bool {
22 const SAMPLE: usize = 1024;
23 const DISTINCT_THRESHOLD: u32 = 200;
24 const MAX_FREQ_DENOM: u32 = 24;
25 if data.len() < SAMPLE {
26 return false;
27 }
28 let mut counts = [0u16; 256];
29 for &b in &data[..SAMPLE] {
30 counts[b as usize] += 1;
31 }
32 let mut distinct: u32 = 0;
33 let mut max_freq: u16 = 0;
34 for &c in &counts {
35 distinct += (c > 0) as u32;
36 max_freq = max_freq.max(c);
37 }
38 distinct >= DISTINCT_THRESHOLD && (max_freq as u32) <= SAMPLE as u32 / MAX_FREQ_DENOM
39}
40
41pub(crate) fn clamp_params_to_src_size(params: &mut strategy::LevelParams, src_len: usize) {
42 if src_len >= 2 {
43 let src_log = 32 - ((src_len as u32) - 1).leading_zeros();
44 params.hash_log = params.hash_log.min(src_log);
45 params.window_log = params.window_log.min(src_log);
46 }
47}
48
49pub fn compress_with_params(
50 input: &[u8],
51 params: &strategy::LevelParams,
52) -> Result<Vec<u8>, CompressError> {
53 let mut params = *params;
54 clamp_params_to_src_size(&mut params, input.len());
55 compress_inner(input, ¶ms)
56}
57
58pub fn compress(input: &[u8], level: i32) -> Result<Vec<u8>, CompressError> {
59 let mut params = strategy::level_params(level).ok_or(CompressError::InvalidLevel(level))?;
60 clamp_params_to_src_size(&mut params, input.len());
61 compress_inner(input, ¶ms)
62}
63
64fn compress_inner(input: &[u8], params: &strategy::LevelParams) -> Result<Vec<u8>, CompressError> {
65 let mut output = Vec::with_capacity(input.len() + 32);
66 compress_frame(input, params, &mut output);
67 Ok(output)
68}
69
70fn compress_frame(input: &[u8], params: &strategy::LevelParams, output: &mut Vec<u8>) {
71 output.extend_from_slice(&ZSTD_MAGIC.to_le_bytes());
72
73 let fcs_size = if input.len() <= 255 {
74 1
75 } else if input.len() <= 0xFFFF + 256 {
76 2
77 } else if input.len() <= 0xFFFFFFFF {
78 4
79 } else {
80 8
81 };
82
83 let fcs_flag = match fcs_size {
84 1 => 0,
85 2 => 1,
86 4 => 2,
87 8 => 3,
88 _ => unreachable!(),
89 };
90
91 let descriptor = 0x20 | 0x04 | (fcs_flag << 6);
92 output.push(descriptor);
93
94 match fcs_size {
95 1 => output.push(input.len() as u8),
96 2 => {
97 let v = (input.len() - 256) as u16;
98 output.extend_from_slice(&v.to_le_bytes());
99 }
100 4 => output.extend_from_slice(&(input.len() as u32).to_le_bytes()),
101 8 => output.extend_from_slice(&(input.len() as u64).to_le_bytes()),
102 _ => unreachable!(),
103 }
104
105 if input.is_empty() {
106 block_encoder::encode_raw_block(&[], true, output);
107 } else {
108 let hash_size = 1usize << params.hash_log;
109 let mut rep_offsets = [1u32, 4, 8];
110 let mut offset = 0;
111 let mut sequences = Vec::with_capacity(MAX_BLOCK_SIZE / 8);
112 let mut workspace = block_encoder::BlockEncodeWorkspace::new();
113
114 match params.strategy {
115 Strategy::Fast => {
116 let mut hash_table = vec![0u32; hash_size];
117 while offset < input.len() {
118 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
119 let block_end = offset + chunk_size;
120 let is_last = block_end >= input.len();
121
122 if block_looks_incompressible(&input[offset..block_end]) {
123 block_encoder::encode_raw_block(&input[offset..block_end], is_last, output);
124 } else {
125 fast::compress_fast_block(
126 input,
127 offset,
128 block_end,
129 params,
130 &rep_offsets,
131 &mut hash_table,
132 &mut sequences,
133 );
134 block_encoder::encode_compressed_block(
135 &input[offset..block_end],
136 &sequences,
137 &mut rep_offsets,
138 is_last,
139 output,
140 &mut workspace,
141 );
142 }
143 offset = block_end;
144 }
145 }
146 Strategy::DFast => {
147 let mut hash_short = vec![0u32; hash_size];
148 let mut hash_long = vec![0u32; hash_size];
149 while offset < input.len() {
150 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
151 let block_end = offset + chunk_size;
152 let is_last = block_end >= input.len();
153
154 if block_looks_incompressible(&input[offset..block_end]) {
155 block_encoder::encode_raw_block(&input[offset..block_end], is_last, output);
156 } else {
157 dfast::compress_dfast_block(
158 input,
159 offset,
160 block_end,
161 params,
162 &rep_offsets,
163 &mut hash_short,
164 &mut hash_long,
165 &mut sequences,
166 );
167 block_encoder::encode_compressed_block(
168 &input[offset..block_end],
169 &sequences,
170 &mut rep_offsets,
171 is_last,
172 output,
173 &mut workspace,
174 );
175 }
176 offset = block_end;
177 }
178 }
179 }
180 }
181
182 let hash = xxh64(input, 0);
183 let checksum = (hash & 0xFFFFFFFF) as u32;
184 output.extend_from_slice(&checksum.to_le_bytes());
185}
186
187pub fn compress_with_dict(
188 input: &[u8],
189 level: i32,
190 dict: &crate::dict::Dictionary,
191) -> Result<Vec<u8>, CompressError> {
192 let mut params = strategy::level_params(level).ok_or(CompressError::InvalidLevel(level))?;
193 clamp_params_to_src_size(&mut params, input.len());
194
195 let mut output = Vec::with_capacity(input.len() + 32);
196
197 output.extend_from_slice(&ZSTD_MAGIC.to_le_bytes());
198
199 let fcs_size = if input.len() <= 255 {
200 1
201 } else if input.len() <= 0xFFFF + 256 {
202 2
203 } else if input.len() <= 0xFFFFFFFF {
204 4
205 } else {
206 8
207 };
208
209 let fcs_flag = match fcs_size {
210 1 => 0,
211 2 => 1,
212 4 => 2,
213 8 => 3,
214 _ => unreachable!(),
215 };
216
217 let dict_id = dict.id();
218 let dict_id_flag = if dict_id <= 0xFF {
219 1u8
220 } else if dict_id <= 0xFFFF {
221 2
222 } else {
223 3
224 };
225
226 let descriptor = 0x20 | 0x04 | (fcs_flag << 6) | dict_id_flag;
227 output.push(descriptor);
228
229 match dict_id_flag {
230 1 => output.push(dict_id as u8),
231 2 => output.extend_from_slice(&(dict_id as u16).to_le_bytes()),
232 3 => output.extend_from_slice(&dict_id.to_le_bytes()),
233 _ => unreachable!(),
234 }
235
236 match fcs_size {
237 1 => output.push(input.len() as u8),
238 2 => {
239 let v = (input.len() - 256) as u16;
240 output.extend_from_slice(&v.to_le_bytes());
241 }
242 4 => output.extend_from_slice(&(input.len() as u32).to_le_bytes()),
243 8 => output.extend_from_slice(&(input.len() as u64).to_le_bytes()),
244 _ => unreachable!(),
245 }
246
247 if input.is_empty() {
248 block_encoder::encode_raw_block(&[], true, &mut output);
249 } else {
250 let prefix = dict.content();
251 let mut rep_offsets = *dict.rep_offsets();
252 let mut workspace = block_encoder::BlockEncodeWorkspace::new();
253
254 if input.len() <= MAX_BLOCK_SIZE {
255 let sequences = match params.strategy {
256 Strategy::Fast => {
257 fast::compress_fast_with_prefix(input, ¶ms, &rep_offsets, prefix)
258 }
259 Strategy::DFast => {
260 dfast::compress_dfast_with_prefix(input, ¶ms, &rep_offsets, prefix)
261 }
262 };
263 block_encoder::encode_compressed_block(
264 input,
265 &sequences,
266 &mut rep_offsets,
267 true,
268 &mut output,
269 &mut workspace,
270 );
271 } else {
272 let mut combined = Vec::with_capacity(prefix.len() + input.len());
273 combined.extend_from_slice(prefix);
274 combined.extend_from_slice(input);
275 let plen = prefix.len();
276 let hash_size = 1usize << params.hash_log;
277 let mut sequences = Vec::new();
278
279 match params.strategy {
280 Strategy::Fast => {
281 let mut hash_table = vec![0u32; hash_size];
282 fast::prefill_hash_table(&combined, plen, params.hash_log, &mut hash_table);
283 let mut offset = 0;
284 while offset < input.len() {
285 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
286 let is_last = offset + chunk_size >= input.len();
287 fast::compress_fast_block(
288 &combined,
289 plen + offset,
290 plen + offset + chunk_size,
291 ¶ms,
292 &rep_offsets,
293 &mut hash_table,
294 &mut sequences,
295 );
296 block_encoder::encode_compressed_block(
297 &input[offset..offset + chunk_size],
298 &sequences,
299 &mut rep_offsets,
300 is_last,
301 &mut output,
302 &mut workspace,
303 );
304 offset += chunk_size;
305 }
306 }
307 Strategy::DFast => {
308 let mut hash_short = vec![0u32; hash_size];
309 let mut hash_long = vec![0u32; hash_size];
310 dfast::prefill_hash_tables(
311 &combined,
312 plen,
313 params.hash_log,
314 &mut hash_short,
315 &mut hash_long,
316 );
317 let mut offset = 0;
318 while offset < input.len() {
319 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
320 let is_last = offset + chunk_size >= input.len();
321 dfast::compress_dfast_block(
322 &combined,
323 plen + offset,
324 plen + offset + chunk_size,
325 ¶ms,
326 &rep_offsets,
327 &mut hash_short,
328 &mut hash_long,
329 &mut sequences,
330 );
331 block_encoder::encode_compressed_block(
332 &input[offset..offset + chunk_size],
333 &sequences,
334 &mut rep_offsets,
335 is_last,
336 &mut output,
337 &mut workspace,
338 );
339 offset += chunk_size;
340 }
341 }
342 }
343 }
344 }
345
346 let hash = xxh64(input, 0);
347 let checksum = (hash & 0xFFFFFFFF) as u32;
348 output.extend_from_slice(&checksum.to_le_bytes());
349
350 Ok(output)
351}
352
353pub fn compress_into(input: &[u8], output: &mut [u8], level: i32) -> Result<usize, CompressError> {
354 let mut params = strategy::level_params(level).ok_or(CompressError::InvalidLevel(level))?;
355 clamp_params_to_src_size(&mut params, input.len());
356 let mut buf = Vec::with_capacity(output.len());
357 compress_frame(input, ¶ms, &mut buf);
358 if buf.len() > output.len() {
359 return Err(CompressError::OutputTooSmall);
360 }
361 output[..buf.len()].copy_from_slice(&buf);
362 Ok(buf.len())
363}