diff --git a/programs/zstdcli.c b/programs/zstdcli.c index 65a0e2fac..1117a2714 100644 --- a/programs/zstdcli.c +++ b/programs/zstdcli.c @@ -24,7 +24,7 @@ int ZSTD_rust_cli_bench(const char* const* fileNames, unsigned nbFiles, int startCLevel, int endCLevel, const ZSTD_compressionParameters* compressionParams, int displayLevel, unsigned nbSeconds, - size_t blockSize, int nbWorkers); + size_t blockSize, int nbWorkers, int mode); const char* ZSTD_rust_cli_expected_version(void) { @@ -41,7 +41,7 @@ int ZSTD_rust_cli_bench(const char* const* fileNames, unsigned nbFiles, int startCLevel, int endCLevel, const ZSTD_compressionParameters* compressionParams, int displayLevel, unsigned nbSeconds, - size_t blockSize, int nbWorkers) + size_t blockSize, int nbWorkers, int mode) { #ifndef ZSTD_NOBENCH BMK_advancedParams_t advancedParams = BMK_initAdvancedParams(); @@ -50,6 +50,11 @@ int ZSTD_rust_cli_bench(const char* const* fileNames, unsigned nbFiles, advancedParams.nbSeconds = nbSeconds; advancedParams.blockSize = blockSize; advancedParams.nbWorkers = nbWorkers; + advancedParams.mode = (BMK_mode_t)mode; + if (advancedParams.mode == BMK_decodeOnly) { + startLevel = 0; + endLevel = 0; + } if (startLevel > ZSTD_maxCLevel()) startLevel = ZSTD_maxCLevel(); if (endLevel > ZSTD_maxCLevel()) endLevel = ZSTD_maxCLevel(); if (endLevel < startLevel) endLevel = startLevel; @@ -67,6 +72,7 @@ int ZSTD_rust_cli_bench(const char* const* fileNames, unsigned nbFiles, (void)fileNames; (void)nbFiles; (void)dictFileName; (void)startCLevel; (void)endCLevel; (void)compressionParams; (void)displayLevel; (void)nbSeconds; (void)blockSize; (void)nbWorkers; + (void)mode; return -1; #endif } diff --git a/rust/src/zstd_cli.rs b/rust/src/zstd_cli.rs index 6095cabe5..c642f4a45 100644 --- a/rust/src/zstd_cli.rs +++ b/rust/src/zstd_cli.rs @@ -22,8 +22,9 @@ //! byte-identical with the C CLI. //! //! Remaining C-only CLI boundaries are called out in `unsupported()` below: -//! tracing and alternate-format selection. Dictionary training uses the -//! temporary `DiB_trainFromFiles` bridge while its algorithms are Rust-owned. +//! a few compatibility-only options. Trace logging uses the existing C trace +//! backend through a narrow lifecycle bridge, while dictionary training uses +//! the temporary `DiB_trainFromFiles` bridge and Rust-owned algorithms. use std::env; use std::ffi::{CStr, CString, OsStr, OsString}; @@ -56,10 +57,17 @@ const STDOUT_MARK: &str = "/*stdout*\\"; const NULL_MARK: &str = "NUL"; #[cfg(not(windows))] const NULL_MARK: &str = "/dev/null"; -#[cfg(feature = "compression")] const ZSTD_SUFFIX: &[u8] = b".zst\0"; +const GZ_SUFFIX: &[u8] = b".gz\0"; +const XZ_SUFFIX: &[u8] = b".xz\0"; +const LZMA_SUFFIX: &[u8] = b".lzma\0"; +const LZ4_SUFFIX: &[u8] = b".lz4\0"; const FIO_ZSTD_COMPRESSION: c_int = 0; +const FIO_GZIP_COMPRESSION: c_int = 1; +const FIO_XZ_COMPRESSION: c_int = 2; +const FIO_LZMA_COMPRESSION: c_int = 3; +const FIO_LZ4_COMPRESSION: c_int = 4; const FIO_PS_AUTO: c_int = 0; const FIO_PS_NEVER: c_int = 1; const FIO_PS_ALWAYS: c_int = 2; @@ -212,6 +220,8 @@ unsafe extern "C" { fn FIO_setNbFilesTotal(ctx: *mut FIO_ctx_t, value: c_int); fn FIO_setHasStdinInput(ctx: *mut FIO_ctx_t, value: c_int); fn FIO_setHasStdoutOutput(ctx: *mut FIO_ctx_t, value: c_int); + fn TRACE_enable(filename: *const c_char); + fn TRACE_finish(); #[cfg(feature = "compression")] fn FIO_compressFilename( @@ -275,6 +285,7 @@ unsafe extern "C" { nb_seconds: c_uint, block_size: usize, nb_workers: c_int, + mode: c_int, ) -> c_int; /// Narrow bridge to `programs/dibio.c`. The file loader and dictionary /// algorithms remain on the C/Rust library side of this boundary. @@ -322,6 +333,48 @@ enum Operation { Train, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum CompressionFormat { + Zstd, + Gzip, + Xz, + Lzma, + Lz4, +} + +impl CompressionFormat { + fn parse(value: &str) -> Result { + match value { + "zstd" => Ok(Self::Zstd), + "gzip" => Ok(Self::Gzip), + "xz" => Ok(Self::Xz), + "lzma" => Ok(Self::Lzma), + "lz4" => Ok(Self::Lz4), + _ => Err(format!("unknown compression format {value:?}")), + } + } + + const fn fio_type(self) -> c_int { + match self { + Self::Zstd => FIO_ZSTD_COMPRESSION, + Self::Gzip => FIO_GZIP_COMPRESSION, + Self::Xz => FIO_XZ_COMPRESSION, + Self::Lzma => FIO_LZMA_COMPRESSION, + Self::Lz4 => FIO_LZ4_COMPRESSION, + } + } + + const fn suffix(self) -> &'static [u8] { + match self { + Self::Zstd => ZSTD_SUFFIX, + Self::Gzip => GZ_SUFFIX, + Self::Xz => XZ_SUFFIX, + Self::Lzma => LZMA_SUFFIX, + Self::Lz4 => LZ4_SUFFIX, + } + } +} + #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] enum TrainingAlgorithm { Cover, @@ -340,6 +393,8 @@ enum Action { #[derive(Debug)] struct Cli { operation: Operation, + bench_decode: bool, + compression_format: CompressionFormat, inputs: Vec, output: Option, dictionary: Option, @@ -359,6 +414,7 @@ struct Cli { progress: i32, workers: Option, single_thread: bool, + auto_threads_logical: bool, block_size: Option, mem_limit: Option, ldm: bool, @@ -389,15 +445,17 @@ struct Cli { max_dict_size: usize, dictionary_id: u32, dictionary_selectivity: u32, + trace: Option, cover_params: ZDICT_cover_params_t, fast_cover_params: ZDICT_fastCover_params_t, - unsupported_program: Option, } impl Cli { fn new(program_name: &str) -> Self { let mut cli = Self { operation: Operation::Compress, + bench_decode: false, + compression_format: CompressionFormat::Zstd, inputs: Vec::new(), output: None, dictionary: None, @@ -417,6 +475,7 @@ impl Cli { progress: FIO_PS_AUTO, workers: None, single_thread: false, + auto_threads_logical: false, block_size: None, mem_limit: None, ldm: false, @@ -447,9 +506,9 @@ impl Cli { max_dict_size: DEFAULT_MAX_DICT_SIZE, dictionary_id: 0, dictionary_selectivity: DEFAULT_DICT_SELECTIVITY, + trace: None, cover_params: ZDICT_cover_params_t::default(), fast_cover_params: default_fast_cover_params(), - unsupported_program: None, }; match program_name { @@ -463,8 +522,46 @@ impl Cli { cli.pass_through = Some(1); cli.display_level = 1; } - "gzip" | "gunzip" | "gzcat" | "lzma" | "unlzma" | "xz" | "unxz" | "lz4" | "unlz4" => { - cli.unsupported_program = Some(program_name.to_owned()) + "gzip" => { + cli.compression_format = CompressionFormat::Gzip; + cli.remove_source = true; + } + "gunzip" => { + cli.compression_format = CompressionFormat::Gzip; + cli.operation = Operation::Decompress; + cli.remove_source = true; + } + "gzcat" => { + cli.compression_format = CompressionFormat::Gzip; + cli.operation = Operation::Decompress; + cli.output = Some(cstring(STDOUT_MARK).expect("static stdout marker")); + cli.force = true; + cli.force_stdout = true; + cli.pass_through = Some(1); + cli.display_level = 1; + } + "xz" => { + cli.compression_format = CompressionFormat::Xz; + cli.remove_source = true; + } + "unxz" => { + cli.compression_format = CompressionFormat::Xz; + cli.operation = Operation::Decompress; + cli.remove_source = true; + } + "lzma" => { + cli.compression_format = CompressionFormat::Lzma; + cli.remove_source = true; + } + "unlzma" => { + cli.compression_format = CompressionFormat::Lzma; + cli.operation = Operation::Decompress; + cli.remove_source = true; + } + "lz4" => cli.compression_format = CompressionFormat::Lz4, + "unlz4" => { + cli.compression_format = CompressionFormat::Lz4; + cli.operation = Operation::Decompress; } _ => {} } @@ -508,11 +605,16 @@ unsafe fn default_worker_count() -> i32 { } #[cfg(feature = "compression")] -unsafe fn resolved_worker_count(workers: Option, single_thread: bool) -> i32 { +unsafe fn resolved_worker_count( + workers: Option, + single_thread: bool, + auto_threads_logical: bool, +) -> i32 { match workers { /* --single-thread pins zero workers; a bare zero (-T0 or the zstdmt * program name) auto-detects the core count as in the C CLI. */ Some(0) if single_thread => 0, + Some(0) if auto_threads_logical => unsafe { UTIL_countLogicalCores() }.max(1), Some(0) => unsafe { UTIL_countPhysicalCores() }.max(1), Some(workers) => workers, None => unsafe { default_worker_count() }, @@ -564,7 +666,7 @@ fn usage(advanced: bool) { let _ = writeln!(out, "\nImplemented advanced compression controls:"); let _ = writeln!( out, - " --fast[=#], --ultra, --long[=#], --threads=#, --single-thread, --block-size=#" + " --fast[=#], --ultra, --long[=#], --threads=#, --single-thread, --auto-threads={{physical|logical}}, --block-size=#" ); let _ = writeln!( out, @@ -617,7 +719,14 @@ fn usage(advanced: bool) { out, "\nDictionary builder:\n --train Create a dictionary from training files\n --train-cover[=k=#,d=#,steps=#,split=#,shrink[=#]]\n --train-fastcover[=k=#,d=#,f=#,steps=#,split=#,accel=#,shrink[=#]]\n --train-legacy[=s=#] Use the legacy algorithm\n --maxdict=# Maximum dictionary size (default {DEFAULT_MAX_DICT_SIZE})\n --dictID=# Force the dictionary ID (default: random)" ); - let _ = writeln!(out, "\nNot yet migrated: trace and alternate formats."); + let _ = writeln!( + out, + "\n --format=zstd|gzip|xz|lzma|lz4 Select the compression frame format" + ); + let _ = writeln!( + out, + " --trace LOG Log compression/decompression trace data" + ); } } @@ -1000,7 +1109,10 @@ fn parse_long_option( Ok(None) } "--decompress" | "--uncompress" => { - cli.operation = Operation::Decompress; + cli.bench_decode = true; + if cli.operation != Operation::Bench { + cli.operation = Operation::Decompress; + } Ok(None) } "--test" => { @@ -1156,6 +1268,15 @@ fn parse_long_option( cli.workers = Some(parse_worker_count(&value)?); Ok(None) } + "--auto-threads" => { + let value = next_field(attached, args, index, name)?; + cli.auto_threads_logical = match value.as_str() { + "physical" => false, + "logical" => true, + _ => return Err(format!("unknown auto-thread mode {value:?}")), + }; + Ok(None) + } "--single-thread" => { /* As in the C CLI: zero workers plus a latch that suppresses the * automatic core-count resolution, so fileio runs its @@ -1241,6 +1362,16 @@ fn parse_long_option( cli.dictionary_id = parse_u32(&value, "dictionary ID")?; Ok(None) } + "--format" => { + let value = next_field(attached, args, index, name)?; + cli.compression_format = CompressionFormat::parse(&value)?; + Ok(None) + } + "--trace" => { + let value = next_field(attached, args, index, name)?; + cli.trace = Some(cstring(&value)?); + Ok(None) + } "--filelist" => { let value = next_field(attached, args, index, name)?; cli.file_lists.push(cstring(&value)?); @@ -1273,10 +1404,7 @@ fn parse_long_option( } "--max" | "--patch-from" - | "--trace" - | "--format" | "--priority" - | "--auto-threads" | "--fake-stdin-is-console" | "--fake-stdout-is-console" | "--fake-stderr-is-console" @@ -1309,7 +1437,12 @@ fn parse_short_options( offset = digits_end; continue; } - 'd' => cli.operation = Operation::Decompress, + 'd' => { + cli.bench_decode = true; + if cli.operation != Operation::Bench { + cli.operation = Operation::Decompress; + } + } 'z' => cli.operation = Operation::Compress, 't' => cli.operation = Operation::Test, 'b' => cli.operation = Operation::Bench, @@ -1377,7 +1510,20 @@ fn parse_short_options( } break; } - 'p' | 'P' | 's' | 'S' => { + 's' => { + let mut digits_end = offset + 1; + while digits_end < bytes.len() && bytes[digits_end].is_ascii_digit() { + digits_end += 1; + } + if digits_end == offset + 1 { + return Err("missing argument for -s".to_owned()); + } + cli.dictionary_selectivity = + parse_u32(&value[offset + 1..digits_end], "dictionary selectivity")?; + offset = digits_end; + continue; + } + 'p' | 'P' | 'S' => { unsupported(&format!("-{option}"))?; } _ => return Err(format!("unknown option -{option}")), @@ -1429,7 +1575,7 @@ unsafe fn apply_preferences(cli: &Cli, prefs: *mut FIO_prefs_t, ctx: *mut FIO_ct }; unsafe { g_utilDisplayLevel = display_level; - FIO_setCompressionType(prefs, FIO_ZSTD_COMPRESSION); + FIO_setCompressionType(prefs, cli.compression_format.fio_type()); FIO_setNotificationLevel(display_level); FIO_setProgressSetting( if !io::stderr().is_terminal() && cli.progress != FIO_PS_ALWAYS { @@ -1462,7 +1608,10 @@ unsafe fn apply_preferences(cli: &Cli, prefs: *mut FIO_prefs_t, ctx: *mut FIO_ct }), ); #[cfg(feature = "compression")] - FIO_setNbWorkers(prefs, resolved_worker_count(cli.workers, cli.single_thread)); + FIO_setNbWorkers( + prefs, + resolved_worker_count(cli.workers, cli.single_thread, cli.auto_threads_logical), + ); FIO_setLdmFlag(prefs, u32::from(cli.ldm)); FIO_setAdaptiveMode(prefs, i32::from(cli.adapt)); FIO_setRsyncable(prefs, i32::from(cli.rsyncable)); @@ -1694,7 +1843,7 @@ unsafe fn run_compress( output_dir_mirror(cli), output_dir_flat(cli), output, - ZSTD_SUFFIX.as_ptr().cast(), + cli.compression_format.suffix().as_ptr().cast(), dictionary, cli.level, cli.compression_params, @@ -1826,6 +1975,18 @@ fn run_bench(cli: &Cli) -> Result { .dictionary .as_ref() .map_or(ptr::null(), |value| value.as_ptr()); + let workers = { + #[cfg(feature = "compression")] + { + unsafe { + resolved_worker_count(cli.workers, cli.single_thread, cli.auto_threads_logical) + } + } + #[cfg(not(feature = "compression"))] + { + cli.workers.unwrap_or(1) + } + }; let result = unsafe { ZSTD_rust_cli_bench( inputs.as_ptr(), @@ -1838,7 +1999,8 @@ fn run_bench(cli: &Cli) -> Result { cli.bench_nb_seconds.unwrap_or(DEFAULT_BENCH_NB_SECONDS), cli.block_size.unwrap_or(0), // The C CLI benchmarks single-threaded unless -T was given. - cli.workers.unwrap_or(1), + workers, + if cli.bench_decode { 1 } else { 0 }, ) }; if result < 0 { @@ -1862,7 +2024,9 @@ fn run_train(cli: &Cli) -> Result { let default_output = cstring(DEFAULT_DICT_NAME).expect("static dictionary name"); let output = cli.output.as_ref().unwrap_or(&default_output); let inputs: Vec<*const c_char> = cli.inputs.iter().map(|value| value.as_ptr()).collect(); - let workers = unsafe { resolved_worker_count(cli.workers, cli.single_thread) } as c_uint; + let workers = unsafe { + resolved_worker_count(cli.workers, cli.single_thread, cli.auto_threads_logical) + } as c_uint; let z_params = ZDICT_params_t { compressionLevel: cli.level, notificationLevel: cli.display_level as c_uint, @@ -1936,12 +2100,31 @@ fn run_train(cli: &Cli) -> Result { } } -fn run_cli(mut cli: Cli) -> Result { - if let Some(program_name) = &cli.unsupported_program { - return Err(format!( - "{program_name} compatibility mode is not yet implemented by the Rust CLI frontend" - )); +struct TraceGuard { + enabled: bool, +} + +impl TraceGuard { + fn new(filename: Option<&CString>) -> Self { + if let Some(filename) = filename { + unsafe { TRACE_enable(filename.as_ptr()) }; + Self { enabled: true } + } else { + Self { enabled: false } + } } +} + +impl Drop for TraceGuard { + fn drop(&mut self) { + if self.enabled { + unsafe { TRACE_finish() }; + } + } +} + +fn run_cli(mut cli: Cli) -> Result { + let _trace_guard = TraceGuard::new(cli.trace.as_ref()); unsafe { /* The C CLI publishes the display level to util.c before any table * expansion, so traversal warnings obey -q/-v. */ @@ -1966,6 +2149,9 @@ fn run_cli(mut cli: Cli) -> Result { return Err("file information is not supported".to_owned()); } if cli.operation == Operation::Bench { + if cli.compression_format != CompressionFormat::Zstd { + return Err("benchmark mode only supports the zstd format".to_owned()); + } return run_bench(&cli); } if cli.operation == Operation::Train { @@ -2186,6 +2372,46 @@ mod tests { assert_eq!(cli.content_size, 1); assert_eq!(cli.workers, None); assert_eq!(cli.operation, Operation::Compress); + assert_eq!(cli.compression_format, CompressionFormat::Zstd); + } + + #[test] + fn format_option_selects_frame_type_and_suffix() { + let cli = parse(&["zstd", "--format=gzip", "input"]); + + assert_eq!(cli.compression_format, CompressionFormat::Gzip); + assert_eq!(cli.compression_format.fio_type(), FIO_GZIP_COMPRESSION); + assert_eq!(cli.compression_format.suffix(), b".gz\0"); + + let cli = parse(&["zstd", "--format=gzip", "--format=zstd", "input"]); + assert_eq!(cli.compression_format, CompressionFormat::Zstd); + } + + #[test] + fn format_option_rejects_unknown_formats() { + let error = parse_args(vec![ + OsString::from("zstd"), + OsString::from("--format=zip"), + OsString::from("input"), + ]) + .expect_err("unknown formats must be rejected"); + + assert!(error.contains("unknown compression format")); + } + + #[test] + fn trace_option_accepts_attached_and_separate_paths() { + let cli = parse(&["zstd", "--trace=first.trace", "input"]); + assert_eq!( + cli.trace.as_deref().map(CStr::to_bytes), + Some(&b"first.trace"[..]) + ); + + let cli = parse(&["zstd", "--trace", "second.trace", "input"]); + assert_eq!( + cli.trace.as_deref().map(CStr::to_bytes), + Some(&b"second.trace"[..]) + ); } #[test] @@ -2276,6 +2502,28 @@ mod tests { assert!(!cli.single_thread); } + #[test] + fn auto_threads_selects_core_family() { + let physical = parse(&["zstd", "-T0", "--auto-threads=physical", "input"]); + assert_eq!(physical.workers, Some(0)); + assert!(!physical.auto_threads_logical); + + let logical = parse(&["zstd", "-T0", "--auto-threads", "logical", "input"]); + assert_eq!(logical.workers, Some(0)); + assert!(logical.auto_threads_logical); + } + + #[test] + fn auto_threads_rejects_unknown_core_family() { + let error = parse_args(vec![ + OsString::from("zstd"), + OsString::from("--auto-threads=bogus"), + ]) + .expect_err("an unknown core family must be rejected"); + + assert!(error.contains("unknown auto-thread mode")); + } + #[test] fn single_thread_pins_zero_workers() { let cli = parse(&["zstd", "--single-thread", "input"]); @@ -2306,10 +2554,20 @@ mod tests { } #[test] - fn alternate_format_aliases_fail_before_processing_files() { + fn alternate_format_aliases_select_compatibility_modes() { let cli = Cli::new("gzip"); + assert_eq!(cli.compression_format, CompressionFormat::Gzip); + assert_eq!(cli.operation, Operation::Compress); + assert!(cli.remove_source); - assert_eq!(cli.unsupported_program.as_deref(), Some("gzip")); + let cli = Cli::new("unxz"); + assert_eq!(cli.compression_format, CompressionFormat::Xz); + assert_eq!(cli.operation, Operation::Decompress); + assert!(cli.remove_source); + + let cli = Cli::new("lzma"); + assert_eq!(cli.compression_format, CompressionFormat::Lzma); + assert_eq!(cli.operation, Operation::Compress); } #[cfg(unix)] @@ -2405,6 +2663,10 @@ mod tests { assert_eq!(legacy.operation, Operation::Train); assert_eq!(legacy.training_algorithm, TrainingAlgorithm::Legacy); assert_eq!(legacy.dictionary_selectivity, 12); + + let legacy_short = parse(&["zstd", "--train-legacy", "-s5", "sample"]); + assert_eq!(legacy_short.training_algorithm, TrainingAlgorithm::Legacy); + assert_eq!(legacy_short.dictionary_selectivity, 5); } #[test] @@ -2573,4 +2835,12 @@ mod tests { assert_eq!(cli.bench_end_level, Some(0)); assert_eq!(cli.bench_nb_seconds, Some(0)); } + + #[test] + fn bench_decompression_keeps_benchmark_operation() { + let cli = parse(&["zstd", "-b", "-d", "-i0", "input"]); + + assert_eq!(cli.operation, Operation::Bench); + assert!(cli.bench_decode); + } }