use std::ffi::{OsStr, OsString};
use std::fs::{self, File, OpenOptions};
use std::io::{self, BufRead, BufReader, ErrorKind, IsTerminal, Read, Write};
use std::path::{Path, PathBuf};
use structured_zstd::encoding::{
CompressionContext, CompressionLevel, CompressionParameters, LevelParameters,
LiteralCompressionMode, Strategy, StreamingEncoder,
};
type Error = Box<dyn std::error::Error + Send + Sync>;
type Result<T, E = Error> = core::result::Result<T, E>;
macro_rules! eyre {
($($arg:tt)*) => {
<$crate::Error as From<String>>::from(format!($($arg)*))
};
}
macro_rules! bail {
($($arg:tt)*) => {
return Err(eyre!($($arg)*))
};
}
trait WrapErr<T> {
fn wrap_err(self, msg: impl core::fmt::Display) -> Result<T>;
fn wrap_err_with<D: core::fmt::Display, F: FnOnce() -> D>(self, msg: F) -> Result<T>;
}
impl<T, E: core::fmt::Display> WrapErr<T> for core::result::Result<T, E> {
fn wrap_err(self, msg: impl core::fmt::Display) -> Result<T> {
self.map_err(|source| eyre!("{msg}: {source}"))
}
fn wrap_err_with<D: core::fmt::Display, F: FnOnce() -> D>(self, msg: F) -> Result<T> {
self.map_err(|source| eyre!("{}: {source}", msg()))
}
}
macro_rules! display {
($verbosity:expr, $level:expr, $($arg:tt)*) => {
if $verbosity >= $level {
eprintln!($($arg)*);
}
};
}
mod display;
mod identity;
mod inputs;
mod interrupt;
mod priority;
mod progress;
mod synthetic;
mod trace;
use display::{DEFAULT_LEVEL, HumanSize, Progress, confirm};
use inputs::Selection;
use progress::ProgressMonitor;
const ZSTD_SUFFIX: &str = ".zst";
const DECOMPRESS_SUFFIXES: [(&str, &str); 3] = [("zst", ""), ("zstd", ""), ("tzst", "tar")];
const UPSTREAM_VERSION: &str = "1.5.7";
const UPSTREAM_VERSION_NUMBER: u32 = 10507;
const STDOUT_MARK: &str = "/*stdout*\\";
const CLI_MAX_LEVEL_WITHOUT_ULTRA: i32 = 19;
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Mode {
Compress,
Decompress,
Test,
List,
Train,
}
#[derive(Clone, Copy, Default)]
struct FakeConsole {
stdin: bool,
stdout: bool,
stderr: bool,
}
impl FakeConsole {
fn stdin(self) -> bool {
self.stdin || io::stdin().is_terminal()
}
fn stdout(self) -> bool {
self.stdout || io::stdout().is_terminal()
}
fn stderr(self) -> bool {
self.stderr || io::stderr().is_terminal()
}
}
struct Options {
mode: Mode,
level: i32,
dict: Option<PathBuf>,
to_stdout: bool,
output: Option<PathBuf>,
force: bool,
keep: bool,
remove_source: bool,
inputs: Vec<PathBuf>,
max_dict: usize,
dict_id: Option<u32>,
bench: bool,
bench_start: i32,
bench_end: i32,
bench_secs: f64,
bench_param: Option<u32>,
bench_compressibility: Option<u32>,
pause: bool,
bench_realtime: bool,
console: FakeConsole,
trace: Option<PathBuf>,
decode_window: u64,
patch_window_lifted: bool,
bench_separately: bool,
block_size: Option<u64>,
long: bool,
long_window_log: Option<u32>,
memory_limit: Option<u64>,
target_block_size: Option<u32>,
pledged_size: Option<u64>,
size_hint: Option<u64>,
verbosity: i32,
checksum: bool,
content_size_flag: bool,
dict_id_flag: bool,
pass_through: Option<bool>,
exclude_compressed: bool,
recursive: bool,
ignore_read_errors: bool,
follow_links: bool,
force_stdin: bool,
filelists: Vec<PathBuf>,
output_dir: Option<PathBuf>,
output_dir_mirror: Option<PathBuf>,
progress: Progress,
advanced: AdvancedParams,
literals: LiteralCompressionMode,
patch_from: Option<PathBuf>,
show_default_cparams: bool,
trainer: Trainer,
trainer_params: TrainerParams,
selectivity: u32,
cpu: Option<structured_zstd::CpuLevel>,
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Trainer {
FastCover,
Cover,
Legacy,
}
#[derive(Clone, Copy, Default, PartialEq, Eq, Debug)]
struct TrainerParams {
listed: bool,
k: Option<u32>,
d: Option<u32>,
f: Option<u32>,
steps: Option<u32>,
split_percent: Option<u32>,
accel: Option<u32>,
shrink: Option<u32>,
}
#[derive(Clone, Copy, Default, PartialEq, Eq, Debug)]
struct AdvancedParams {
window_log: Option<u32>,
chain_log: Option<u32>,
hash_log: Option<u32>,
search_log: Option<u32>,
min_match: Option<u32>,
target_length: Option<u32>,
strategy: Option<Strategy>,
ldm_hash_log: Option<u32>,
ldm_min_match: Option<u32>,
ldm_bucket_size_log: Option<u32>,
ldm_hash_rate_log: Option<u32>,
}
impl AdvancedParams {
fn is_default(&self) -> bool {
*self == Self::default()
}
}
const DEFAULT_MAX_DICT: usize = 112_640;
const DEFAULT_BENCH_SECONDS: f64 = 3.0;
const MIN_BENCH_BLOCK_SIZE: u64 = 32;
const DEFAULT_LONG_WINDOW_LOG: u32 = 27;
#[derive(Clone, Copy, PartialEq, Eq)]
enum SuffixCase {
Upper,
Any,
}
fn read_leading_number(text: &str, max: u64, case: SuffixCase) -> Result<(u64, &str)> {
let overflow = || {
eyre!(
"numeric value `{text}` overflows {} bits",
64 - max.leading_zeros()
)
};
let bytes = text.as_bytes();
let mut at = 0;
let mut value: u64 = 0;
while let Some(&digit) = bytes.get(at).filter(|b| b.is_ascii_digit()) {
value = value
.checked_mul(10)
.and_then(|v| v.checked_add(u64::from(digit - b'0')))
.filter(|&v| v <= max)
.ok_or_else(overflow)?;
at += 1;
}
let fold = |byte: &u8| match case {
SuffixCase::Upper => *byte,
SuffixCase::Any => byte.to_ascii_uppercase(),
};
let shift = match bytes.get(at).map(fold) {
Some(b'K') => 10,
Some(b'M') => 20,
Some(b'G') => 30,
_ => return Ok((value, &text[at..])),
};
value = value
.checked_mul(1 << shift)
.filter(|&v| v <= max)
.ok_or_else(overflow)?;
at += 1;
at += usize::from(bytes.get(at) == Some(&b'i'));
at += usize::from(bytes.get(at).map(fold) == Some(b'B'));
Ok((value, &text[at..]))
}
fn parse_number(text: &str, max: u64) -> Result<u64> {
if !text.starts_with(|c: char| c.is_ascii_digit()) {
bail!("`{text}`: expected a number");
}
let (value, rest) = read_leading_number(text, max, SuffixCase::Any)?;
if !rest.is_empty() {
bail!(
"`{text}`: only numeric values with optional suffixes K, KB, KiB, M, MB, MiB, G, GB, GiB are allowed"
);
}
Ok(value)
}
fn parse_size(text: &str) -> Result<u64> {
parse_number(text, usize::MAX as u64)
}
fn parse_u32(text: &str) -> Result<u32> {
let value = parse_number(text, u64::from(u32::MAX))?;
Ok(u32::try_from(value).expect("parse_number caps the value at u32::MAX"))
}
fn cluster_u32(chars: &[char], at: usize) -> Result<(u32, usize)> {
let rest: String = chars[at..].iter().collect();
let (value, tail) = read_leading_u32(&rest)?;
Ok((value, chars.len() - tail.chars().count()))
}
fn read_leading_u32(text: &str) -> Result<(u32, &str)> {
let (value, rest) = read_leading_number(text, u64::from(u32::MAX), SuffixCase::Upper)?;
Ok((
u32::try_from(value).expect("read_leading_number caps the value at u32::MAX"),
rest,
))
}
fn parse_memory_limit(text: &str) -> Result<Option<u64>> {
let value = parse_size(text)?;
let bytes = if text.trim().ends_with(|c: char| c.is_ascii_digit()) {
value
.checked_mul(1 << 20)
.ok_or_else(|| eyre!("memory limit `{text}` MiB does not fit in 64 bits"))?
} else {
value
};
Ok((bytes != 0).then_some(bytes))
}
const DECODER_AUXILIARY_BYTES: u64 = 1 << 20;
fn decode_window_ceiling(opts: &Options) -> u64 {
use structured_zstd::decoding::{MAX_DECODER_WINDOW_SIZE, MAXIMUM_ALLOWED_WINDOW_SIZE};
if let Some(limit) = opts.memory_limit {
return limit.clamp(MAXIMUM_ALLOWED_WINDOW_SIZE, MAX_DECODER_WINDOW_SIZE);
}
if opts.patch_window_lifted {
return MAX_DECODER_WINDOW_SIZE;
}
match opts.advanced.window_log.or(opts.long_window_log) {
Some(log) => (1u64 << log).clamp(MAXIMUM_ALLOWED_WINDOW_SIZE, MAX_DECODER_WINDOW_SIZE),
None => MAXIMUM_ALLOWED_WINDOW_SIZE,
}
}
fn check_memory_limit(requested: u64, whole_file_bytes: u64, copies: u64) -> Result<()> {
const WINDOW: u64 = structured_zstd::decoding::MAXIMUM_ALLOWED_WINDOW_SIZE;
const AUXILIARY: u64 = DECODER_AUXILIARY_BYTES;
let buffers = whole_file_bytes
.checked_mul(copies)
.and_then(|buffers| buffers.checked_add(WINDOW + AUXILIARY));
let Some(floor) = buffers else {
bail!(
"requested memory limit {requested} B cannot cover this run: the files it \
would hold add up to more than any machine can address."
);
};
if requested < floor {
if floor > WINDOW + AUXILIARY {
let held = if copies == 1 {
String::new()
} else {
format!(", each held in {copies} forms at once")
};
bail!(
"requested memory limit {requested} B does not cover this run: a {} MiB \
window, about {} MiB of literal, block, sequence and table buffers, and \
{} MiB of whole-file buffers{held}.",
WINDOW / (1 << 20),
AUXILIARY / (1 << 20),
(floor - WINDOW - AUXILIARY) / (1 << 20),
);
}
bail!(
"requested memory limit {requested} B is below what decompression can need \
here: a {} MiB window plus about {} MiB of literal, block, sequence and \
table buffers. This build cannot be tightened below that.",
WINDOW / (1 << 20),
AUXILIARY / (1 << 20),
);
}
Ok(())
}
fn check_window_log(log: u32, compressing: bool) -> Result<()> {
use structured_zstd::encoding::CParameter;
let bounds = CParameter::WindowLog.bounds();
let upper = i64::from(if compressing {
max_window_log()
} else {
structured_zstd::decoding::MAX_DECODER_WINDOW_SIZE.ilog2()
});
if i64::from(log) < bounds.lower_bound || i64::from(log) > upper {
bail!(
"window log {log} is outside the supported range {}..={upper}",
bounds.lower_bound,
);
}
Ok(())
}
fn max_window_log() -> u32 {
use structured_zstd::encoding::CParameter;
let encodable = u32::try_from(CParameter::WindowLog.bounds().upper_bound)
.expect("the window-log bound is a small positive number");
let decodable = structured_zstd::decoding::MAX_DECODER_WINDOW_SIZE.ilog2();
encodable.min(decodable)
}
fn max_compression_params() -> AdvancedParams {
use structured_zstd::encoding::CParameter;
let upper = |parameter: CParameter| {
u32::try_from(parameter.bounds().upper_bound)
.expect("every compression-parameter bound is a small positive number")
};
let lower = |parameter: CParameter| {
u32::try_from(parameter.bounds().lower_bound)
.expect("every compression-parameter bound is a small positive number")
};
AdvancedParams {
window_log: Some(max_window_log()),
chain_log: Some(upper(CParameter::ChainLog)),
hash_log: Some(upper(CParameter::HashLog)),
search_log: Some(upper(CParameter::SearchLog)),
min_match: Some(lower(CParameter::MinMatch)),
target_length: Some(upper(CParameter::TargetLength)),
strategy: Some(Strategy::Btultra2),
ldm_hash_log: Some(upper(CParameter::LdmHashLog)),
ldm_min_match: Some(16),
ldm_bucket_size_log: Some(upper(CParameter::LdmBucketSizeLog)),
ldm_hash_rate_log: None,
}
}
fn parse_adapt_params(params: &str) -> Result<()> {
if params.is_empty() {
bail!("--adapt= needs parameters, e.g. --adapt=min=1,max=9");
}
for field in params.split(',') {
let (key, value) = field
.split_once('=')
.ok_or_else(|| eyre!("--adapt parameter `{field}` is not `key=value`"))?;
if key != "min" && key != "max" {
bail!("--adapt has no `{key}` parameter; expected `min` or `max`");
}
value
.parse::<i32>()
.map_err(|_| eyre!("--adapt {key} must be a number, got `{value}`"))?;
}
Ok(())
}
fn apply_advanced_params(text: &str, params: &mut AdvancedParams) -> Result<()> {
if text.is_empty() {
return Ok(());
}
for field in text.split(',') {
let (key, value) = field
.split_once('=')
.ok_or_else(|| eyre!("--zstd parameter `{field}` is not `key=value`"))?;
let (number, tail) = read_leading_u32(value)
.wrap_err_with(|| format!("--zstd parameter `{key}` has an invalid value"))?;
if !tail.is_empty() {
bail!("--zstd parameter `{key}` has an invalid value `{value}`");
}
let set = (number != 0).then_some(number);
match key {
"windowLog" | "wlog" => params.window_log = set,
"chainLog" | "clog" => params.chain_log = set,
"hashLog" | "hlog" => params.hash_log = set,
"searchLog" | "slog" => params.search_log = set,
"minMatch" | "mml" => params.min_match = set,
"targetLength" | "tlen" => params.target_length = set,
"strategy" | "strat" => {
params.strategy = match set {
None => None,
Some(ordinal) => {
Some(Strategy::from_ordinal(ordinal).ok_or_else(|| {
eyre!("--zstd strategy {ordinal} is out of range 1..=9")
})?)
}
}
}
"overlapLog" | "ovlog" => {}
"ldmHashLog" | "lhlog" => params.ldm_hash_log = set,
"ldmMinMatch" | "lmml" => params.ldm_min_match = set,
"ldmBucketSizeLog" | "lblog" => params.ldm_bucket_size_log = set,
"ldmHashRateLog" | "lhrlog" => params.ldm_hash_rate_log = set,
_ => bail!("--zstd has no `{key}` parameter"),
}
}
Ok(())
}
#[cfg(test)]
fn parse_advanced_params(text: &str) -> Result<AdvancedParams> {
let mut params = AdvancedParams::default();
apply_advanced_params(text, &mut params)?;
Ok(params)
}
fn parse_trainer_params(text: &str, fastcover: bool) -> Result<TrainerParams> {
let flag = if fastcover {
"--train-fastcover"
} else {
"--train-cover"
};
let mut params = TrainerParams {
listed: true,
..TrainerParams::default()
};
for field in text.split(',') {
if field == "shrink" || field.starts_with("shrink=") {
let mut bound = 1;
if let Some(text) = field.strip_prefix("shrink=") {
let (value, tail) = read_leading_u32(text)
.wrap_err_with(|| format!("{flag} shrink bound `{text}` is invalid"))?;
if !tail.is_empty() {
bail!("{flag} shrink bound `{text}` is invalid");
}
bound = value;
}
params.shrink = Some(bound);
continue;
}
let (key, value) = field
.split_once('=')
.ok_or_else(|| eyre!("{flag} parameter `{field}` is not `key=value`"))?;
let (number, tail) = read_leading_u32(value)
.wrap_err_with(|| format!("{flag} parameter `{key}` has an invalid value"))?;
if !tail.is_empty() {
bail!("{flag} parameter `{key}` has an invalid value `{value}`");
}
let set = (number != 0).then_some(number);
match key {
"k" => params.k = set,
"d" => params.d = set,
"steps" => params.steps = set,
"split" => params.split_percent = set,
"f" if fastcover => params.f = set,
"accel" if fastcover => params.accel = set,
_ => bail!("{flag} has no `{key}` parameter"),
}
}
Ok(params)
}
fn refuse_console_stdout(opts: &Options, stdout_is_terminal: bool) -> Result<()> {
if opts.mode == Mode::Compress
&& !opts.bench
&& opts.output.is_none()
&& !opts.to_stdout
&& reads_stdin(&opts.inputs)
&& stdout_is_terminal
&& !opts.force
{
bail!("stdout is a console, aborting");
}
Ok(())
}
enum Parsed {
Run(Box<Options>),
Handled { pause: bool },
}
struct ParseFailure {
error: Error,
verbosity: i32,
pause: bool,
}
fn wait_for_enter(stdin: &mut impl Read) {
eprintln!("Press enter to continue... ");
let mut byte = [0u8; 1];
drop(stdin.read(&mut byte));
}
fn main() {
let raw: Vec<OsString> = std::env::args_os().collect();
std::process::exit(run_command(&raw, &mut io::stdin()));
}
fn run_command(raw: &[OsString], stdin: &mut impl Read) -> i32 {
let mut pause = false;
let status = command_status(raw, &mut pause);
if pause {
wait_for_enter(stdin);
}
status
}
fn command_status(raw: &[OsString], pause: &mut bool) -> i32 {
let prog = raw
.first()
.map(|arg| arg.to_string_lossy().into_owned())
.unwrap_or_else(|| "zstd".to_string());
let preset = program_preset(&prog);
let default_level = level_from_env(std::env::var_os("ZSTD_CLEVEL").as_deref(), DEFAULT_LEVEL);
check_threads_env(std::env::var_os("ZSTD_NBTHREADS").as_deref(), DEFAULT_LEVEL);
let args = raw.get(1..).unwrap_or_default();
let options = match parse_args(args, &preset, default_level) {
Ok(Parsed::Run(options)) => *options,
Ok(Parsed::Handled { pause: asked }) => {
*pause = asked;
return 0;
}
Err(failure) => {
*pause = failure.pause;
display!(failure.verbosity, 1, "zstd: {}", failure.error);
if failure.verbosity >= DEFAULT_LEVEL {
let mut stderr = io::stderr().lock();
let _ = write_short_usage(&mut stderr, &preset.name);
}
return 1;
}
};
let verbosity = options.verbosity;
*pause = options.pause;
if let Some(level) = options.cpu
&& let Err(err) = structured_zstd::set_cpu_ceiling(level)
{
display!(verbosity, 1, "zstd: --cpu={level}: {err}");
return 1;
}
display!(
verbosity,
4,
"kernels: entropy {}, match finding {}",
structured_zstd::active_cpu_kernel_name(),
structured_zstd::active_match_kernel_name()
);
match run(options) {
Ok(0) => 0,
Ok(_failed_inputs) => 1,
Err(err) => {
display!(verbosity, 1, "zstd: {err}");
1
}
}
}
struct ProgramPreset {
name: String,
mode: Mode,
to_stdout: bool,
force: bool,
pass_through: Option<bool>,
verbosity: i32,
}
fn program_preset(prog: &str) -> ProgramPreset {
let name = Path::new(prog)
.file_name()
.and_then(|n| n.to_str())
.unwrap_or(prog);
let stem = name.strip_suffix(".exe").unwrap_or(name);
let plain = ProgramPreset {
name: name.to_string(),
mode: Mode::Compress,
to_stdout: false,
force: false,
pass_through: None,
verbosity: DEFAULT_LEVEL,
};
match stem {
"unzstd" => ProgramPreset {
mode: Mode::Decompress,
..plain
},
"zstdcat" | "zcat" => ProgramPreset {
mode: Mode::Decompress,
to_stdout: true,
force: true,
pass_through: Some(true),
verbosity: 1,
..plain
},
_ => plain,
}
}
fn level_from_env(value: Option<&OsStr>, verbosity: i32) -> i32 {
let Some(value) = value else {
return CompressionLevel::DEFAULT_LEVEL;
};
let text = value.to_string_lossy();
let (sign, digits) = match text.strip_prefix('-') {
Some(rest) => (-1, rest),
None => (1, text.strip_prefix('+').unwrap_or(&text)),
};
if digits.starts_with(|c: char| c.is_ascii_digit()) {
match read_leading_u32(digits) {
Ok((magnitude, "")) => {
return i32::try_from(magnitude)
.map(|magnitude| sign * magnitude)
.unwrap_or(if sign < 0 { i32::MIN } else { i32::MAX })
.max(CompressionLevel::MIN_LEVEL);
}
Err(_) => {
display!(
verbosity,
2,
"Ignore environment variable setting ZSTD_CLEVEL={text}: numeric value too large"
);
return CompressionLevel::DEFAULT_LEVEL;
}
Ok(_) => {}
}
}
display!(
verbosity,
2,
"Ignore environment variable setting ZSTD_CLEVEL={text}: not a valid integer value"
);
CompressionLevel::DEFAULT_LEVEL
}
fn check_threads_env(value: Option<&OsStr>, verbosity: i32) {
let Some(value) = value else {
return;
};
let text = value.to_string_lossy();
if text.starts_with(|c: char| c.is_ascii_digit()) {
match read_leading_u32(&text) {
Ok((_, "")) => return,
Err(_) => {
display!(
verbosity,
2,
"Ignore environment variable setting ZSTD_NBTHREADS={text}: numeric value too large"
);
return;
}
Ok(_) => {}
}
}
display!(
verbosity,
2,
"Ignore environment variable setting ZSTD_NBTHREADS={text}: not a valid unsigned value"
);
}
fn option_value<'a>(
long: &str,
name: &str,
arg_os: &OsStr,
rest: &mut impl Iterator<Item = (usize, &'a OsString)>,
) -> Result<Option<PathBuf>> {
if long == name {
let Some((_, value)) = rest.next() else {
bail!("error: missing command argument for --{name}");
};
if value.to_string_lossy().starts_with('-') {
bail!("error: command cannot be separated from its argument by another command");
}
return Ok(Some(PathBuf::from(value)));
}
if long.len() > name.len() && long.starts_with(name) && long.as_bytes()[name.len()] == b'=' {
return Ok(Some(attached_path(arg_os, 2 + name.len() + 1)));
}
Ok(None)
}
fn option_text<'a>(
long: &str,
name: &str,
arg_os: &OsStr,
rest: &mut impl Iterator<Item = (usize, &'a OsString)>,
) -> Result<Option<String>> {
Ok(option_value(long, name, arg_os, rest)?
.map(|value| value.as_os_str().to_string_lossy().into_owned()))
}
fn first_option_text<'a>(
long: &str,
names: &[&str],
arg_os: &OsStr,
rest: &mut impl Iterator<Item = (usize, &'a OsString)>,
) -> Result<Option<String>> {
for name in names {
if let Some(value) = option_text(long, name, arg_os, rest)? {
return Ok(Some(value));
}
}
Ok(None)
}
fn parse_args(
args: &[OsString],
preset: &ProgramPreset,
default_level: i32,
) -> Result<Parsed, ParseFailure> {
let mut verbosity = preset.verbosity;
let mut pause = false;
parse_args_into(args, preset, default_level, &mut verbosity, &mut pause).map_err(|error| {
ParseFailure {
error,
verbosity,
pause,
}
})
}
fn parse_args_into(
args: &[OsString],
preset: &ProgramPreset,
default_level: i32,
verbosity: &mut i32,
pause: &mut bool,
) -> Result<Parsed> {
let mut opts = Options {
mode: preset.mode,
level: default_level,
dict: None,
to_stdout: preset.to_stdout,
output: None,
force: preset.force,
keep: false,
remove_source: false,
inputs: Vec::new(),
max_dict: DEFAULT_MAX_DICT,
dict_id: None,
bench: false,
bench_start: default_level,
bench_end: default_level,
bench_secs: DEFAULT_BENCH_SECONDS,
bench_param: None,
bench_compressibility: None,
pause: false,
bench_realtime: false,
console: FakeConsole::default(),
trace: None,
decode_window: structured_zstd::decoding::MAXIMUM_ALLOWED_WINDOW_SIZE,
patch_window_lifted: false,
bench_separately: false,
block_size: None,
long: false,
long_window_log: None,
memory_limit: None,
target_block_size: None,
pledged_size: None,
size_hint: None,
verbosity: preset.verbosity,
checksum: true,
content_size_flag: true,
dict_id_flag: true,
pass_through: preset.pass_through,
exclude_compressed: false,
recursive: false,
ignore_read_errors: false,
follow_links: preset.force,
force_stdin: preset.force,
filelists: Vec::new(),
output_dir: None,
output_dir_mirror: None,
progress: Progress::Auto,
advanced: AdvancedParams::default(),
literals: LiteralCompressionMode::Auto,
patch_from: None,
show_default_cparams: false,
trainer: Trainer::FastCover,
trainer_params: TrainerParams::default(),
selectivity: 0,
cpu: None,
};
let mut ultra = false;
let mut bench_end = None;
let mut iter = args.iter().enumerate().peekable();
let mut positional_only = false;
while let Some((idx, arg_os)) = iter.next() {
let arg = arg_os.to_string_lossy();
let arg = arg.as_ref();
if positional_only || arg == "-" || !arg.starts_with('-') {
opts.inputs.push(PathBuf::from(arg_os));
continue;
}
if arg == "--" {
positional_only = true;
continue;
}
if let Some(long) = arg.strip_prefix("--") {
match long {
"compress" => select_mode(&mut opts, Mode::Compress),
"decompress" | "uncompress" => select_mode(&mut opts, Mode::Decompress),
"test" => select_mode(&mut opts, Mode::Test),
"list" => select_mode(&mut opts, Mode::List),
"train" => select_mode(&mut opts, Mode::Train),
"train-fastcover" => {
select_mode(&mut opts, Mode::Train);
opts.trainer = Trainer::FastCover;
opts.trainer_params = TrainerParams::default();
}
"train-cover" => {
select_mode(&mut opts, Mode::Train);
opts.trainer = Trainer::Cover;
opts.trainer_params = TrainerParams::default();
}
"train-legacy" => {
select_mode(&mut opts, Mode::Train);
opts.trainer = Trainer::Legacy;
}
"stdout" | "to-stdout" => {
opts.to_stdout = true;
opts.output = None;
}
"force" => {
opts.force = true;
opts.force_stdin = true;
opts.follow_links = true;
}
"keep" => opts.keep = true,
"rm" => opts.remove_source = true,
"ultra" => ultra = true,
"max" => {
if usize::BITS < 64 {
bail!("--max is incompatible with 32-bit mode");
}
ultra = true;
opts.long = true;
opts.advanced = max_compression_params();
}
"quiet" => *verbosity -= 1,
"verbose" => *verbosity += 1,
"check" => opts.checksum = true,
"no-check" => opts.checksum = false,
"content-size" => opts.content_size_flag = true,
"no-content-size" => opts.content_size_flag = false,
"no-dictID" => opts.dict_id_flag = false,
"pass-through" => opts.pass_through = Some(true),
"no-pass-through" => opts.pass_through = Some(false),
"exclude-compressed" => opts.exclude_compressed = true,
"ignore-read-errors" => opts.ignore_read_errors = true,
"show-default-cparams" => opts.show_default_cparams = true,
"progress" => opts.progress = Progress::Always,
"no-progress" => opts.progress = Progress::Never,
"version" => {
print_version(*verbosity);
return Ok(Parsed::Handled { pause: *pause });
}
"help" => {
print_help(*verbosity, &preset.name);
return Ok(Parsed::Handled { pause: *pause });
}
"single-thread"
| "adapt"
| "sparse"
| "no-sparse"
| "asyncio"
| "no-asyncio"
| "mmap-dict"
| "no-mmap-dict"
| "row-match-finder"
| "no-row-match-finder" => {}
"trace-file-stat" => bail!(
"--trace-file-stat traces the reference implementation's own file-system \
calls; this build has none to trace"
),
"priority=rt" => opts.bench_realtime = true,
"fake-stdin-is-console" => opts.console.stdin = true,
"fake-stdout-is-console" => opts.console.stdout = true,
"fake-stderr-is-console" => opts.console.stderr = true,
"compress-literals" => opts.literals = LiteralCompressionMode::Enable,
"no-compress-literals" => opts.literals = LiteralCompressionMode::Disable,
_ => {
if long == "fast" {
opts.level = -1;
} else if let Some(v) = long.strip_prefix("fast=") {
let (n, _tail) = read_leading_u32(v).wrap_err("invalid --fast level")?;
if n == 0 {
bail!("--fast level must be at least 1, got 0");
}
let capped = n.min(CompressionLevel::MIN_LEVEL.unsigned_abs());
opts.level = -i32::try_from(capped)
.expect("capped at |MIN_LEVEL|, which is an i32 magnitude");
} else if long.starts_with("use-dict=") {
opts.dict = Some(attached_path(arg_os, "--use-dict=".len()));
} else if let Some(v) = option_text(long, "maxdict", arg_os, &mut iter)? {
opts.max_dict =
usize::try_from(parse_u32(&v).wrap_err("invalid --maxdict size")?)
.expect("a u32 fits in usize on the platforms this builds for");
} else if let Some(v) = option_text(long, "dictID", arg_os, &mut iter)? {
let id = parse_u32(&v).wrap_err("invalid --dictID")?;
opts.dict_id = (id != 0).then_some(id);
} else if let Some(v) = option_text(long, "stream-size", arg_os, &mut iter)? {
opts.pledged_size = Some(parse_size(&v).wrap_err("invalid --stream-size")?);
} else if let Some(v) = option_text(long, "size-hint", arg_os, &mut iter)? {
opts.size_hint = Some(parse_size(&v).wrap_err("invalid --size-hint")?);
} else if let Some(v) = first_option_text(
long,
&["memory", "memlimit", "memlimit-decompress"],
arg_os,
&mut iter,
)? {
opts.memory_limit =
parse_memory_limit(&v).wrap_err("invalid memory limit")?;
opts.patch_window_lifted = false;
} else if let Some(params) = long.strip_prefix("adapt=") {
parse_adapt_params(params)?;
} else if let Some(v) = option_text(long, "auto-threads", arg_os, &mut iter)? {
if v != "physical" && v != "logical" {
bail!("--auto-threads must be `physical` or `logical`, got `{v}`");
}
} else if let Some(v) =
option_text(long, "target-compressed-block-size", arg_os, &mut iter)?
{
let target =
parse_size(&v).wrap_err("invalid --target-compressed-block-size")?;
opts.target_block_size = Some(u32::try_from(target).map_err(|_| {
eyre!("--target-compressed-block-size={v} is too large")
})?);
} else if let Some(v) = option_text(long, "threads", arg_os, &mut iter)? {
parse_u32(&v).wrap_err("invalid --threads")?;
} else if let Some(v) = first_option_text(
long,
&["block-size", "split", "jobsize"],
arg_os,
&mut iter,
)? {
let size = parse_size(&v).wrap_err("invalid --block-size")?;
opts.block_size = (size != 0).then_some(size);
} else if let Some(list) = option_value(long, "filelist", arg_os, &mut iter)? {
opts.filelists.push(list);
} else if let Some(dir) =
option_value(long, "output-dir-flat", arg_os, &mut iter)?
{
if dir.as_os_str().is_empty() {
bail!(
"error: output dir cannot be empty string (did you mean to pass '.' instead?)"
);
}
opts.output_dir = Some(dir);
} else if let Some(dir) =
option_value(long, "output-dir-mirror", arg_os, &mut iter)?
{
if dir.as_os_str().is_empty() {
bail!(
"error: output dir cannot be empty string (did you mean to pass '.' instead?)"
);
}
opts.output_dir_mirror = Some(dir);
} else if let Some(v) = long.strip_prefix("format=") {
if v != "zstd" {
bail!("--format={v} is not supported; this build only writes zstd");
}
} else if let Some(v) = long.strip_prefix("zstd=") {
apply_advanced_params(v, &mut opts.advanced)?;
} else if let Some(v) = long.strip_prefix("cpu=") {
opts.cpu = Some(v.parse().map_err(|err| eyre!("--cpu={v}: {err}"))?);
} else if let Some(v) = long.strip_prefix("train-cover=") {
select_mode(&mut opts, Mode::Train);
opts.trainer = Trainer::Cover;
opts.trainer_params = parse_trainer_params(v, false)?;
} else if let Some(v) = long.strip_prefix("train-fastcover=") {
select_mode(&mut opts, Mode::Train);
opts.trainer = Trainer::FastCover;
opts.trainer_params = parse_trainer_params(v, true)?;
} else if let Some(v) = long.strip_prefix("train-legacy=") {
select_mode(&mut opts, Mode::Train);
opts.trainer = Trainer::Legacy;
let value = v
.strip_prefix("selectivity=")
.or_else(|| v.strip_prefix("s="))
.ok_or_else(|| {
eyre!("--train-legacy takes `s=#` or `selectivity=#`, got `{v}`")
})?;
let (selectivity, tail) = read_leading_u32(value)
.wrap_err("invalid --train-legacy selectivity")?;
if !tail.is_empty() {
bail!("invalid --train-legacy selectivity `{value}`");
}
opts.selectivity = selectivity;
} else if let Some(reference) =
option_value(long, "patch-from", arg_os, &mut iter)?
{
opts.patch_from = Some(reference);
opts.patch_window_lifted = true;
ultra = true;
} else if let Some(reference) =
option_value(long, "patch-apply", arg_os, &mut iter)?
{
select_mode(&mut opts, Mode::Decompress);
opts.patch_from = Some(reference);
opts.patch_window_lifted = true;
opts.memory_limit = None;
} else if let Some(path) = option_value(long, "trace", arg_os, &mut iter)? {
opts.trace = Some(path);
} else if long == "rsyncable" {
bail!("--rsyncable is not compatible with single-thread mode");
} else if long == "long" {
opts.long = true;
opts.long_window_log = Some(DEFAULT_LONG_WINDOW_LOG);
} else if let Some(v) = long.strip_prefix("long=") {
let log: u32 = v.parse().wrap_err("invalid --long window log")?;
opts.long = true;
opts.long_window_log = Some(log);
} else {
bail!("unknown option: --{long}");
}
}
}
continue;
}
let chars: Vec<char> = arg[1..].chars().collect();
let mut ci = 0;
while ci < chars.len() {
let c = chars[ci];
match c {
'd' => select_mode(&mut opts, Mode::Decompress),
'z' => select_mode(&mut opts, Mode::Compress),
't' => select_mode(&mut opts, Mode::Test),
'l' => select_mode(&mut opts, Mode::List),
'b' => opts.bench = true,
'p' if !chars.get(ci + 1).is_some_and(char::is_ascii_digit) => *pause = true,
'e' | 'i' | 'p' | 'P' | 'B' | 'T' | 's' => {
let (value, next) = cluster_u32(&chars, ci + 1)
.wrap_err_with(|| format!("invalid -{c} value (in {arg})"))?;
match c {
'e' => {
bench_end = Some(
i32::try_from(value)
.wrap_err_with(|| format!("invalid -e level `{value}`"))?,
);
}
'i' => opts.bench_secs = f64::from(value),
'p' => opts.bench_param = Some(value),
'P' => opts.bench_compressibility = Some(value),
'B' => opts.block_size = (value != 0).then_some(u64::from(value)),
'T' => {}
_ => {
debug_assert_eq!(c, 's');
opts.selectivity = value;
}
}
ci = next;
continue;
}
'n' => {}
'c' => {
opts.to_stdout = true;
opts.output = None;
}
'f' => {
opts.force = true;
opts.force_stdin = true;
opts.follow_links = true;
}
'k' => opts.keep = true,
'S' => opts.bench_separately = true,
'q' => *verbosity -= 1,
'v' => *verbosity += 1,
'C' => opts.checksum = true,
'r' => opts.recursive = true,
'M' => {
let (_, next) = cluster_u32(&chars, ci + 1)
.wrap_err_with(|| format!("invalid -M memory limit (in {arg})"))?;
let number: String = chars[ci + 1..next].iter().collect();
if !number.is_empty() {
opts.memory_limit =
parse_memory_limit(&number).wrap_err("invalid -M memory limit")?;
opts.patch_window_lifted = false;
}
ci = next;
continue;
}
'V' => {
print_version(*verbosity);
return Ok(Parsed::Handled { pause: *pause });
}
'H' => {
print_help(*verbosity, &preset.name);
return Ok(Parsed::Handled { pause: *pause });
}
'h' => {
let mut stdout = io::stdout().lock();
write_short_usage(&mut stdout, &preset.name)
.wrap_err("failed to write usage")?;
return Ok(Parsed::Handled { pause: *pause });
}
'D' | 'o' => {
let value = if ci + 1 == chars.len() {
iter.next()
.map(|(_, v)| PathBuf::from(v))
.ok_or_else(|| eyre!("option -{c} requires a value"))?
} else {
attached_path(arg_os, ci + 2)
};
if c == 'D' {
opts.dict = Some(value);
} else {
opts.output = Some(value);
opts.to_stdout = false;
}
ci = chars.len();
continue;
}
'0'..='9' => {
let (level, next) = cluster_u32(&chars, ci)
.wrap_err_with(|| format!("invalid level (in {arg})"))?;
opts.level = i32::try_from(level)
.wrap_err_with(|| format!("invalid level `{level}`"))?;
ci = next;
continue;
}
_ => bail!("unknown flag: -{c} (in {arg})"),
}
ci += 1;
}
let _ = idx;
}
opts.verbosity = *verbosity;
opts.pause = *pause;
if let Some(limit) = opts.memory_limit
&& decodes(&opts)
{
check_memory_limit(limit, 0, 0)?;
}
validate_level(opts.level)?;
opts.bench_start = opts.level;
opts.bench_end = bench_end.map_or(opts.bench_start, |end: i32| end.max(opts.bench_start));
let (lowest_level, highest_level) = if opts.bench {
(
opts.bench_start.min(opts.bench_end),
opts.bench_start.max(opts.bench_end),
)
} else {
(opts.level, opts.level)
};
validate_level(lowest_level)?;
validate_level(highest_level)?;
if !ultra && highest_level > CLI_MAX_LEVEL_WITHOUT_ULTRA {
display!(
*verbosity,
2,
"Warning : compression level higher than max, reduced to {CLI_MAX_LEVEL_WITHOUT_ULTRA}"
);
opts.level = opts.level.min(CLI_MAX_LEVEL_WITHOUT_ULTRA);
opts.bench_start = opts.bench_start.min(CLI_MAX_LEVEL_WITHOUT_ULTRA);
opts.bench_end = opts.bench_end.min(CLI_MAX_LEVEL_WITHOUT_ULTRA);
}
for log in [opts.long_window_log, opts.advanced.window_log]
.into_iter()
.flatten()
{
check_window_log(log, compresses(&opts))?;
}
if compresses(&opts) {
frame_parameters(
CompressionLevel::from_level(opts.level),
&FrameSettings::from_options(&opts),
)?;
}
if opts.patch_from.is_some() && opts.dict.is_some() {
bail!("error : can't use -D and --patch-from=# at the same time");
}
Ok(Parsed::Run(Box::new(opts)))
}
fn attached_path(arg: &std::ffi::OsStr, at: usize) -> PathBuf {
let bytes = arg.as_encoded_bytes();
debug_assert!(
bytes[..at].is_ascii(),
"the split point must follow ASCII, or it is not a boundary"
);
let rest = unsafe { std::ffi::OsStr::from_encoded_bytes_unchecked(&bytes[at..]) };
PathBuf::from(rest)
}
fn select_mode(opts: &mut Options, mode: Mode) {
opts.mode = mode;
opts.bench = false;
}
fn decodes(opts: &Options) -> bool {
opts.bench || matches!(opts.mode, Mode::Decompress | Mode::Test)
}
fn compresses(opts: &Options) -> bool {
opts.bench || matches!(opts.mode, Mode::Compress)
}
fn validate_level(level: i32) -> Result<()> {
let (min, max) = (CompressionLevel::MIN_LEVEL, CompressionLevel::MAX_LEVEL);
if !(min..=max).contains(&level) {
bail!("compression level {level} out of range [{min}, {max}]");
}
Ok(())
}
fn print_version(verbosity: i32) {
if verbosity < DEFAULT_LEVEL {
println!("{UPSTREAM_VERSION}");
return;
}
println!(
"zstd version {UPSTREAM_VERSION} (structured-zstd v{})",
env!("CARGO_PKG_VERSION")
);
if verbosity >= 3 {
println!("*** supports: zstd");
}
}
fn write_short_usage(out: &mut impl Write, program: &str) -> io::Result<()> {
writeln!(
out,
"Compress or decompress the INPUT file(s); reads from STDIN if INPUT is `-` or not provided.\n"
)?;
writeln!(
out,
"Usage: {program} [OPTIONS...] [INPUT... | -] [-o OUTPUT]\n"
)?;
writeln!(out, "Options:")?;
writeln!(
out,
" -o OUTPUT Write output to a single file, OUTPUT."
)?;
writeln!(
out,
" -k, --keep Preserve INPUT file(s). [Default]"
)?;
writeln!(
out,
" --rm Remove INPUT file(s) after successful (de)compression.\n"
)?;
writeln!(
out,
" -# Desired compression level, where `#` is a number between 1 and {CLI_MAX_LEVEL_WITHOUT_ULTRA};"
)?;
writeln!(
out,
" lower numbers provide faster compression, higher numbers yield"
)?;
writeln!(
out,
" better compression ratios. [Default: {}]\n",
CompressionLevel::DEFAULT_LEVEL
)?;
writeln!(
out,
" -d, --decompress Perform decompression."
)?;
writeln!(
out,
" -D DICT Use DICT as the dictionary for compression or decompression.\n"
)?;
writeln!(
out,
" -f, --force Disable input and output checks. Allows overwriting existing files,"
)?;
writeln!(
out,
" receiving input from the console, printing output to STDOUT, and"
)?;
writeln!(
out,
" operating on links, block devices, etc. Unrecognized formats will be"
)?;
writeln!(
out,
" passed through as-is.\n"
)?;
writeln!(
out,
" -h Display short usage and exit."
)?;
writeln!(
out,
" -H, --help Display full help and exit."
)?;
writeln!(
out,
" -V, --version Display the program version and exit.\n"
)
}
fn print_help(verbosity: i32, program: &str) {
print_version(verbosity.max(DEFAULT_LEVEL));
println!();
let mut stdout = io::stdout().lock();
let _ = write_short_usage(&mut stdout, program);
let _ = stdout.write_all(HELP_ADVANCED.as_bytes());
}
const HELP_ADVANCED: &str = "\
Advanced options:
-c, --stdout Write to STDOUT (even if it is a console) and keep the INPUT file(s).
-v, --verbose Enable verbose output; pass multiple times to increase verbosity.
-q, --quiet Suppress warnings; pass twice to suppress errors.
--trace LOG Log tracing information to LOG.
--[no-]progress Forcibly show/hide the progress counter. NOTE: Any (de)compressed
output to terminal will mix with progress counter text.
-r Operate recursively on directories.
--ignore-read-errors With -r, report a directory that cannot be opened and go on
without failing the run (exit 0, as upstream zstd).
--filelist LIST Read a list of files to operate on from LIST.
--output-dir-flat DIR Store processed files in DIR.
--output-dir-mirror DIR Store processed files in DIR, respecting original directory structure.
--[no-]check Add XXH64 integrity checksums during compression. [Default: Add, Validate]
If `-d` is present, ignore/validate checksums during decompression.
-- Treat remaining arguments after `--` as files.
--cpu=LEVEL Use kernels up to LEVEL only: scalar, sse2, sse4.2, bmi2, avx2, avx512
(x86) or neon, sve (aarch64). [Default: all the CPU supports]
Advanced compression options:
--ultra Enable levels beyond 19, up to 22; requires more memory.
--max Compress with every parameter at its maximum; the window stops at 30,
the widest this build writes. Requires a lot of memory.
--fast[=#] Use to very fast compression levels. [Default: 1]
--long[=#] Enable long distance matching with window log #. [Default: 27]
Available from level 16 up (or with --zstd=strat=7..9), where
long-distance matching runs; up to 30 when compressing. When
decompressing, admits windows up to 2^#.
--patch-from=REF Use REF as the reference point for Zstandard's diff engine.
--zstd=wlog=#,clog=#,hlog=#,slog=#,mml=#,tlen=#,strat=#[,lhlog=#,lmml=#,lblog=#,lhrlog=#]
Override the level's compression parameters knob by knob.
--exclude-compressed Only compress files that are not already compressed.
--show-default-cparams Print the parameters the level selects for each input.
--stream-size=# Specify size of streaming input from STDIN.
--size-hint=# Optimize compression parameters for streaming input of approximately size #.
--target-compressed-block-size=#
Generate compressed blocks of approximately # size.
--no-dictID Don't write `dictID` into the header (dictionary compression only).
--[no-]content-size Write the input size into the frame header when it is known. [Default: Write]
--[no-]compress-literals Force (un)compressed literals.
--format=zstd Compress files to the `.zst` format. [Default]
Advanced decompression options:
-l Print information about Zstandard-compressed files.
--test Test compressed file integrity.
-M# Set the memory usage limit to # megabytes.
--[no-]pass-through Pass through uncompressed files as-is. [Default: Disabled; Enabled for zstdcat]
Dictionary builder:
--train Create a dictionary from a training set of files.
--train-cover[=k=#,d=#,steps=#,split=#,shrink[=#]]
Use the cover algorithm (with optional arguments).
--train-fastcover[=k=#,d=#,f=#,steps=#,split=#,accel=#,shrink[=#]]
Use the fast cover algorithm (with optional arguments).
--train-legacy[=s=#] Use the legacy algorithm with selectivity #. [Default: 9]
-B# Cut each file into samples of size #; otherwise each file
is one sample of up to 128 KiB. At least 5 samples are needed.
-o NAME Use NAME as dictionary name. [Default: dictionary]
--maxdict=# Limit dictionary to specified size #. [Default: 112640]
--dictID=# Force dictionary ID to #. [Default: Random]
Benchmark options:
-b# Perform benchmarking with compression level #. [Default: 3]
-e# Test all compression levels up to #; starting level is `-b#`. [Default: 1]
-i# Set the minimum evaluation to time # seconds. [Default: 3]
-B# Cut file into independent chunks of size #. [Default: No chunking]
-S Output one benchmark result per input file. [Default: Consolidated result]
-D dictionary Benchmark using dictionary
Environment: ZSTD_CLEVEL sets the default compression level; ZSTD_NBTHREADS is read and validated.
Accepted for compatibility, with no effect here: -T#/--threads=#, --single-thread,
--auto-threads, -B# and --block-size=# when compressing (the job size of a
multi-threaded run), --adapt, --zstd=ovlog=#, --[no-]sparse,
--[no-]asyncio, --[no-]mmap-dict, --[no-]row-match-finder (compression runs
single-threaded).
Rejected rather than ignored, because they would change the result: --format=
other than zstd, --rsyncable (needs worker threads), and -M/--memory below the
enforced ceiling when decoding.
Trainer shrink[=#] takes effect: the smallest dictionary within # percent
(default 1) of the full one is kept.
A new output file keeps its source's permissions.
";
const STRATEGY_NAMES: [&str; 9] = [
"ZSTD_fast",
"ZSTD_dfast",
"ZSTD_greedy",
"ZSTD_lazy",
"ZSTD_lazy2",
"ZSTD_btlazy2",
"ZSTD_btopt",
"ZSTD_btultra",
"ZSTD_btultra2",
];
fn write_default_cparams(
out: &mut impl Write,
name: &str,
size: Option<u64>,
dictionary_size: usize,
level: i32,
) -> io::Result<()> {
match size {
Some(bytes) => writeln!(out, "{name} ({bytes} bytes)")?,
None => writeln!(out, "{name} (src size unknown)")?,
}
let params =
LevelParameters::for_level(level, size.filter(|&bytes| bytes != 0), dictionary_size);
let ordinal = params.strategy.ordinal();
let strategy = STRATEGY_NAMES[ordinal as usize - 1];
writeln!(out, " - windowLog : {}", params.window_log)?;
writeln!(out, " - chainLog : {}", params.chain_log)?;
writeln!(out, " - hashLog : {}", params.hash_log)?;
writeln!(out, " - searchLog : {}", params.search_log)?;
writeln!(out, " - minMatch : {}", params.min_match)?;
writeln!(out, " - targetLength : {}", params.target_length)?;
writeln!(out, " - strategy : {strategy} ({ordinal})")
}
fn show_default_cparams(opts: &Options) -> Result<()> {
let dictionary_size = match dictionary_path(opts) {
Some(path) => fs::metadata(path)
.wrap_err_with(|| format!("failed to inspect dictionary file {}", path.display()))?
.len(),
None => 0,
};
let dictionary_size = usize::try_from(dictionary_size)
.map_err(|_| eyre!("dictionary of {dictionary_size} bytes does not fit in memory"))?;
let mut err = io::stderr().lock();
let stdin = [PathBuf::from("-")];
let inputs: &[PathBuf] = if opts.inputs.is_empty() {
&stdin
} else {
&opts.inputs
};
for input in inputs {
let (name, size) = if input == Path::new("-") {
(STDIN_MARK.to_string(), None)
} else {
let size = fs::metadata(input)
.ok()
.filter(fs::Metadata::is_file)
.map(|metadata| metadata.len());
(input.display().to_string(), size)
};
write_default_cparams(&mut err, &name, size, dictionary_size, opts.level)
.wrap_err("failed to write the default parameters")?;
}
Ok(())
}
fn dictionary_path(opts: &Options) -> Option<&Path> {
opts.dict.as_deref().or(opts.patch_from.as_deref())
}
fn patch_window_log(source_size: u64) -> Result<u32> {
use structured_zstd::encoding::CParameter;
let file_window_log = u64::BITS - source_size.max(1).leading_zeros();
let lower = u32::try_from(CParameter::WindowLog.bounds().lower_bound)
.expect("the window log lower bound is a small positive number");
let upper = max_window_log();
if file_window_log > upper {
bail!(
"Can't handle files larger than {} MiB with --patch-from: the patch would \
need a window past the largest this build writes and reads back",
(1u64 << upper) >> 20
);
}
Ok(file_window_log.max(lower))
}
fn patch_needs_long_distance(
level: i32,
advanced: &AdvancedParams,
source_size: u64,
reference_size: u64,
) -> bool {
let file_window_log = u64::BITS - source_size.max(1).leading_zeros();
let reference_size = usize::try_from(reference_size).unwrap_or(usize::MAX);
let params = LevelParameters::for_level(level, Some(source_size), reference_size);
let chain_log = advanced.chain_log.unwrap_or(params.chain_log);
let strategy = advanced.strategy.unwrap_or(params.strategy);
let cycle_log = chain_log - u32::from(strategy >= Strategy::Btlazy2);
file_window_log > cycle_log
}
fn patch_run_needs_long_distance(opts: &Options, source_size: u64, reference_size: u64) -> bool {
let levels = if opts.bench {
opts.bench_start..=opts.bench_end
} else {
opts.level..=opts.level
};
levels
.into_iter()
.any(|level| patch_needs_long_distance(level, &opts.advanced, source_size, reference_size))
}
fn reads_dictionary(opts: &Options) -> bool {
opts.bench || !matches!(opts.mode, Mode::List | Mode::Train)
}
fn load_dictionary(opts: &Options) -> Result<Option<Vec<u8>>> {
let Some(path) = dictionary_path(opts) else {
return Ok(None);
};
if !reads_dictionary(opts) {
return Ok(None);
}
let named = fs::metadata(path)
.wrap_err_with(|| format!("failed to inspect dictionary file {}", path.display()))?;
if !named.is_file() {
bail!("-D needs a regular file: {} is not one", path.display());
}
let weigh = |size: u64| -> Result<usize> {
let capacity = usize::try_from(size).map_err(|_| {
eyre!(
"dictionary file {} is {size} bytes, more than this machine can address",
path.display()
)
})?;
if let Some(limit) = opts.memory_limit
&& decodes(opts)
{
check_memory_limit(limit, size, 2)?;
}
Ok(capacity)
};
weigh(named.len())?;
let file = File::open(path)
.wrap_err_with(|| format!("failed to open dictionary file {}", path.display()))?;
let opened = file
.metadata()
.wrap_err_with(|| format!("failed to inspect dictionary file {}", path.display()))?;
if !opened.is_file() {
bail!("-D needs a regular file: {} is not one", path.display());
}
let size = opened.len();
let capacity = weigh(size)?;
let mut bytes = Vec::with_capacity(capacity);
file.take(size + 1)
.read_to_end(&mut bytes)
.wrap_err_with(|| format!("failed to read dictionary file {}", path.display()))?;
if bytes.len() as u64 > size {
bail!("{} grew while it was being read; run again", path.display());
}
Ok(Some(bytes))
}
#[derive(Default)]
struct Codecs {
encoder: Option<structured_zstd::encoding::EncoderDictionary>,
decoder: Option<structured_zstd::decoding::DictionaryHandle>,
compressor: Option<(FrameSettings, CompressionContext)>,
decompressor: structured_zstd::decoding::FrameDecoder,
chunk: Vec<u8>,
dictionary_size: u64,
trace: Option<trace::Trace>,
}
impl Codecs {
fn open_trace(&mut self, opts: &Options, single_pass: bool) -> Result<()> {
let Some(path) = &opts.trace else {
return Ok(());
};
refuse_trace_alias(opts, path)?;
match trace::Trace::open(path, UPSTREAM_VERSION_NUMBER, single_pass) {
Ok(mut trace) => {
trace.reports_failures = opts.verbosity >= 1;
self.trace = Some(trace);
}
Err(err) => display!(
opts.verbosity,
2,
"zstd: --trace {}: {err}; nothing is traced",
path.display()
),
}
Ok(())
}
}
fn refuse_trace_alias(opts: &Options, trace: &Path) -> Result<()> {
let trace = &dangling_link_target(trace);
let names = |other: &Path| -> Result<bool> {
let other = &dangling_link_target(other);
Ok(names_the_same_file(trace, other)?
|| (trace.exists() && other.exists() && paths_point_to_same_file(trace, other)?))
};
let refuse = |role: &str, other: &Path| -> Result<()> {
if names(other)? {
bail!(
"--trace {} is also {role} {}",
trace.display(),
other.display()
);
}
Ok(())
};
if let Some(dictionary) = dictionary_path(opts)
&& reads_dictionary(opts)
{
refuse("the dictionary", dictionary)?;
}
for input in opts.inputs.iter().filter(|input| *input != Path::new("-")) {
refuse("the input", input)?;
}
for list in &opts.filelists {
refuse("the file list", list)?;
}
if is_stdin(trace) {
bail!(
"--trace {} is also standard input, which this run reads",
trace.display()
);
}
if writes_to_stdout(opts) && is_stdout(trace) {
bail!(
"--trace {} is also standard output, which this run writes to",
trace.display()
);
}
if (opts.verbosity >= 1
|| opts.show_default_cparams
|| opts.pause
|| opts.progress == Progress::Always)
&& is_stderr(trace)
{
bail!(
"--trace {} is also standard error, which this run writes to",
trace.display()
);
}
if opts.bench || opts.to_stdout || matches!(opts.mode, Mode::Test | Mode::List) {
return Ok(());
}
if opts.mode == Mode::Train {
return refuse("the dictionary being trained", &training_output(opts));
}
if let Some(output) = &opts.output {
let writes_output = opts.inputs.is_empty()
|| opts
.inputs
.iter()
.any(|input| !skips_as_compressed(opts, input));
return if writes_output {
refuse("the output", output)
} else {
Ok(())
};
}
if matches!(opts.mode, Mode::Compress | Mode::Decompress) {
for input in opts.inputs.iter().filter(|input| *input != Path::new("-")) {
if skips_as_compressed(opts, input) {
continue;
}
if let Ok(output) = derive_output_path(opts, input) {
refuse("the output", &output)?;
for directory in output.ancestors().skip(1) {
if directory.as_os_str().is_empty() || directory.exists() {
break;
}
refuse("the output directory", directory)?;
}
}
}
}
Ok(())
}
fn writes_to_stdout(opts: &Options) -> bool {
if opts.bench {
return opts.verbosity >= 1;
}
match opts.mode {
Mode::List => true,
Mode::Compress | Mode::Decompress => {
opts.to_stdout
|| (opts.output.is_none()
&& (opts.inputs.is_empty()
|| opts.inputs.iter().any(|input| input == Path::new("-"))))
}
Mode::Test | Mode::Train => false,
}
}
fn skips_as_compressed(opts: &Options, input: &Path) -> bool {
compresses(opts) && opts.exclude_compressed && inputs::has_compressed_extension(input)
}
fn is_stdout(path: &Path) -> bool {
names_stream(path, &["/dev/stdout", "/dev/fd/1", "/proc/self/fd/1"])
|| identity::of_stdout().is_some_and(|stdout| identity::of_path(path) == Some(stdout))
}
fn is_stderr(path: &Path) -> bool {
names_stream(path, &["/dev/stderr", "/dev/fd/2", "/proc/self/fd/2"])
|| identity::of_stderr().is_some_and(|stderr| identity::of_path(path) == Some(stderr))
}
fn is_stdin(path: &Path) -> bool {
names_stream(path, &["/dev/stdin", "/dev/fd/0", "/proc/self/fd/0"])
|| (fs::metadata(path)
.is_ok_and(|metadata| metadata.is_file() || is_fifo_type(&metadata.file_type()))
&& identity::of_stdin().is_some_and(|stdin| identity::of_path(path) == Some(stdin)))
}
fn names_stream(path: &Path, names: &[&str]) -> bool {
cfg!(unix) && names.iter().any(|name| path == Path::new(name))
}
fn training_output(opts: &Options) -> PathBuf {
opts.output
.clone()
.unwrap_or_else(|| PathBuf::from("dictionary"))
}
fn dangling_link_target(path: &Path) -> PathBuf {
let mut current = path.to_path_buf();
for _ in 0..40 {
if current.exists() {
break;
}
let Ok(target) = fs::read_link(¤t) else {
break;
};
current = match current.parent() {
Some(parent) if target.is_relative() => parent.join(target),
_ => target,
};
}
current
}
fn trace_frame(trace: &mut Option<trace::Trace>, frame: &trace::Frame) {
if let Some(open) = trace.as_mut()
&& let Err(err) = open.record(frame)
{
if open.reports_failures {
eprintln!("zstd: --trace: {err}; no further frames are traced");
}
*trace = None;
}
}
const IO_CHUNK: usize = 128 * 1024;
impl Codecs {
fn prepare(
raw: Option<&[u8]>,
raw_content: bool,
for_compression: bool,
for_decoding: bool,
) -> Result<Self> {
use structured_zstd::decoding::{Dictionary, DictionaryHandle};
use structured_zstd::encoding::EncoderDictionary;
let Some(raw) = raw.filter(|raw| !raw.is_empty()) else {
return Ok(Self::default());
};
let mut prepared = Self {
dictionary_size: raw.len() as u64,
..Self::default()
};
if for_compression {
let dictionary = if raw_content {
Dictionary::from_raw_content(0, raw.to_vec())
.map(EncoderDictionary::from_dictionary)
} else {
EncoderDictionary::from_serialized_or_raw_content(raw)
};
prepared.encoder =
Some(dictionary.map_err(|err| eyre!("invalid dictionary: {err:?}"))?);
}
if for_decoding {
let dictionary = if raw_content {
Dictionary::from_raw_content(0, raw.to_vec())
} else {
Dictionary::from_serialized_or_raw_content(raw)
};
prepared.decoder = Some(DictionaryHandle::from_dictionary(
dictionary.map_err(|err| eyre!("failed to parse dictionary: {err:?}"))?,
));
}
Ok(prepared)
}
fn compressor(&mut self, settings: &FrameSettings) -> Result<&mut CompressionContext> {
let key = FrameSettings {
pledged_size: None,
..*settings
};
if self
.compressor
.as_ref()
.is_none_or(|(built_for, _)| *built_for != key)
{
self.compressor = None;
self.compressor = Some((key, new_compressor(&key, self.encoder.as_ref())?));
}
let (_, context) = self
.compressor
.as_mut()
.expect("a context for these settings is in place");
Ok(context)
}
}
const STDIN_MARK: &str = "/*stdin*\\";
fn reads_stdin(inputs: &[PathBuf]) -> bool {
inputs.is_empty() || inputs.iter().any(|input| input == Path::new("-"))
}
fn writes_stdout(opts: &Options) -> bool {
opts.to_stdout
|| (opts.output.is_none() && opts.inputs.iter().all(|input| input == Path::new("-")))
}
fn run(mut opts: Options) -> Result<usize> {
let Selection {
files,
explicit,
unreadable_dirs,
} = inputs::select_inputs(
std::mem::take(&mut opts.inputs),
&opts.filelists,
opts.recursive,
opts.follow_links,
opts.verbosity,
)?;
let skipped = if opts.ignore_read_errors {
0
} else {
unreadable_dirs
};
let streams =
matches!(opts.mode, Mode::Compress | Mode::Decompress | Mode::Test) && !opts.bench;
if files.is_empty() && explicit && streams {
display!(
opts.verbosity,
1,
"please provide correct input file(s) or non-empty directories -- ignored"
);
return Ok(skipped);
}
if files.is_empty() && explicit && opts.bench {
bail!("-b requires one or more regular input files to benchmark");
}
opts.inputs = files;
Ok(run_selected(opts)? + skipped)
}
fn run_selected(mut opts: Options) -> Result<usize> {
let streams =
matches!(opts.mode, Mode::Compress | Mode::Decompress | Mode::Test) && !opts.bench;
if streams && reads_stdin(&opts.inputs) && !opts.force_stdin && opts.console.stdin() {
bail!("stdin is a console, aborting");
}
refuse_console_stdout(&opts, opts.console.stdout())?;
let has_stdout_output = matches!(opts.mode, Mode::Compress | Mode::Decompress)
&& !opts.bench
&& writes_stdout(&opts);
if has_stdout_output && opts.verbosity == DEFAULT_LEVEL {
opts.verbosity = 1;
}
if !opts.console.stderr() && opts.progress != Progress::Always {
opts.progress = Progress::Never;
}
if has_stdout_output && opts.remove_source {
display!(
opts.verbosity,
3,
"Note: src files are not removed when output is stdout"
);
opts.remove_source = false;
}
if opts.mode == Mode::Test {
opts.remove_source = false;
}
if opts.bench && opts.inputs.is_empty() {
opts.dict = None;
opts.patch_from = None;
}
if opts.patch_from.is_some() {
if opts.inputs.len() > 1 {
bail!("error : can't use --patch-from=# on multiple files");
}
if compresses(&opts) {
let source_size = match (opts.inputs.first(), opts.pledged_size) {
(Some(input), _) if input != Path::new("-") => fs::metadata(input)
.map_err(|err| eyre!("can't stat {} : {err}", input.display()))?
.len(),
(_, Some(size)) => size,
_ => bail!("Using --patch-from with stdin requires --stream-size"),
};
opts.advanced.window_log = Some(patch_window_log(source_size)?);
let reference = opts.patch_from.as_deref().expect("set in this branch");
let reference_size = fs::metadata(reference)
.map_err(|err| eyre!("can't stat {} : {err}", reference.display()))?
.len();
if patch_run_needs_long_distance(&opts, source_size, reference_size) {
if !opts.long {
display!(opts.verbosity, 2, "long mode automatically triggered");
}
opts.long = true;
}
}
}
let dict_bytes = load_dictionary(&opts)?;
opts.decode_window = decode_window_ceiling(&opts);
if opts.bench {
run_benchmark(&opts, dict_bytes)?;
return Ok(0);
}
let mut codecs = Codecs::prepare(
dict_bytes.as_deref(),
opts.patch_from.is_some(),
compresses(&opts),
decodes(&opts),
)?;
codecs.open_trace(&opts, false)?;
drop(dict_bytes);
if opts.mode == Mode::Train {
train_dictionary(&opts)?;
return Ok(0);
}
if opts.mode == Mode::List {
return list_files(&opts);
}
if opts.show_default_cparams {
if opts.mode == Mode::Decompress {
bail!("error : can't use --show-default-cparams in decompression mode");
}
if opts.mode == Mode::Compress {
show_default_cparams(&opts)?;
}
}
if let (Some(output), Some(dict)) = (&opts.output, dictionary_path(&opts))
&& !opts.to_stdout
&& opts.mode == Mode::Compress
&& names_the_same_file(output, dict)?
{
bail!(
"{} would be written over the dictionary {}",
output.display(),
dict.display(),
);
}
if let Some(output) = &opts.output
&& !opts.to_stdout
&& opts.mode != Mode::Test
{
for input in &opts.inputs {
if input != Path::new("-")
&& (names_the_same_file(output, input)?
|| (output.exists()
&& input.exists()
&& paths_point_to_same_file(input, output)?))
{
bail!(
"{} would be written over {}, which is also an input",
output.display(),
input.display(),
);
}
}
}
let total = opts.inputs.len().max(1);
match (opts.to_stdout, &opts.output) {
(true, _) => process_concatenated(&opts, &mut codecs, None, total),
(false, Some(output)) if opts.mode != Mode::Test => {
process_concatenated(&opts, &mut codecs, Some(output), total)
}
_ => process_separately(&opts, &mut codecs, total),
}
}
fn process_concatenated(
opts: &Options,
codecs: &mut Codecs,
output: Option<&Path>,
total: usize,
) -> Result<usize> {
let mut remove_source = opts.remove_source;
if total > 1 {
match output {
None => display!(
opts.verbosity,
2,
"zstd: WARNING: all input files will be processed and concatenated into stdout."
),
Some(output) => display!(
opts.verbosity,
2,
"zstd: WARNING: all input files will be processed and concatenated into a single output file: {}",
output.display()
),
}
display!(
opts.verbosity,
2,
"The concatenated output CANNOT regenerate original file names nor directory structure."
);
if remove_source {
display!(
opts.verbosity,
2,
"Since it's a destructive operation, input files will not be removed."
);
remove_source = false;
}
if output.is_some() && !opts.force {
if opts.verbosity <= 1 {
display!(
opts.verbosity,
1,
"Concatenating multiple processed inputs into a single output loses file metadata."
);
display!(opts.verbosity, 1, "Aborting.");
return Ok(total);
}
if !confirm(
"Proceed? (y/n): ",
"Aborting...",
reads_stdin(&opts.inputs),
&mut io::stdin().lock(),
) {
return Ok(total);
}
}
}
let mut tally = Tally::default();
let inputs: Vec<&Path> = if opts.inputs.is_empty() {
vec![Path::new("-")]
} else {
opts.inputs.iter().map(PathBuf::as_path).collect()
};
match output {
None => {
let mut sink = io::stdout().lock();
for input in inputs {
let outcome = stream_input_to(opts, codecs, input, &mut sink, STDOUT_MARK, total)?;
tally.record(outcome);
}
}
Some(output) => {
if let [input] = inputs[..] {
let (source, metadata) = if input == Path::new("-") {
(None, None)
} else {
match open_input(opts, input) {
Ok(Some((source, metadata))) => (Some(source), Some(metadata)),
Ok(None) => return Ok(0),
Err(err) => {
display!(opts.verbosity, 1, "zstd: {err}");
return Ok(1);
}
}
};
let name = if source.is_some() {
input.display().to_string()
} else {
STDIN_MARK.to_string()
};
let written =
write_output_file(opts, output, metadata.as_ref(), |sink| match source {
Some(source) => stream_opened(
opts,
codecs,
source,
metadata.as_ref().expect("an opened input has metadata"),
sink,
),
None => stream_stdin(opts, codecs, sink),
})?;
match written {
Some((processed, placement)) => {
file_summary(
opts,
total,
&name,
&output.display().to_string(),
&processed,
);
if remove_source && input != Path::new("-") {
remove_source_if_requested(opts, input, placement)?;
}
tally.record(Outcome::Done(processed));
}
None => tally.record(Outcome::Refused),
}
} else {
let written = write_output_file_if(
opts,
output,
None,
|sink| {
let mut tally = Tally::default();
for input in &inputs {
let outcome = stream_input_to(
opts,
codecs,
input,
&mut *sink,
&output.display().to_string(),
total,
)?;
tally.record(outcome);
}
Ok(tally)
},
|tally| tally.processed > 0,
)?;
match written {
Some((inner, _)) => tally = inner,
None => tally.failed = total,
}
}
}
}
multi_summary(opts, total, &tally);
Ok(tally.failed)
}
fn process_separately(opts: &Options, codecs: &mut Codecs, total: usize) -> Result<usize> {
if matches!(opts.mode, Mode::Compress | Mode::Decompress) {
for input in &opts.inputs {
if input == Path::new("-") {
continue;
}
let Ok(output) = derive_output_path(opts, input) else {
continue;
};
if let Some(dict) = dictionary_path(opts)
&& opts.mode == Mode::Compress
&& names_the_same_file(&output, dict)?
{
bail!(
"{} would be written over the dictionary {}",
input.display(),
dict.display(),
);
}
if let Some(dict) = dictionary_path(opts)
&& opts.remove_source
&& !opts.keep
&& (names_the_same_file(input, dict)?
|| (dict.exists() && paths_point_to_same_file(input, dict)?))
{
bail!(
"--rm would delete {}, which is also the dictionary {} needed to read the result",
input.display(),
dict.display(),
);
}
for other in &opts.inputs {
if other != Path::new("-") && names_the_same_file(&output, other)? {
bail!(
"{} would be written over {}, which is also an input",
input.display(),
other.display(),
);
}
}
}
}
let mut tally = Tally::default();
if opts.inputs.is_empty() {
let outcome = stream_input_to(
opts,
codecs,
Path::new("-"),
io::stdout().lock(),
STDOUT_MARK,
total,
)?;
tally.record(outcome);
}
for input in &opts.inputs {
let outcome = if input == Path::new("-") {
stream_input_to(opts, codecs, input, io::stdout().lock(), STDOUT_MARK, total)?
} else {
match process_file(opts, input, codecs, total) {
Ok(outcome) => outcome,
Err(err) => {
display!(opts.verbosity, 1, "zstd: {err}");
Outcome::Refused
}
}
};
tally.record(outcome);
}
for name in flat_output_collisions(opts) {
display!(
opts.verbosity,
2,
"WARNING: Two files have same filename: {}",
Path::new(&name).display()
);
}
multi_summary(opts, total, &tally);
Ok(tally.failed)
}
fn flat_output_collisions(opts: &Options) -> Vec<OsString> {
if opts.output_dir.is_none()
|| opts.output_dir_mirror.is_some()
|| !matches!(opts.mode, Mode::Compress | Mode::Decompress)
{
return Vec::new();
}
let outputs: Vec<PathBuf> = opts
.inputs
.iter()
.filter(|input| input.as_path() != Path::new("-"))
.filter_map(|input| derive_output_path(opts, input).ok())
.collect();
inputs::shared_file_names(&outputs)
}
enum Outcome {
Done(Processed),
Skipped,
Refused,
}
struct Processed {
read: u64,
written: u64,
}
#[derive(Default)]
struct Tally {
processed: usize,
failed: usize,
read: u64,
written: u64,
}
impl Tally {
fn record(&mut self, outcome: Outcome) {
match outcome {
Outcome::Done(processed) => {
self.processed += 1;
self.read += processed.read;
self.written += processed.written;
}
Outcome::Skipped => {}
Outcome::Refused => self.failed += 1,
}
}
}
fn file_summary(
opts: &Options,
total: usize,
name: &str,
destination: &str,
processed: &Processed,
) {
if total > 1 && opts.verbosity < 3 {
return;
}
let verbose = opts.verbosity > 3;
match opts.mode {
Mode::Compress => {
let read = HumanSize::new(processed.read, verbose);
let written = HumanSize::new(processed.written, verbose);
if processed.read == 0 {
display!(
opts.verbosity,
2,
"{name:<20} : ({read:>6} => {written:>6}, {destination})"
);
} else {
display!(
opts.verbosity,
2,
"{name:<20} :{:>6.2}% ({read:>6} => {written:>6}, {destination})",
processed.written as f64 / processed.read as f64 * 100.0
);
}
}
Mode::Decompress | Mode::Test => {
display!(opts.verbosity, 2, "{name:<20}: {} bytes", processed.written);
}
Mode::List | Mode::Train => {}
}
}
fn multi_summary(opts: &Options, total: usize, tally: &Tally) {
if tally.processed < 1 || total <= 1 {
return;
}
let verbose = opts.verbosity > 3;
match opts.mode {
Mode::Compress => {
let read = HumanSize::new(tally.read, verbose);
let written = HumanSize::new(tally.written, verbose);
if tally.read == 0 {
display!(
opts.verbosity,
2,
"{:>3} files compressed : ({} => {})",
tally.processed,
read.columns(6),
written.columns(6)
);
} else {
display!(
opts.verbosity,
2,
"{:>3} files compressed : {:.2}% ({} => {})",
tally.processed,
tally.written as f64 / tally.read as f64 * 100.0,
read.columns(6),
written.columns(6)
);
}
}
Mode::Decompress | Mode::Test => {
display!(
opts.verbosity,
2,
"{} files decompressed : {:>6} bytes total",
tally.processed,
tally.written
);
}
Mode::List | Mode::Train => {}
}
}
fn run_benchmark(opts: &Options, dict: Option<Vec<u8>>) -> Result<()> {
let synthetic = opts.inputs.is_empty();
if opts.inputs.iter().any(|input| input == Path::new("-")) {
bail!("-b cannot benchmark stdin; name a file (`./-` for one called `-`)");
}
let mut sum = 0u64;
let mut largest = 0u64;
let mut sizes = Vec::with_capacity(opts.inputs.len().max(1));
if synthetic {
let size = opts.block_size.unwrap_or(synthetic::DEFAULT_SIZE as u64);
sum = size;
largest = size;
sizes.push(size);
}
for input in &opts.inputs {
let metadata = fs::metadata(input)
.wrap_err_with(|| format!("failed to inspect {}", input.display()))?;
if !metadata.is_file() {
bail!("-b needs regular files: {} is not one", input.display());
}
sum = sum
.checked_add(metadata.len())
.ok_or_else(|| eyre!("the inputs add up to more than any machine can address"))?;
largest = largest.max(metadata.len());
sizes.push(metadata.len());
}
let mut held = None;
if let Some(limit) = opts.memory_limit {
let subject: &[u64] = if opts.bench_separately {
let at = sizes
.iter()
.position(|&size| size == largest)
.expect("the largest size is one of the sizes");
&sizes[at..=at]
} else {
&sizes
};
let inputs = if opts.bench_separately { largest } else { sum };
let extent = bench_frames_extent(subject, opts.block_size);
let frame = extent.map(|(room, _)| room);
let widest_chunk = extent.map_or(0, |(_, widest)| widest);
let dictionary = dict
.as_ref()
.map(|bytes| structured_zstd::encoding::DictionarySizes::raw_content(bytes.len()));
let mut encoder = 0u64;
for level in opts.bench_start..=opts.bench_end {
let settings = FrameSettings {
level,
..FrameSettings::from_options(opts)
};
let compression_level = CompressionLevel::from_level(level);
let bytes = match frame_parameters(compression_level, &settings)? {
Some(parameters) => {
structured_zstd::encoding::estimated_compression_workspace_bytes_for_parameters(
¶meters,
Some(widest_chunk),
dictionary,
)
}
None => structured_zstd::encoding::estimated_compression_workspace_bytes_for_run(
compression_level,
Some(widest_chunk),
None,
false,
dictionary,
),
};
encoder = encoder.max(bytes as u64);
}
let buffers = frame
.and_then(|frame| {
inputs
.checked_mul(2)?
.checked_add(frame)?
.checked_add(encoder)
})
.ok_or_else(|| eyre!("the inputs add up to more than any machine can address"))?;
check_memory_limit(limit, buffers, 1)?;
let mut total = buffers;
if let Some(bytes) = &dict {
total = (bytes.len() as u64)
.checked_mul(3)
.and_then(|dictionaries| buffers.checked_add(dictionaries))
.ok_or_else(|| eyre!("the inputs add up to more than any machine can address"))?;
check_memory_limit(limit, total, 1)?;
}
held = Some(total);
}
let decode = bench_decode_settings(opts, held);
let codecs = &mut Codecs::prepare(dict.as_deref(), opts.patch_from.is_some(), true, true)?;
codecs.open_trace(opts, true)?;
drop(dict);
if synthetic {
let size = usize::try_from(sizes[0]).expect("a benchmark block size fits usize");
let (label, data) = match opts.bench_compressibility {
Some(percent) => (
format!("Synthetic {percent}%"),
synthetic::compressible(size, percent),
),
None => ("Lorem ipsum".to_string(), synthetic::lorem(size)),
};
let data = data.map_err(|err| {
eyre!("-b: not enough memory for {size} bytes of generated data: {err}")
})?;
return benchmark_one(opts, codecs, &decode, &label, &data, &sizes);
}
if opts.bench_separately {
for (input, size) in opts.inputs.iter().zip(&sizes) {
let data =
read_inputs_bounded(std::slice::from_ref(input), std::slice::from_ref(size))?;
benchmark_one(
opts,
codecs,
&decode,
&input.display().to_string(),
&data,
std::slice::from_ref(size),
)?;
}
return Ok(());
}
let data = read_inputs_bounded(&opts.inputs, &sizes)?;
let label = match opts.inputs.as_slice() {
[only] => only.display().to_string(),
many => format!(" {} files", many.len()),
};
benchmark_one(opts, codecs, &decode, &label, &data, &sizes)
}
fn bench_decode_settings(opts: &Options, held: Option<u64>) -> DecodeSettings {
use structured_zstd::decoding::{MAX_DECODER_WINDOW_SIZE, MAXIMUM_ALLOWED_WINDOW_SIZE};
let settings = DecodeSettings::from_options(opts);
let (Some(limit), Some(held)) = (opts.memory_limit, held) else {
return settings;
};
let overhead = held.checked_add(DECODER_AUXILIARY_BYTES);
DecodeSettings {
max_window: overhead
.and_then(|overhead| limit.checked_sub(overhead))
.unwrap_or(0)
.clamp(MAXIMUM_ALLOWED_WINDOW_SIZE, MAX_DECODER_WINDOW_SIZE),
memory_overhead: overhead,
..settings
}
}
fn read_inputs_bounded(inputs: &[PathBuf], sizes: &[u64]) -> Result<Vec<u8>> {
let mut total = 0usize;
for size in sizes {
total = usize::try_from(*size)
.ok()
.and_then(|size| total.checked_add(size))
.ok_or_else(|| eyre!("the inputs add up to more than this machine can hold"))?;
}
let mut data = Vec::with_capacity(total);
for (input, size) in inputs.iter().zip(sizes) {
let file = File::open(input)
.wrap_err_with(|| format!("failed to open input file {}", input.display()))?;
let before = data.len();
file.take(size + 1)
.read_to_end(&mut data)
.wrap_err_with(|| format!("failed to read {}", input.display()))?;
if (data.len() - before) as u64 != *size {
bail!(
"{} changed size while it was being read; run again",
input.display()
);
}
}
Ok(data)
}
fn bench_chunk_lengths(
file_sizes: &[u64],
block_size: Option<u64>,
) -> impl Iterator<Item = u64> + '_ {
let block = block_size.filter(|&size| size >= MIN_BENCH_BLOCK_SIZE);
file_sizes.iter().flat_map(move |&size| {
let piece = block.unwrap_or(size.max(1));
(0..size.div_ceil(piece)).map(move |at| piece.min(size - at * piece))
})
}
fn bench_frames_extent(file_sizes: &[u64], block_size: Option<u64>) -> Option<(u64, u64)> {
let block = block_size.filter(|&size| size >= MIN_BENCH_BLOCK_SIZE);
let bound = |len: u64| -> Option<u64> {
u64::try_from(structured_zstd::encoding::compress_bound(
usize::try_from(len).ok()?,
))
.ok()
};
let mut room = 0u64;
let mut widest = 0u64;
for &size in file_sizes {
if size == 0 {
continue;
}
let piece = block.unwrap_or(size);
let whole = size / piece;
let tail = size % piece;
room = room.checked_add(whole.checked_mul(bound(piece)?)?)?;
if tail > 0 {
room = room.checked_add(bound(tail)?)?;
}
widest = widest.max(size.min(piece));
}
Some((room, widest))
}
fn benchmark_one(
opts: &Options,
codecs: &mut Codecs,
decode: &DecodeSettings,
label: &str,
data: &[u8],
file_sizes: &[u64],
) -> Result<()> {
use std::time::Instant;
if data.is_empty() {
bail!("-b: {label} is empty");
}
let mb = data.len() as f64 / 1e6;
let name = bench_display_name(label);
display!(
opts.verbosity,
3,
"Benchmarking {label} from level {} to {}",
opts.bench_start,
opts.bench_end
);
if let Some(header) = bench_quiet_header(opts, data.len()) {
println!("{header}");
}
debug_assert_eq!(
bench_chunk_lengths(file_sizes, opts.block_size).sum::<u64>(),
data.len() as u64,
"the frames cover the input exactly"
);
let frames_bound = bench_frames_extent(file_sizes, opts.block_size)
.and_then(|(room, _)| usize::try_from(room).ok())
.ok_or_else(|| eyre!("-b: {label} is more than this machine can hold compressed"))?;
let mut compressed = Vec::new();
compressed
.try_reserve_exact(frames_bound)
.map_err(|err| eyre!("-b: not enough memory for the frames of {label}: {err}"))?;
let mut decoded = Vec::new();
decoded
.try_reserve_exact(data.len())
.map_err(|err| eyre!("-b: not enough memory to decode {label}: {err}"))?;
let _realtime = if opts.bench_realtime {
display!(opts.verbosity, 2, "Note : switching to real-time priority ");
let raised = priority::raise_to_realtime();
if raised.is_none() {
display!(
opts.verbosity,
3,
"Note : real-time priority was refused; measuring at the ordinary priority"
);
}
raised
} else {
None
};
for level in opts.bench_start..=opts.bench_end {
validate_level(level)?;
let settings = FrameSettings {
level,
size_hint: None,
checksum: false,
..FrameSettings::from_options(opts)
};
let mut best_compress = f64::MAX;
let start = Instant::now();
loop {
compressed.clear();
let t = Instant::now();
let mut rest = data;
for chunk in bench_chunk_lengths(file_sizes, opts.block_size) {
let (piece, tail) = rest.split_at(chunk as usize);
rest = tail;
compress_stream(
piece,
&mut compressed,
&FrameSettings {
pledged_size: Some(chunk),
..settings
},
codecs,
)?;
}
best_compress = best_compress.min(t.elapsed().as_secs_f64());
if start.elapsed().as_secs_f64() >= opts.bench_secs {
break;
}
}
let mut best_decompress = f64::MAX;
let start = Instant::now();
loop {
decoded.clear();
let t = Instant::now();
decompress_stream(compressed.as_slice(), &mut decoded, codecs, decode)?;
best_decompress = best_decompress.min(t.elapsed().as_secs_f64());
if start.elapsed().as_secs_f64() >= opts.bench_secs {
break;
}
}
if decoded != data {
bail!("-b: level {level} did not decode {label} back to its input");
}
let c_speed = if best_compress > 0.0 {
mb / best_compress
} else {
f64::INFINITY
};
let d_speed = if best_decompress > 0.0 {
mb / best_decompress
} else {
f64::INFINITY
};
let result = BenchResult {
level,
input: data.len() as u64,
output: compressed.len() as u64,
compress_mb_s: c_speed,
decompress_mb_s: d_speed,
};
if opts.verbosity >= DEFAULT_LEVEL {
println!("{}", result.line(&name));
} else if opts.verbosity == 1 {
println!("{}", result.quiet_line(&name, opts.bench_param));
}
}
Ok(())
}
fn bench_quiet_header(opts: &Options, input_len: usize) -> Option<String> {
(opts.verbosity == 1 && opts.bench_param.is_none_or(|param| param == 0)).then(|| {
format!(
"bench {UPSTREAM_VERSION} : input {input_len} bytes, {} seconds, {} KB blocks",
opts.bench_secs as u64,
opts.block_size.unwrap_or(0) >> 10
)
})
}
struct BenchResult {
level: i32,
input: u64,
output: u64,
compress_mb_s: f64,
decompress_mb_s: f64,
}
impl BenchResult {
fn ratio(&self) -> f64 {
self.input as f64 / self.output as f64
}
fn line(&self, name: &str) -> String {
let ratio = self.ratio();
let ratio_digits = 1 + usize::from(ratio < 100.0) + usize::from(ratio < 10.0);
let speed_digits = if self.compress_mb_s < 10.0 { 2 } else { 1 };
format!(
"{:>2}#{:<17}:{:>10} ->{:>10} (x{:>5.ratio_digits$}), {:>6.speed_digits$} MB/s, {:>6.1} MB/s",
self.level,
name,
self.input,
self.output,
ratio,
self.compress_mb_s,
self.decompress_mb_s,
)
}
fn quiet_line(&self, name: &str, param: Option<u32>) -> String {
let mut line = format!(
"-{:<3}{:>11} ({:>5.3}) {:>6.2} MB/s {:>6.1} MB/s {}",
self.level,
self.output,
self.ratio(),
self.compress_mb_s,
self.decompress_mb_s,
name,
);
if let Some(param) = param.filter(|¶m| param != 0) {
line.push_str(&format!(" (param={param})"));
}
line
}
}
fn bench_display_name(label: &str) -> String {
let name = Path::new(label)
.file_name()
.map_or(label, |name| name.to_str().unwrap_or(label));
let chars = name.chars().count();
if chars > 17 {
name.chars().skip(chars - 17).collect()
} else {
name.to_string()
}
}
fn train_dictionary(opts: &Options) -> Result<()> {
use structured_zstd::dictionary::{
FinalizeOptions, create_legacy_dict_from_slice, optimize_cover_dict,
optimize_fastcover_dict, train_cover_dict, train_fastcover_dict,
};
if opts.inputs.iter().any(|input| input == Path::new("-")) {
bail!("--train cannot read stdin; name the sample files (`./-` for one called `-`)");
}
if opts.inputs.is_empty() {
bail!("--train requires one or more sample files");
}
if opts.to_stdout {
bail!("--train cannot write to stdout; name the dictionary with -o");
}
if opts.max_dict < structured_zstd::dictionary::MIN_TRAINED_DICT_SIZE {
bail!(
"--maxdict must be at least {} bytes",
structured_zstd::dictionary::MIN_TRAINED_DICT_SIZE
);
}
enum Plan {
FastCover(structured_zstd::dictionary::FastCoverOptions),
Cover(structured_zstd::dictionary::CoverOptions),
Legacy,
}
let plan = match opts.trainer {
Trainer::FastCover => Plan::FastCover(fastcover_options(&opts.trainer_params)?),
Trainer::Cover => Plan::Cover(cover_options(&opts.trainer_params)?),
Trainer::Legacy => Plan::Legacy,
};
let tuning_fits = match &plan {
Plan::FastCover(options) => {
structured_zstd::dictionary::check_fastcover_options(options, opts.max_dict)
}
Plan::Cover(options) => {
structured_zstd::dictionary::check_cover_options(options, opts.max_dict)
}
Plan::Legacy => Ok(()),
};
tuning_fits.map_err(|err| eyre!("dictionary training failed: {err}"))?;
let output = training_output(opts);
let in_place = writes_in_place(&output)?;
if !in_place {
ensure_regular_output_destination(&output)?;
if output.exists() && !opts.force {
bail!("{} already exists; use -f to overwrite", output.display());
}
}
for sample in &opts.inputs {
let clashes = names_the_same_file(sample, &output)?
|| (output.exists() && paths_point_to_same_file(sample, &output)?);
if clashes {
bail!(
"{} is a training sample; the dictionary cannot be written over it",
sample.display()
);
}
}
for input in &opts.inputs {
let metadata = fs::metadata(input)
.wrap_err_with(|| format!("failed to inspect {}", input.display()))?;
if !metadata.is_file() {
bail!(
"--train needs regular files: {} is not one",
input.display()
);
}
}
let finalize = FinalizeOptions {
dict_id: opts.dict_id,
level: opts.level,
};
let failed = |err: std::io::Error| eyre!("dictionary training failed: {err}");
let enough_samples = |count: usize| match &plan {
Plan::FastCover(options) => {
structured_zstd::dictionary::check_fastcover_sample_count(options, count)
}
Plan::Cover(options) => {
structured_zstd::dictionary::check_cover_sample_count(options, count)
}
Plan::Legacy => Ok(()),
};
let set = load_training_samples(&opts.inputs, opts.block_size, opts.memory_limit, |count| {
enough_samples(count).map_err(failed)
})?;
let dict = match plan {
Plan::Legacy => {
let mut dict = Vec::new();
create_legacy_dict_from_slice(
&set.corpus,
&set.sizes,
&mut dict,
opts.max_dict,
opts.selectivity,
finalize,
)
.map_err(failed)?;
dict
}
Plan::Cover(options) if options.k != 0 && options.d != 0 => {
train_cover_dict(&set.corpus, &set.sizes, opts.max_dict, &options, finalize)
.map_err(failed)?
}
Plan::Cover(options) => {
let (dict, chosen) =
optimize_cover_dict(&set.corpus, &set.sizes, opts.max_dict, &options, finalize)
.map_err(failed)?;
display!(
opts.verbosity,
2,
"k={}\nd={}\nsteps={}\nsplit={}",
chosen.k,
chosen.d,
chosen.steps,
split_percent(chosen.split_point)
);
dict
}
Plan::FastCover(options) if options.cover.k != 0 && options.cover.d != 0 => {
train_fastcover_dict(&set.corpus, &set.sizes, opts.max_dict, &options, finalize)
.map_err(failed)?
}
Plan::FastCover(options) => {
let (dict, chosen) =
optimize_fastcover_dict(&set.corpus, &set.sizes, opts.max_dict, &options, finalize)
.map_err(failed)?;
display!(
opts.verbosity,
2,
"k={}\nd={}\nf={}\nsteps={}\nsplit={}\naccel={}",
chosen.cover.k,
chosen.cover.d,
chosen.f,
chosen.cover.steps,
split_percent(chosen.cover.split_point),
chosen.accel
);
dict
}
};
let sources = set.sources;
if in_place {
let mut sink = open_in_place(&output)?;
sink.write_all(&dict)
.and_then(|()| sink.flush())
.wrap_err_with(|| format!("failed to write dictionary {}", output.display()))?;
} else {
place_trained_dictionary(&output, &dict, &sources)?;
}
display!(
opts.verbosity,
2,
"trained {} ({}) from {} sample file(s)",
output.display(),
HumanSize::new(dict.len() as u64, false),
opts.inputs.len()
);
Ok(())
}
const TRAINING_SAMPLE_MAX: u64 = 128 << 10;
const TRAINING_DATA_MAX: u64 = 2 << 30;
const TRAINING_SAMPLES_MIN: usize = 5;
struct TrainingSet {
corpus: Vec<u8>,
sizes: Vec<usize>,
sources: Vec<fs::Metadata>,
}
fn shuffle_training_files<T>(files: &mut [T]) {
let mut seed: u32 = 0xFD2F_B528;
let mut next = || {
seed = (seed.wrapping_mul(2_654_435_761) ^ 2_246_822_519).rotate_left(13);
seed >> 5
};
for i in (1..files.len()).rev() {
let j = (next() % (i as u32 + 1)) as usize;
files.swap(i, j);
}
}
fn training_extent(file_sizes: &[u64], block_size: Option<u64>) -> u128 {
let mut samples = 0u128;
for &size in file_sizes {
if size == 0 {
continue;
}
samples += u128::from(block_size.map_or(1, |block| size.div_ceil(block)));
}
samples
}
const TRAINING_SIZE_ENTRY: u64 = core::mem::size_of::<usize>() as u64;
fn plan_training_load(
file_sizes: &[u64],
block_size: Option<u64>,
budget: u64,
limit: u128,
) -> (Vec<u64>, u64) {
let mut plan = vec![0u64; file_sizes.len()];
let mut bytes = 0u64;
let mut entries = 0u64;
'files: for (planned, &size) in plan.iter_mut().zip(file_sizes) {
if u128::from(entries) >= limit {
break;
}
let mut taken = 0u64;
while taken < size {
let piece = match block_size {
Some(block) => (size - taken).min(block),
None => size.min(TRAINING_SAMPLE_MAX),
};
let held = bytes + (entries + 1) * TRAINING_SIZE_ENTRY;
if held > budget || piece > budget - held {
if taken == 0 {
break 'files;
}
break;
}
bytes += piece;
entries += 1;
taken += piece;
*planned += 1;
if block_size.is_none() || u128::from(entries) >= limit {
break;
}
}
}
(plan, bytes)
}
fn load_training_samples(
inputs: &[PathBuf],
block_size: Option<u64>,
memory_limit: Option<u64>,
enough_samples: impl Fn(usize) -> Result<()>,
) -> Result<TrainingSet> {
let mut order: Vec<&PathBuf> = inputs.iter().collect();
shuffle_training_files(&mut order);
let mut file_sizes = Vec::with_capacity(order.len());
for input in &order {
file_sizes.push(
fs::metadata(input)
.wrap_err_with(|| format!("failed to inspect {}", input.display()))?
.len(),
);
}
let samples = training_extent(&file_sizes, block_size);
if samples < TRAINING_SAMPLES_MIN as u128 {
bail!(
"{samples} training sample(s) is too few; provide one file per sample, or \
split files into fixed-size samples with -B#"
);
}
let budget = TRAINING_DATA_MAX.min(memory_limit.unwrap_or(u64::MAX));
let (plan, bytes) = plan_training_load(&file_sizes, block_size, budget, samples);
let retained: u64 = plan.iter().sum();
enough_samples(usize::try_from(retained).unwrap_or(usize::MAX))?;
let corpus_room = usize::try_from(bytes)
.map_err(|_| eyre!("{bytes} bytes of samples is more than this machine can hold"))?;
let entries = usize::try_from(retained)
.map_err(|_| eyre!("{retained} samples is more than this machine can hold"))?;
let mut set = TrainingSet {
corpus: Vec::with_capacity(corpus_room),
sizes: Vec::with_capacity(entries),
sources: Vec::new(),
};
for ((input, &size), &pieces) in order.into_iter().zip(&file_sizes).zip(&plan) {
if pieces == 0 {
continue;
}
let file = File::open(input)
.wrap_err_with(|| format!("failed to open training sample {}", input.display()))?;
let metadata = file
.metadata()
.wrap_err_with(|| format!("failed to inspect {}", input.display()))?;
if metadata.len() != size {
bail!(
"{} changed while it was being read; run again",
input.display()
);
}
let mut reader = file;
let mut taken = 0u64;
for _ in 0..pieces {
let piece = match block_size {
Some(block) => (size - taken).min(block),
None => size.min(TRAINING_SAMPLE_MAX),
};
let before = set.corpus.len();
(&mut reader)
.take(piece)
.read_to_end(&mut set.corpus)
.wrap_err_with(|| format!("failed to read training sample {}", input.display()))?;
if (set.corpus.len() - before) as u64 != piece {
bail!(
"{} changed while it was being read; run again",
input.display()
);
}
set.sizes.push(piece as usize);
taken += piece;
}
set.sources.push(metadata);
}
Ok(set)
}
fn place_trained_dictionary(output: &Path, dict: &[u8], samples: &[fs::Metadata]) -> Result<()> {
let (temp_path, mut temp_file) =
interrupt::create_guarded(|| create_temporary_output_file(output))?;
let placed = (|| {
let sample_permissions = match strictest_sample_permissions(samples, &temp_path) {
Ok(permissions) => permissions,
Err(err) => {
let _ = fs::remove_file(&temp_path);
return Err(err);
}
};
if let Some(permissions) = sample_permissions.clone()
&& let Err(err) = fs::set_permissions(&temp_path, permissions)
{
let _ = fs::remove_file(&temp_path);
return Err(err).wrap_err("failed to apply the samples' permissions to the dictionary");
}
let written = temp_file
.write_all(dict)
.and_then(|()| temp_file.flush())
.wrap_err_with(|| format!("failed to write dictionary {}", output.display()));
if let Err(err) = written {
let _ = fs::remove_file(&temp_path);
return Err(err);
}
drop(temp_file);
replace_output_file(&temp_path, output, sample_permissions)
})();
interrupt::clear();
placed
}
fn cover_options(params: &TrainerParams) -> Result<structured_zstd::dictionary::CoverOptions> {
cover_tuning(params, "--train-cover")
}
fn split_percent(split_point: f64) -> u32 {
(split_point * 100.0).round() as u32
}
fn cover_tuning(
params: &TrainerParams,
flag: &str,
) -> Result<structured_zstd::dictionary::CoverOptions> {
use structured_zstd::dictionary::CoverOptions;
let mut options = if params.listed {
CoverOptions {
k: 0,
d: 0,
steps: 0,
split_point: 0.0,
shrink: None,
}
} else {
CoverOptions::default()
};
if let Some(split) = params.split_percent {
if split > 100 {
bail!("{flag} split is a percentage, got {split}");
}
options.split_point = f64::from(split) / 100.0;
}
if let Some(k) = params.k {
let d = params.d.unwrap_or(8);
if k < d {
bail!("{flag} k must be at least d, got k={k} d={d}");
}
options.k = k;
}
options.d = params.d.unwrap_or(options.d);
options.steps = params.steps.unwrap_or(options.steps);
options.shrink = params.shrink;
Ok(options)
}
fn fastcover_options(
params: &TrainerParams,
) -> Result<structured_zstd::dictionary::FastCoverOptions> {
use structured_zstd::dictionary::FastCoverOptions;
const FLAG: &str = "--train-fastcover";
if let Some(d) = params.d
&& d != 6
&& d != 8
{
bail!("{FLAG} d must be 6 or 8, got {d}");
}
let mut options = FastCoverOptions {
cover: cover_tuning(params, FLAG)?,
..FastCoverOptions::default()
};
if params.listed {
options.f = 0;
options.accel = 0;
} else {
options.cover.split_point = FastCoverOptions::default().cover.split_point;
}
if let Some(f) = params.f {
if f > 31 {
bail!("{FLAG} f must be in 1..=31, got {f}");
}
options.f = f;
}
if let Some(accel) = params.accel {
if accel > 10 {
bail!("{FLAG} accel must be in 1..=10, got {accel}");
}
options.accel = accel;
}
Ok(options)
}
#[cfg(unix)]
fn strictest_sample_permissions(
samples: &[fs::Metadata],
destination: &Path,
) -> Result<Option<fs::Permissions>> {
use std::os::unix::fs::{MetadataExt, PermissionsExt};
if samples.is_empty() {
return Ok(None);
}
let samples_by_group: Vec<(u32, u32)> = samples
.iter()
.map(|sample| (sample.permissions().mode(), sample.gid()))
.collect();
let output_gid = fs::metadata(destination)
.wrap_err_with(|| format!("failed to inspect {}", destination.display()))?
.gid();
Ok(Some(fs::Permissions::from_mode(output_mode_for_sources(
&samples_by_group,
output_gid,
))))
}
#[cfg(unix)]
fn output_mode_for_sources(sources: &[(u32, u32)], output_gid: u32) -> u32 {
let mut mode = 0o0777;
let mut group = None;
let mut one_group = true;
for (source_mode, source_gid) in sources {
mode &= *source_mode;
match group {
None => group = Some(*source_gid),
Some(seen) => one_group &= seen == *source_gid,
}
}
if !one_group || group != Some(output_gid) {
mode &= !0o070;
}
mode
}
#[cfg(not(unix))]
fn strictest_sample_permissions(
_samples: &[fs::Metadata],
_destination: &Path,
) -> Result<Option<fs::Permissions>> {
Ok(None)
}
#[cfg(unix)]
fn permissions_from_source(
source: &fs::Metadata,
destination: &Path,
) -> Result<Option<fs::Permissions>> {
use std::os::unix::fs::{MetadataExt, PermissionsExt};
let output_gid = fs::metadata(destination)
.wrap_err_with(|| format!("failed to inspect {}", destination.display()))?
.gid();
Ok(Some(fs::Permissions::from_mode(output_mode_for_sources(
&[(source.permissions().mode(), source.gid())],
output_gid,
))))
}
#[cfg(not(unix))]
fn permissions_from_source(
source: &fs::Metadata,
_destination: &Path,
) -> Result<Option<fs::Permissions>> {
Ok(Some(source.permissions()))
}
fn read_filling<R: Read>(reader: &mut R, buf: &mut [u8]) -> Result<usize> {
let mut filled = 0;
while filled < buf.len() {
match reader.read(&mut buf[filled..]) {
Ok(0) => break,
Ok(n) => filled += n,
Err(err) if err.kind() == ErrorKind::Interrupted => continue,
Err(err) => return Err(err).wrap_err("failed to read frame header"),
}
}
Ok(filled)
}
fn list_files(opts: &Options) -> Result<usize> {
if opts.inputs.is_empty() {
bail!("No files given");
}
if opts.inputs.iter().any(|input| input == Path::new("-")) {
bail!("--list does not support reading from standard input");
}
let verbose = opts.verbosity > DEFAULT_LEVEL;
if !verbose {
println!("Frames Skips Compressed Uncompressed Ratio Check Filename");
}
let mut total = Some(ListTotal::default());
let mut failed = 0;
for input in &opts.inputs {
match list_file(input, verbose, opts.verbosity) {
Ok(summary) => {
if let Some(sum) = total.as_mut()
&& let Err(err) = sum.add(&summary)
{
display!(opts.verbosity, 1, "zstd: {}: {err}", input.display());
failed += 1;
total = None;
}
}
Err(err) => {
display!(opts.verbosity, 1, "zstd: {err}");
failed += 1;
}
}
}
if let Some(total) = total
&& opts.inputs.len() > 1
&& !verbose
{
total.print();
}
Ok(failed)
}
#[derive(Default)]
struct ListTotal {
frames: u64,
skips: u64,
compressed: u64,
decompressed: u64,
decompressed_unknown: bool,
without_check: bool,
files: usize,
}
impl ListTotal {
fn add(&mut self, summary: &ArchiveSummary) -> Result<()> {
let overflow = || eyre!("the listed files total more than 2^64 bytes");
self.frames = self
.frames
.checked_add(summary.frames)
.ok_or_else(overflow)?;
self.skips = self.skips.checked_add(summary.skips).ok_or_else(overflow)?;
self.compressed = self
.compressed
.checked_add(summary.compressed)
.ok_or_else(overflow)?;
match summary.decompressed {
Some(decompressed) if !self.decompressed_unknown => {
self.decompressed = self
.decompressed
.checked_add(decompressed)
.ok_or_else(overflow)?;
}
Some(_) => {}
None => self.decompressed_unknown = true,
}
self.without_check |= !summary.check;
self.files += 1;
Ok(())
}
fn print(&self) {
println!("----------------------------------------------------------------- ");
let compressed = HumanSize::new(self.compressed, false);
let check = if self.without_check { "" } else { "XXH64" };
if self.decompressed_unknown {
println!(
"{:>6} {:>5} {} {:>5} {} files",
self.frames,
self.skips,
compressed.columns(6),
check,
self.files
);
} else {
let ratio = if self.compressed == 0 {
0.0
} else {
self.decompressed as f64 / self.compressed as f64
};
println!(
"{:>6} {:>5} {} {} {:>5.3} {:>5} {} files",
self.frames,
self.skips,
compressed.columns(6),
HumanSize::new(self.decompressed, false).columns(8),
ratio,
check,
self.files
);
}
}
}
const MAX_FRAME_HEADER_LEN: usize = 18;
#[derive(Debug)]
struct ArchiveSummary {
frames: u64,
skips: u64,
compressed: u64,
decompressed: Option<u64>,
check: bool,
dict_id: Option<u32>,
dict_ids_agree: bool,
window_size: u64,
checksum: Option<[u8; 4]>,
}
fn summarize_archive(path: &Path) -> Result<ArchiveSummary> {
use std::io::{Seek, SeekFrom};
use structured_zstd::decoding::errors::ReadFrameHeaderError;
use structured_zstd::decoding::{FrameContentSize, read_frame_header_info};
let mut file =
File::open(path).wrap_err_with(|| format!("failed to open {}", path.display()))?;
let compressed = file
.metadata()
.wrap_err_with(|| format!("failed to stat {}", path.display()))?
.len();
if compressed == 0 {
bail!("{}: not a zstd frame: empty file", path.display());
}
let mut offset = 0u64;
let mut frames = 0u64;
let mut data_frames = 0u64;
let mut skips = 0u64;
let mut decompressed = Some(0u64);
let mut check = false;
let mut dict_id = None;
let mut dict_ids_agree = true;
let mut window_size = 0u64;
let mut checksum = None;
while offset < compressed {
file.seek(SeekFrom::Start(offset))?;
let mut header_buf = [0u8; MAX_FRAME_HEADER_LEN];
let header_read = read_filling(&mut file, &mut header_buf)?;
let info = match read_frame_header_info(&header_buf[..header_read], false) {
Ok(info) => info,
Err(ReadFrameHeaderError::SkipFrame { length, .. }) => {
offset = offset
.checked_add(8 + u64::from(length))
.filter(|end| *end <= compressed)
.ok_or_else(|| eyre!("{}: truncated skippable frame", path.display()))?;
skips += 1;
frames += 1;
continue;
}
Err(err) => bail!(
"File \"{}\" not compressed by zstd ({err:?})",
path.display()
),
};
let block_size_max = info.window_size.min(128 * 1024);
let mut block_offset = offset + info.header_size as u64;
loop {
file.seek(SeekFrom::Start(block_offset))?;
let mut block_header = [0u8; 3];
file.read_exact(&mut block_header)
.map_err(|_| eyre!("{}: truncated mid-frame", path.display()))?;
let raw = u32::from(block_header[0])
| (u32::from(block_header[1]) << 8)
| (u32::from(block_header[2]) << 16);
let last_block = (raw & 1) != 0;
let block_type = (raw >> 1) & 0b11;
let block_size = u64::from(raw >> 3);
if block_size > block_size_max {
bail!("{}: block exceeds Block_Maximum_Size", path.display());
}
let on_disk = match block_type {
1 => 1,
0 | 2 => block_size,
_ => bail!("{}: reserved block type", path.display()),
};
block_offset = block_offset
.checked_add(3 + on_disk)
.filter(|end| *end <= compressed)
.ok_or_else(|| eyre!("{}: truncated mid-frame", path.display()))?;
if last_block {
break;
}
}
let frame_end = if info.content_checksum {
let end = block_offset
.checked_add(4)
.filter(|end| *end <= compressed)
.ok_or_else(|| eyre!("{}: truncated content checksum", path.display()))?;
file.seek(SeekFrom::Start(block_offset))?;
let mut stored = [0u8; 4];
file.read_exact(&mut stored)
.map_err(|_| eyre!("{}: truncated content checksum", path.display()))?;
checksum = Some(stored);
end
} else {
block_offset
};
window_size = window_size.max(info.window_size);
match info.content_size {
FrameContentSize::Known(n) => {
decompressed = match decompressed {
Some(total) => Some(total.checked_add(n).ok_or_else(|| {
eyre!(
"{}: declared content sizes total more than 2^64 bytes",
path.display()
)
})?),
None => None,
}
}
FrameContentSize::Unknown => decompressed = None,
}
check |= info.content_checksum;
if data_frames == 0 {
dict_id = info.dictionary_id;
} else if dict_id != info.dictionary_id {
dict_ids_agree = false;
}
data_frames += 1;
frames += 1;
offset = frame_end;
}
Ok(ArchiveSummary {
frames,
skips,
compressed,
decompressed,
check,
dict_id: dict_ids_agree.then_some(dict_id).flatten(),
dict_ids_agree,
window_size,
checksum,
})
}
fn list_file(path: &Path, verbose: bool, verbosity: i32) -> Result<ArchiveSummary> {
let metadata =
fs::metadata(path).wrap_err_with(|| format!("Error : {} is not a file", path.display()))?;
if !metadata.is_file() {
bail!("Error : {} is not a file", path.display());
}
let summary = summarize_archive(path)?;
if !summary.dict_ids_agree {
display!(
verbosity,
2,
"WARNING: File contains multiple frames with different dictionary IDs. Showing dictID 0 instead"
);
}
let ratio = summary
.decompressed
.map(|decompressed| decompressed as f64 / summary.compressed as f64);
let check = if summary.check { "XXH64" } else { "None" };
let compressed = HumanSize::new(summary.compressed, false);
if !verbose {
match (summary.decompressed, ratio) {
(Some(decompressed), Some(ratio)) => println!(
"{:>6} {:>5} {} {} {ratio:>5.3} {check:>5} {}",
summary.frames,
summary.skips,
compressed.columns(6),
HumanSize::new(decompressed, false).columns(8),
path.display()
),
_ => println!(
"{:>6} {:>5} {} {check:>5} {}",
summary.frames,
summary.skips,
compressed.columns(6),
path.display()
),
}
return Ok(summary);
}
let data_frames = summary.frames - summary.skips;
println!("{} ", path.display());
println!("# Zstandard Frames: {data_frames}");
if summary.skips > 0 {
println!("# Skippable Frames: {}", summary.skips);
}
println!("DictID: {}", summary.dict_id.unwrap_or(0));
println!(
"Window Size: {} ({} B)",
HumanSize::new(summary.window_size, false),
summary.window_size
);
println!("Compressed Size: {compressed} ({} B)", summary.compressed);
if let (Some(decompressed), Some(ratio)) = (summary.decompressed, ratio) {
println!(
"Decompressed Size: {} ({decompressed} B)",
HumanSize::new(decompressed, false)
);
println!("Ratio: {ratio:.4}");
}
match summary.checksum {
Some(stored) if summary.check && data_frames == 1 => println!(
"Check: {check} {:02x}{:02x}{:02x}{:02x}",
stored[3], stored[2], stored[1], stored[0]
),
_ => println!("Check: {check}"),
}
println!();
Ok(summary)
}
fn remove_source_if_requested(opts: &Options, input: &Path, placement: Placement) -> Result<()> {
if !opts.remove_source || opts.keep || opts.to_stdout {
return Ok(());
}
if placement != Placement::File {
display!(
opts.verbosity,
3,
"Note: {} is not removed: its output is not a regular file",
input.display()
);
return Ok(());
}
interrupt::clear();
fs::remove_file(input).wrap_err("failed to remove source file after success")?;
Ok(())
}
struct CountingWriter<W: Write> {
inner: W,
written: u64,
}
impl<W: Write> Write for CountingWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let written = self.inner.write(buf)?;
self.written += written as u64;
Ok(written)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
}
trait BytesWritten {
fn bytes_written(&self) -> u64;
}
impl<W: Write> BytesWritten for CountingWriter<W> {
fn bytes_written(&self) -> u64 {
self.written
}
}
impl BytesWritten for Vec<u8> {
fn bytes_written(&self) -> u64 {
self.len() as u64
}
}
impl<T: BytesWritten + ?Sized> BytesWritten for &mut T {
fn bytes_written(&self) -> u64 {
(**self).bytes_written()
}
}
fn is_fifo_type(kind: &fs::FileType) -> bool {
#[cfg(unix)]
{
use std::os::unix::fs::FileTypeExt;
kind.is_fifo()
}
#[cfg(not(unix))]
{
let _ = kind;
false
}
}
fn is_block_device_type(kind: &fs::FileType) -> bool {
#[cfg(unix)]
{
use std::os::unix::fs::FileTypeExt;
kind.is_block_device()
}
#[cfg(not(unix))]
{
let _ = kind;
false
}
}
fn open_input(opts: &Options, input: &Path) -> Result<Option<(File, fs::Metadata)>> {
let named = fs::metadata(input)
.map_err(|err| eyre!("can't stat {} : {err} -- ignored", input.display()))?;
if named.is_dir() {
bail!("{} is a directory -- ignored", input.display());
}
let admissible = |kind: &fs::FileType| {
kind.is_file() || is_fifo_type(kind) || (opts.force && is_block_device_type(kind))
};
if !admissible(&named.file_type()) {
bail!("{} is not a regular file -- ignored", input.display());
}
if skips_as_compressed(opts, input) {
display!(
opts.verbosity,
4,
"File is already compressed : {}",
input.display()
);
return Ok(None);
}
let source = File::open(input).wrap_err_with(|| input.display().to_string())?;
let metadata = source
.metadata()
.wrap_err_with(|| format!("failed to inspect {}", input.display()))?;
if metadata.is_dir() || !admissible(&metadata.file_type()) {
bail!("{} is not a regular file -- ignored", input.display());
}
Ok(Some((source, metadata)))
}
fn stream_opened<W: Write>(
opts: &Options,
codecs: &mut Codecs,
source: File,
metadata: &fs::Metadata,
sink: W,
) -> Result<Processed> {
let source_size = metadata.len();
let pledged_size = metadata.is_file().then_some(source_size);
let reader = progress_monitor(
opts,
BufReader::with_capacity(IO_CHUNK, source),
pledged_size,
);
let decode = DecodeSettings::for_file(opts);
stream(opts, codecs, reader, pledged_size, &decode, sink)
}
fn stream_stdin<W: Write>(opts: &Options, codecs: &mut Codecs, sink: W) -> Result<Processed> {
let stdin = io::stdin();
let decode = DecodeSettings::for_stdin(opts);
stream(
opts,
codecs,
stdin_monitor(opts, stdin.lock()),
None,
&decode,
sink,
)
}
fn progress_monitor<R: Read>(opts: &Options, reader: R, total: Option<u64>) -> ProgressMonitor<R> {
let shown = opts.progress.shown(opts.verbosity, opts.console.stderr());
ProgressMonitor::new(reader, total, shown)
}
fn stdin_monitor<R: Read>(opts: &Options, reader: R) -> ProgressMonitor<BufReader<R>> {
progress_monitor(
opts,
BufReader::with_capacity(IO_CHUNK, reader),
opts.pledged_size,
)
}
fn stream<R: BufRead, W: Write>(
opts: &Options,
codecs: &mut Codecs,
mut reader: ProgressMonitor<R>,
pledged_size: Option<u64>,
decode: &DecodeSettings,
mut sink: W,
) -> Result<Processed> {
let written = match opts.mode {
Mode::Compress => compress_stream(
&mut reader,
CountingWriter {
inner: &mut sink,
written: 0,
},
&FrameSettings {
pledged_size: pledged_size.or(opts.pledged_size),
..FrameSettings::from_options(opts)
},
codecs,
)?,
Mode::Decompress => decompress_stream(&mut reader, &mut sink, codecs, decode)?,
Mode::Test => decompress_stream(&mut reader, io::sink(), codecs, decode)?,
Mode::List | Mode::Train => unreachable!("list / train never stream"),
};
sink.flush().wrap_err("failed to flush output")?;
Ok(Processed {
read: reader.read,
written,
})
}
fn stream_input_to<W: Write>(
opts: &Options,
codecs: &mut Codecs,
input: &Path,
mut sink: W,
destination: &str,
total: usize,
) -> Result<Outcome> {
let (name, processed) = if input == Path::new("-") {
(
STDIN_MARK.to_string(),
stream_stdin(opts, codecs, &mut sink)?,
)
} else {
match open_input(opts, input) {
Ok(Some((source, metadata))) => (
input.display().to_string(),
stream_opened(opts, codecs, source, &metadata, &mut sink)?,
),
Ok(None) => return Ok(Outcome::Skipped),
Err(err) => {
display!(opts.verbosity, 1, "zstd: {err}");
return Ok(Outcome::Refused);
}
}
};
file_summary(opts, total, &name, destination, &processed);
Ok(Outcome::Done(processed))
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Placement {
File,
Stream,
Discarded,
}
fn write_output_file<T>(
opts: &Options,
output: &Path,
source: Option<&fs::Metadata>,
fill: impl FnOnce(&mut File) -> Result<T>,
) -> Result<Option<(T, Placement)>> {
write_output_file_if(opts, output, source, fill, |_| true)
}
fn write_output_file_if<T>(
opts: &Options,
output: &Path,
source: Option<&fs::Metadata>,
fill: impl FnOnce(&mut File) -> Result<T>,
keep: impl FnOnce(&T) -> bool,
) -> Result<Option<(T, Placement)>> {
if writes_in_place(output)? {
let mut sink = open_in_place(output)?;
let value = fill(&mut sink)?;
sink.flush().wrap_err("failed to flush output")?;
return Ok(Some((value, Placement::Stream)));
}
ensure_regular_output_destination(output)?;
if output.exists() && !opts.force {
if opts.verbosity <= 1 {
display!(
opts.verbosity,
1,
"zstd: {} already exists; not overwritten",
output.display()
);
return Ok(None);
}
eprint!("zstd: {} already exists; ", output.display());
if !confirm(
"overwrite (y/n) ? ",
"Not overwritten",
reads_stdin(&opts.inputs),
&mut io::stdin().lock(),
) {
return Ok(None);
}
}
let (temp_path, temp_file) =
interrupt::create_guarded(|| create_temporary_output_file(output))?;
let abandon = |temp_path: &Path| {
let _ = fs::remove_file(temp_path);
interrupt::clear();
};
let source_permissions = match source {
Some(metadata) => match permissions_from_source(metadata, &temp_path) {
Ok(permissions) => permissions,
Err(err) => {
abandon(&temp_path);
return Err(err);
}
},
None => None,
};
if let Some(permissions) = source_permissions.clone()
&& let Err(err) = fs::set_permissions(&temp_path, permissions)
{
abandon(&temp_path);
return Err(err).wrap_err("failed to apply the source's permissions to the output");
}
let result: Result<T> = (|| {
let mut sink = temp_file;
let value = fill(&mut sink)?;
sink.flush().wrap_err("failed to flush output")?;
Ok(value)
})();
let value = match result {
Ok(value) => value,
Err(err) => {
abandon(&temp_path);
return Err(err);
}
};
if !keep(&value) {
abandon(&temp_path);
return Ok(Some((value, Placement::Discarded)));
}
let replaced = replace_output_file(&temp_path, output, source_permissions);
interrupt::clear();
replaced.map(|()| Some((value, Placement::File)))
}
fn derive_output_path(opts: &Options, input: &Path) -> Result<PathBuf> {
if let Some(out) = &opts.output {
return Ok(out.clone());
}
let beside = match opts.mode {
Mode::Compress => add_extension(input, ZSTD_SUFFIX),
Mode::Decompress => {
let replacement = input.extension().and_then(|extension| {
DECOMPRESS_SUFFIXES
.iter()
.find(|(known, _)| extension == *known)
.map(|(_, replacement)| *replacement)
});
let Some(replacement) = replacement else {
bail!(
"{}: unknown suffix (.zst/.tzst/.zstd expected). Can't derive the output \
file name. Specify it with -o dstFileName. Ignoring.",
input.display()
);
};
input.with_extension(replacement)
}
Mode::Test | Mode::List | Mode::Train => {
unreachable!("test / list / train modes never write an output file")
}
};
if let Some(root) = &opts.output_dir_mirror {
let verb = if opts.mode == Mode::Compress {
"compress"
} else {
"decompress"
};
let directory = inputs::mirrored_output_dir(input, root).ok_or_else(|| {
eyre!(
"--output-dir-mirror cannot {verb} '{}' into '{}'",
input.display(),
root.display()
)
})?;
return Ok(directory.join(beside.file_name().unwrap_or(beside.as_os_str())));
}
if let Some(directory) = &opts.output_dir {
return Ok(inputs::flat_output_path(&beside, directory));
}
Ok(beside)
}
fn process_file(
opts: &Options,
input: &Path,
codecs: &mut Codecs,
total: usize,
) -> Result<Outcome> {
let Some((source, metadata)) = open_input(opts, input)? else {
return Ok(Outcome::Skipped);
};
let name = input.display().to_string();
if opts.mode == Mode::Test {
let processed = stream_opened(opts, codecs, source, &metadata, io::sink())?;
file_summary(opts, total, &name, "", &processed);
return Ok(Outcome::Done(processed));
}
let output = derive_output_path(opts, input)?;
ensure_distinct_paths(input, &output)?;
if let Some(root) = &opts.output_dir_mirror {
inputs::create_mirrored_dirs(input, root)?;
}
let Some((processed, placement)) = write_output_file(opts, &output, Some(&metadata), |sink| {
stream_opened(opts, codecs, source, &metadata, sink)
})?
else {
return Ok(Outcome::Refused);
};
file_summary(
opts,
total,
&name,
&output.display().to_string(),
&processed,
);
remove_source_if_requested(opts, input, placement)?;
Ok(Outcome::Done(processed))
}
#[derive(Clone, Copy, PartialEq)]
struct FrameSettings {
level: i32,
pledged_size: Option<u64>,
size_hint: Option<u64>,
long: bool,
long_window_log: Option<u32>,
target_block_size: Option<u32>,
checksum: bool,
content_size_flag: bool,
dict_id_flag: bool,
advanced: AdvancedParams,
literals: LiteralCompressionMode,
}
impl Default for FrameSettings {
fn default() -> Self {
Self {
level: CompressionLevel::DEFAULT_LEVEL,
pledged_size: None,
size_hint: None,
long: false,
long_window_log: None,
target_block_size: None,
checksum: true,
content_size_flag: true,
dict_id_flag: true,
advanced: AdvancedParams::default(),
literals: LiteralCompressionMode::Auto,
}
}
}
impl FrameSettings {
fn from_options(opts: &Options) -> Self {
Self {
level: opts.level,
pledged_size: opts.pledged_size,
size_hint: opts.size_hint,
long: opts.long,
long_window_log: opts.long_window_log,
target_block_size: opts.target_block_size,
checksum: opts.checksum,
content_size_flag: opts.content_size_flag,
dict_id_flag: opts.dict_id_flag,
advanced: opts.advanced,
literals: opts.literals,
}
}
}
fn frame_parameters(
level: CompressionLevel,
settings: &FrameSettings,
) -> Result<Option<CompressionParameters>> {
let advanced = settings.advanced;
if !settings.long && advanced.is_default() && settings.literals == LiteralCompressionMode::Auto
{
return Ok(None);
}
let mut builder = CompressionParameters::builder(level);
let window_log = advanced
.window_log
.or(settings.long.then_some(settings.long_window_log).flatten());
if let Some(log) = window_log {
builder = builder.window_log(log);
}
if let Some(log) = advanced.chain_log {
builder = builder.chain_log(log);
}
if let Some(log) = advanced.hash_log {
builder = builder.hash_log(log);
}
if let Some(log) = advanced.search_log {
builder = builder.search_log(log);
}
if let Some(length) = advanced.min_match {
builder = builder.min_match(length);
}
if let Some(length) = advanced.target_length {
builder = builder.target_length(length);
}
if let Some(strategy) = advanced.strategy {
builder = builder.strategy(strategy);
}
if settings.long {
builder = builder.enable_long_distance_matching(true);
if let Some(log) = advanced.ldm_hash_log {
builder = builder.ldm_hash_log(log);
}
if let Some(length) = advanced.ldm_min_match {
builder = builder.ldm_min_match(length);
}
if let Some(log) = advanced.ldm_bucket_size_log {
builder = builder.ldm_bucket_size_log(log);
}
if let Some(log) = advanced.ldm_hash_rate_log {
builder = builder.ldm_hash_rate_log(log);
}
}
builder = builder.literal_compression(settings.literals);
builder
.build()
.map(Some)
.map_err(|err| eyre!("invalid compression parameters: {err}"))
}
#[derive(Clone, Copy)]
struct DecodeSettings {
verify_checksum: bool,
pass_through: bool,
max_window: u64,
memory_overhead: Option<u64>,
}
impl Default for DecodeSettings {
fn default() -> Self {
Self {
verify_checksum: true,
pass_through: false,
max_window: structured_zstd::decoding::MAXIMUM_ALLOWED_WINDOW_SIZE,
memory_overhead: Some(0),
}
}
}
impl DecodeSettings {
fn with_destination(opts: &Options, to_stdout: bool) -> Self {
Self {
verify_checksum: opts.checksum,
pass_through: opts.mode != Mode::Test
&& opts.pass_through.unwrap_or(opts.force && to_stdout),
max_window: opts.decode_window,
memory_overhead: Some(0),
}
}
fn from_options(opts: &Options) -> Self {
Self::with_destination(opts, writes_stdout(opts))
}
fn for_stdin(opts: &Options) -> Self {
Self::with_destination(opts, opts.to_stdout || opts.output.is_none())
}
fn for_file(opts: &Options) -> Self {
Self::with_destination(opts, opts.to_stdout)
}
}
fn new_compressor(
settings: &FrameSettings,
dictionary: Option<&structured_zstd::encoding::EncoderDictionary>,
) -> Result<CompressionContext> {
let &FrameSettings {
level,
size_hint,
target_block_size,
checksum,
content_size_flag,
dict_id_flag,
..
} = settings;
let compression_level = CompressionLevel::from_level(level);
let mut context = CompressionContext::new(compression_level);
context
.set_content_checksum(checksum)
.wrap_err("failed to set the content checksum flag")?;
context
.set_content_size_flag(content_size_flag)
.wrap_err("failed to set the content size flag")?;
context
.set_dictionary_id_flag(dict_id_flag)
.wrap_err("failed to set the dictionary ID flag")?;
if let Some(target) = target_block_size.filter(|target| *target != 0) {
context
.set_target_block_size(Some(target))
.wrap_err("failed to set the block-size target")?;
}
if let Some(params) = frame_parameters(compression_level, settings)? {
context
.set_parameters(¶ms)
.wrap_err("failed to apply the compression parameters")?;
}
if let Some(size) = size_hint.filter(|size| *size != 0) {
context
.set_source_size_hint(size)
.wrap_err("failed to set source size hint")?;
}
if let Some(prepared) = dictionary {
context
.set_encoder_dictionary(prepared.clone())
.wrap_err("failed to load dictionary for compression")?;
}
Ok(context)
}
fn compress_stream<R: BufRead, W: Write + BytesWritten>(
reader: R,
writer: W,
settings: &FrameSettings,
codecs: &mut Codecs,
) -> Result<u64> {
let compressed = compress_frame(reader, writer, settings, codecs);
if compressed.is_err() {
codecs.compressor = None;
}
compressed
}
fn compress_frame<R: BufRead, W: Write + BytesWritten>(
reader: R,
writer: W,
settings: &FrameSettings,
codecs: &mut Codecs,
) -> Result<u64> {
let before = writer.bytes_written();
if codecs.trace.is_none() {
let writer = encode_frame(reader, writer, settings, codecs)?;
return Ok(writer.bytes_written() - before);
}
let started = std::time::Instant::now();
let mut source = TimedReader {
inner: reader,
spent: std::time::Duration::ZERO,
consumed: 0,
};
let timed = TimedWriter {
inner: writer,
spent: std::time::Duration::ZERO,
};
let timed = encode_frame(&mut source, timed, settings, codecs)?;
let consumed = source.consumed;
let duration = started.elapsed() - timed.spent - source.spent;
let written = timed.bytes_written() - before;
let dictionary_size = codecs
.encoder
.as_ref()
.map_or(0, |_| codecs.dictionary_size);
trace_frame(
&mut codecs.trace,
&trace::Frame {
method: trace::Method::Compress,
level: if settings.level == 0 {
structured_zstd::encoding::CompressionLevel::DEFAULT_LEVEL
} else {
settings.level
},
dictionary_size,
uncompressed_size: consumed,
compressed_size: written,
duration,
},
);
Ok(written)
}
struct TimedWriter<W> {
inner: W,
spent: std::time::Duration,
}
impl<W: Write> Write for TimedWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let started = std::time::Instant::now();
let written = self.inner.write(buf);
self.spent += started.elapsed();
written
}
fn flush(&mut self) -> io::Result<()> {
let started = std::time::Instant::now();
let flushed = self.inner.flush();
self.spent += started.elapsed();
flushed
}
}
impl<W: BytesWritten> BytesWritten for TimedWriter<W> {
fn bytes_written(&self) -> u64 {
self.inner.bytes_written()
}
}
struct TimedReader<R> {
inner: R,
spent: std::time::Duration,
consumed: u64,
}
impl<R: Read> Read for TimedReader<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let started = std::time::Instant::now();
let read = self.inner.read(buf);
self.spent += started.elapsed();
if let Ok(n) = read {
self.consumed += n as u64;
}
read
}
}
impl<R: BufRead> BufRead for TimedReader<R> {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
let started = std::time::Instant::now();
let filled = self.inner.fill_buf();
self.spent += started.elapsed();
filled
}
fn consume(&mut self, amount: usize) {
let started = std::time::Instant::now();
self.inner.consume(amount);
self.spent += started.elapsed();
self.consumed += amount as u64;
}
}
fn encode_frame<R: BufRead, W: Write>(
mut reader: R,
writer: W,
settings: &FrameSettings,
codecs: &mut Codecs,
) -> Result<W> {
let mut encoder = StreamingEncoder::with_context(writer, codecs.compressor(settings)?);
if let Some(size) = settings.pledged_size {
encoder
.set_pledged_content_size(size)
.wrap_err("failed to set pledged content size")?;
}
loop {
let piece = match reader.fill_buf() {
Ok(piece) => piece,
Err(err) if err.kind() == io::ErrorKind::Interrupted => continue,
Err(err) => return Err(err).wrap_err("streaming compression failed"),
};
if piece.is_empty() {
break;
}
let len = piece.len();
encoder
.write_all(piece)
.wrap_err("streaming compression failed")?;
reader.consume(len);
}
encoder.finish().wrap_err("failed to finalize zstd frame")
}
const FRAME_MAGIC: u32 = 0xFD2F_B528;
const SKIPPABLE_MAGIC_BASE: u32 = 0x184D_2A50;
fn decompress_stream<R: Read, W: Write>(
reader: R,
mut writer: W,
codecs: &mut Codecs,
settings: &DecodeSettings,
) -> Result<u64> {
let Codecs {
decoder: handle,
decompressor,
chunk,
dictionary_size,
trace,
..
} = codecs;
let dictionary_size = if handle.is_some() {
*dictionary_size
} else {
0
};
decompressor
.set_max_window_size(settings.max_window)
.map_err(|err| eyre!("{err}"))?;
let mut source = BufReader::new(reader);
let mut frames = 0u64;
let mut written = 0u64;
loop {
let mut magic = [0u8; 4];
let read = read_filling(&mut source, &mut magic)?;
if read == 0 {
if frames == 0 && !settings.pass_through {
bail!("unexpected end of file");
}
return Ok(written);
}
let is_frame = read == 4 && {
let magic = u32::from_le_bytes(magic);
magic == FRAME_MAGIC || magic & 0xFFFF_FFF0 == SKIPPABLE_MAGIC_BASE
};
if !is_frame {
if settings.pass_through && frames == 0 {
writer
.write_all(&magic[..read])
.wrap_err("failed to write the passed-through input")?;
written += read as u64;
written += io::copy(&mut source, &mut writer)
.wrap_err("failed to pass the input through")?;
return Ok(written);
}
if read < 4 {
bail!("unknown header");
}
bail!("unsupported format");
}
frames += 1;
let stream = Frames {
magic,
source: &mut source,
handle: handle.as_ref(),
decompressor,
chunk,
trace,
dictionary_size,
settings,
};
written += if stream.trace.is_some() {
decode_frames(
stream,
&mut writer,
&core::cell::Cell::new(Some(std::time::Duration::ZERO)),
)?
} else {
decode_frames(stream, &mut writer, &Untimed)?
};
}
}
struct Frames<'a, R> {
magic: [u8; 4],
source: &'a mut R,
handle: Option<&'a structured_zstd::decoding::DictionaryHandle>,
decompressor: &'a mut structured_zstd::decoding::FrameDecoder,
chunk: &'a mut Vec<u8>,
trace: &'a mut Option<trace::Trace>,
dictionary_size: u64,
settings: &'a DecodeSettings,
}
fn decode_frames<R: Read, W: Write, C: ReadClock>(
frames: Frames<'_, R>,
writer: &mut W,
read_time: &C,
) -> Result<u64> {
use structured_zstd::decoding::StreamingDecoder;
use structured_zstd::decoding::errors::FrameDecoderError;
let Frames {
magic,
source,
handle,
decompressor,
chunk,
trace,
dictionary_size,
settings,
} = frames;
let mut stream = Watched {
inner: io::Cursor::new(magic).chain(source),
failed: false,
read_time,
};
let mut decoder = match handle {
Some(h) => StreamingDecoder::new_with_decoder_and_dictionary_handle(
&mut stream,
&mut *decompressor,
h,
)
.map_err(|err| eyre!("invalid dictionary: {err}"))?,
None => StreamingDecoder::new_with_decoder(&mut stream, &mut *decompressor),
};
decoder
.decoder_mut()
.set_content_checksum(if settings.verify_checksum {
structured_zstd::decoding::ContentChecksum::Verify
} else {
structured_zstd::decoding::ContentChecksum::None
});
chunk.resize(IO_CHUNK, 0);
let err = match drain_decoded(
&mut decoder,
writer,
chunk,
trace,
read_time,
dictionary_size,
) {
Ok(n) => return Ok(n),
Err(err) => err,
};
if decoder.get_ref().failed {
bail!("failed to read the input: {err}");
}
match err
.get_ref()
.and_then(|e| e.downcast_ref::<FrameDecoderError>())
{
Some(&FrameDecoderError::WindowSizeTooBig {
requested, limit, ..
}) => bail!(
"{}",
window_too_large_message(requested, limit, settings.memory_overhead)
),
Some(frame_err) => match after_frame_message(frame_err) {
Some(message) => bail!("{message}"),
None => bail!("invalid zstd frame: {frame_err:?}"),
},
None => Err(err).wrap_err("streaming decompression failed"),
}
}
fn drain_decoded<
R: Read,
D: core::borrow::BorrowMut<structured_zstd::decoding::FrameDecoder>,
C: ReadClock,
>(
decoder: &mut structured_zstd::decoding::StreamingDecoder<R, D>,
writer: &mut impl Write,
buffer: &mut [u8],
trace: &mut Option<trace::Trace>,
read_time: &C,
dictionary_size: u64,
) -> io::Result<u64> {
let mut spent = std::time::Duration::ZERO;
let mut written = 0u64;
let read_sum = || read_time.spent();
let since = |mark: std::time::Instant, read_at_mark: std::time::Duration| {
mark.elapsed() - (read_sum() - read_at_mark)
};
while trace.is_some() {
let mut mark = std::time::Instant::now();
let mut read_at_mark = read_sum();
let read = decoder.read_reporting_frames(buffer, |frame| {
let duration = spent + since(mark, read_at_mark);
spent = std::time::Duration::ZERO;
trace_frame(
trace,
&trace::Frame {
method: trace::Method::Decompress,
level: 0,
dictionary_size,
uncompressed_size: frame.decompressed_size,
compressed_size: frame.compressed_size,
duration,
},
);
if trace.is_none() {
read_time.stop();
}
mark = std::time::Instant::now();
read_at_mark = read_sum();
});
spent += since(mark, read_at_mark);
let n = match read {
Ok(0) => return Ok(written),
Ok(n) => n,
Err(err) if err.kind() == io::ErrorKind::Interrupted => continue,
Err(err) => return Err(err),
};
writer.write_all(&buffer[..n])?;
written += n as u64;
}
read_time.stop();
loop {
let n = match decoder.read(buffer) {
Ok(0) => return Ok(written),
Ok(n) => n,
Err(err) if err.kind() == io::ErrorKind::Interrupted => continue,
Err(err) => return Err(err),
};
writer.write_all(&buffer[..n])?;
written += n as u64;
}
}
trait ReadClock {
fn time<T>(&self, read: impl FnOnce() -> T) -> T;
fn spent(&self) -> std::time::Duration;
fn stop(&self);
}
struct Untimed;
impl ReadClock for Untimed {
#[inline(always)]
fn time<T>(&self, read: impl FnOnce() -> T) -> T {
read()
}
fn spent(&self) -> std::time::Duration {
std::time::Duration::ZERO
}
fn stop(&self) {}
}
impl ReadClock for core::cell::Cell<Option<std::time::Duration>> {
fn time<T>(&self, read: impl FnOnce() -> T) -> T {
match self.get() {
Some(sum) => {
let started = std::time::Instant::now();
let result = read();
self.set(Some(sum + started.elapsed()));
result
}
None => read(),
}
}
fn spent(&self) -> std::time::Duration {
self.get().unwrap_or_default()
}
fn stop(&self) {
self.set(None);
}
}
struct Watched<'a, R, C> {
inner: R,
failed: bool,
read_time: &'a C,
}
impl<R: Read, C: ReadClock> Read for Watched<'_, R, C> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let inner = &mut self.inner;
let result = self.read_time.time(|| inner.read(buf));
if let Err(e) = &result
&& e.kind() != io::ErrorKind::Interrupted
{
self.failed = true;
}
result
}
}
fn window_too_large_message(window: u64, limit: u64, memory_overhead: Option<u64>) -> String {
use structured_zstd::decoding::MAX_DECODER_WINDOW_SIZE;
let head = format!(
"Decoding error (36) : Frame requires too much memory for decoding\n\
Window size larger than maximum : {window} > {limit}"
);
let window_log = window.next_power_of_two().ilog2();
if window <= MAX_DECODER_WINDOW_SIZE {
match memory_overhead.and_then(|overhead| window.checked_add(overhead)) {
Some(memory) => format!(
"{head}\nUse --long={window_log} or --memory={}MB",
memory.div_ceil(1 << 20)
),
None => format!("{head}\nUse --long={window_log}"),
}
} else {
format!(
"{head}\nWindow log larger than ZSTD_WINDOWLOG_MAX={}; not supported",
MAX_DECODER_WINDOW_SIZE.ilog2()
)
}
}
fn after_frame_message(
err: &structured_zstd::decoding::errors::FrameDecoderError,
) -> Option<&'static str> {
use structured_zstd::decoding::errors::{
FrameDecoderError, FrameDescriptorError, ReadFrameHeaderError,
};
match err {
FrameDecoderError::ReadFrameHeaderError(ReadFrameHeaderError::InvalidFrameDescriptor(
FrameDescriptorError::ReservedBitSet,
)) => Some("Decoding error (36) : Unsupported frame parameter"),
FrameDecoderError::ReadFrameHeaderError(ReadFrameHeaderError::BadMagicNumber(_)) => {
Some("unsupported format")
}
FrameDecoderError::ReadFrameHeaderError(ReadFrameHeaderError::MagicNumberReadError(e))
if e.kind() == io::ErrorKind::UnexpectedEof =>
{
Some("unknown header")
}
FrameDecoderError::FailedToSkipFrame => Some("skippable frame is truncated"),
_ => None,
}
}
fn ensure_distinct_paths(input: &Path, output: &Path) -> Result<()> {
let canonical_input = match fs::canonicalize(input) {
Ok(path) => path,
Err(err) if err.kind() == ErrorKind::NotFound => {
return Err(err).wrap_err("failed to open input file");
}
Err(err) => {
return Err(err).wrap_err("failed to canonicalize input file");
}
};
if output.exists() {
let canonical_output =
fs::canonicalize(output).wrap_err("failed to canonicalize existing output file")?;
if canonical_input == canonical_output || paths_point_to_same_file(input, output)? {
return Err(eyre!(
"input and output paths refer to the same file: {input:?} -> {output:?}"
));
}
}
Ok(())
}
fn names_the_same_file(left: &Path, right: &Path) -> Result<bool> {
fn resolved(path: &Path) -> Option<PathBuf> {
if let Ok(whole) = fs::canonicalize(path) {
return Some(whole);
}
let parent = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty());
let directory = fs::canonicalize(parent.unwrap_or_else(|| Path::new("."))).ok()?;
Some(directory.join(path.file_name()?))
}
match (resolved(left), resolved(right)) {
(Some(left), Some(right)) => Ok(left == right),
_ => Ok(false),
}
}
fn paths_point_to_same_file(input: &Path, output: &Path) -> Result<bool> {
fs::metadata(input).wrap_err("failed to inspect input file metadata")?;
fs::metadata(output).wrap_err("failed to inspect existing output file metadata")?;
Ok(identity::of_path(input).is_some_and(|input| identity::of_path(output) == Some(input)))
}
#[cfg(windows)]
fn create_temporary_output_path(output: &Path) -> Result<PathBuf> {
let (path, file) = create_temporary_output_file(output)?;
drop(file);
fs::remove_file(&path).wrap_err("failed to reserve temporary output path")?;
Ok(path)
}
fn create_temporary_output_file(output: &Path) -> Result<(PathBuf, File)> {
let parent = output.parent().unwrap_or_else(|| Path::new("."));
let file_name = output
.file_name()
.and_then(|name| name.to_str())
.unwrap_or("output.zst");
for attempt in 0..u16::MAX {
let candidate = parent.join(format!(
".{file_name}.tmp.{}.{}",
std::process::id(),
attempt
));
match OpenOptions::new()
.create_new(true)
.write(true)
.open(&candidate)
{
Ok(file) => return Ok((candidate, file)),
Err(err) if err.kind() == ErrorKind::AlreadyExists => continue,
Err(err) => {
return Err(err).wrap_err_with(|| output.display().to_string());
}
}
}
Err(eyre!("failed to allocate unique temporary output file"))
}
fn replace_output_file(
temporary_output_path: &Path,
output: &Path,
source_permissions: Option<fs::Permissions>,
) -> Result<()> {
let output_kind = match output_destination_kind(output).inspect_err(|_err| {
let _ = fs::remove_file(temporary_output_path);
})? {
Some(kind) => kind,
None => {
return match fs::rename(temporary_output_path, output) {
Ok(()) => Ok(()),
Err(err) => {
let _ = fs::remove_file(temporary_output_path);
Err(err).wrap_err("failed to move temporary output file into final location")
}
};
}
};
if !output_kind.is_file() {
let _ = fs::remove_file(temporary_output_path);
return Err(eyre!(
"output path exists and is not a regular file: {output:?}"
));
}
let original_permissions = match source_permissions {
Some(permissions) => permissions,
None => fs::metadata(output)
.wrap_err("failed to read existing output file metadata")
.inspect_err(|_err| {
let _ = fs::remove_file(temporary_output_path);
})?
.permissions(),
};
if let Err(err) = fs::set_permissions(temporary_output_path, original_permissions.clone()) {
let _ = fs::remove_file(temporary_output_path);
return Err(err).wrap_err("failed to apply the output's permissions to temporary file");
}
#[cfg(not(windows))]
{
match fs::rename(temporary_output_path, output) {
Ok(()) => Ok(()),
Err(err) => {
let _ = fs::remove_file(temporary_output_path);
Err(err).wrap_err("failed to move temporary output file into final location")
}
}
}
#[cfg(windows)]
{
let backup_output_path = match create_temporary_output_path(output) {
Ok(path) => path,
Err(err) => {
let _ = fs::remove_file(temporary_output_path);
return Err(err).wrap_err("failed to allocate backup output path");
}
};
if let Err(err) = fs::rename(output, &backup_output_path) {
let _ = fs::remove_file(temporary_output_path);
return Err(err).wrap_err("failed to move existing output file into backup location");
}
if let Err(err) = fs::rename(temporary_output_path, output) {
let restore_result = fs::rename(&backup_output_path, output);
let _ = fs::remove_file(temporary_output_path);
if let Err(restore_err) = restore_result {
return Err(err).wrap_err(format!(
"failed to move temporary output file into final location; also failed to restore backup from {backup_output_path:?}: {restore_err}"
));
}
return Err(err).wrap_err("failed to move temporary output file into final location");
}
let _ = fs::remove_file(&backup_output_path);
Ok(())
}
}
fn output_destination_kind(output: &Path) -> Result<Option<std::fs::FileType>> {
match fs::symlink_metadata(output) {
Ok(metadata) => Ok(Some(metadata.file_type())),
Err(err) if err.kind() == ErrorKind::NotFound => Ok(None),
Err(err) => Err(err).wrap_err("failed to inspect existing output path"),
}
}
fn writes_in_place(output: &Path) -> Result<bool> {
match fs::metadata(output) {
Ok(metadata) => Ok(!metadata.is_file() && !metadata.is_dir()),
Err(err) if err.kind() == ErrorKind::NotFound => Ok(false),
Err(err) => Err(err).wrap_err("failed to inspect existing output path"),
}
}
fn open_in_place(output: &Path) -> Result<File> {
OpenOptions::new()
.write(true)
.truncate(true)
.open(output)
.wrap_err_with(|| output.display().to_string())
}
fn ensure_regular_output_destination(output: &Path) -> Result<()> {
if output_destination_kind(output)?.is_some_and(|kind| !kind.is_file()) {
return Err(eyre!(
"output path exists and is not a regular file: {output:?}"
));
}
Ok(())
}
fn add_extension<P: AsRef<Path>>(path: &Path, extension: P) -> PathBuf {
let mut output = path.to_path_buf().into_os_string();
output.push(extension.as_ref().as_os_str());
output.into()
}
#[cfg(test)]
mod tests;