From e11622202bdf12e849b1764d3502e088dd782096 Mon Sep 17 00:00:00 2001 From: Dennis Ranke Date: Thu, 24 Aug 2023 00:00:20 +0200 Subject: [PATCH] add support to read/write from/to stdin/stdout --- src/main.rs | 130 +++++++++++++++++++++++++++++++++++----------------- 1 file changed, 89 insertions(+), 41 deletions(-) diff --git a/src/main.rs b/src/main.rs index 24ceb5a..cfeeaaf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -68,7 +68,7 @@ fn main() -> Result<()> { Short(n) if n.is_ascii_digit() => level = n as u8 - b'0', Short('h') | Long("help") => print_help(0), Long("version") => { - println!("{}", env!("CARGO_PKG_VERSION")); + eprintln!("{}", env!("CARGO_PKG_VERSION")); process::exit(0); } Long("max-unpacked-size") => max_unpacked_size = parser.value()?.parse()?, @@ -78,38 +78,8 @@ fn main() -> Result<()> { } } - let infile = infile.unwrap_or_else(|| print_help(1)); - enum OutFileType { - Packed, - Unpacked, - Heatmap, - } - let outfile = |tpe: OutFileType| { - outfile.clone().unwrap_or_else(|| { - let mut name = infile.clone(); - match tpe { - OutFileType::Packed => { - let mut filename = name - .file_name() - .unwrap_or_else(|| OsStr::new("")) - .to_os_string(); - filename.push(".upk"); - name.set_file_name(filename); - } - OutFileType::Unpacked => { - if name.extension().filter(|&e| e == "upk").is_some() { - name.set_extension(""); - } else { - name.set_extension("bin"); - } - } - OutFileType::Heatmap => { - name.set_extension("heatmap"); - } - } - name - }) - }; + let infile = IoTarget::from_filename(infile); + let outfile = |tpe: OutFileType| infile.output(tpe, &outfile); if config.parity_contexts != 1 && config.parity_contexts != 2 && config.parity_contexts != 4 { eprintln!("--parity has to be 1, 2, or 4"); @@ -117,8 +87,7 @@ fn main() -> Result<()> { } if !unpack && !calculate_margin && !create_heatmap { - let mut data = vec![]; - File::open(&infile)?.read_to_end(&mut data)?; + let mut data = infile.read()?; if reverse { data.reverse(); } @@ -145,16 +114,15 @@ fn main() -> Result<()> { packed_data.reverse(); } - println!( + eprintln!( "Compressed {} bytes to {} bytes ({}%)", data.len(), packed_data.len(), packed_data.len() as f32 * 100. / data.len() as f32 ); - File::create(outfile(OutFileType::Packed))?.write_all(&packed_data)?; + outfile(OutFileType::Packed).write(&packed_data)?; } else { - let mut data = vec![]; - File::open(&infile)?.read_to_end(&mut data)?; + let mut data = infile.read()?; if reverse { data.reverse(); } @@ -163,7 +131,7 @@ fn main() -> Result<()> { if reverse { unpacked_data.reverse(); } - File::create(outfile(OutFileType::Unpacked))?.write_all(&unpacked_data)?; + outfile(OutFileType::Unpacked).write(&unpacked_data)?; } if create_heatmap { let mut heatmap = upkr::create_heatmap(&data, &config, max_unpacked_size)?; @@ -182,7 +150,7 @@ fn main() -> Result<()> { .min(127.) as u8; heatmap_bin.push((cost << 1) | heatmap.is_literal(i) as u8); } - File::create(outfile(OutFileType::Heatmap))?.write_all(&heatmap_bin)?; + outfile(OutFileType::Heatmap).write(&heatmap_bin)?; } } } @@ -194,6 +162,81 @@ fn main() -> Result<()> { Ok(()) } +enum OutFileType { + Packed, + Unpacked, + Heatmap, +} + +enum IoTarget { + StdInOut, + File(PathBuf), +} + +impl IoTarget { + fn from_filename(filename: Option) -> IoTarget { + if let Some(path) = filename { + if path.as_os_str() == "-" { + IoTarget::StdInOut + } else { + IoTarget::File(path) + } + } else { + IoTarget::StdInOut + } + } + + fn read(&self) -> Result> { + let mut buffer = vec![]; + match *self { + IoTarget::StdInOut => std::io::stdin().read_to_end(&mut buffer)?, + IoTarget::File(ref path) => File::open(path)?.read_to_end(&mut buffer)?, + }; + Ok(buffer) + } + + fn write(&self, data: &[u8]) -> Result<()> { + match *self { + IoTarget::StdInOut => std::io::stdout().write_all(data)?, + IoTarget::File(ref path) => File::create(path)?.write_all(data)?, + }; + Ok(()) + } + + fn output(&self, tpe: OutFileType, outname: &Option) -> IoTarget { + if outname.is_some() { + return IoTarget::from_filename(outname.clone()); + } + match *self { + IoTarget::StdInOut => IoTarget::StdInOut, + IoTarget::File(ref path) => { + let mut name = path.clone(); + match tpe { + OutFileType::Packed => { + let mut filename = name + .file_name() + .unwrap_or_else(|| OsStr::new("")) + .to_os_string(); + filename.push(".upk"); + name.set_file_name(filename); + } + OutFileType::Unpacked => { + if name.extension().filter(|&e| e == "upk").is_some() { + name.set_extension(""); + } else { + name.set_extension("bin"); + } + } + OutFileType::Heatmap => { + name.set_extension("heatmap"); + } + } + IoTarget::File(name) + } + } + } +} + fn print_help(exit_code: i32) -> ! { eprintln!("Usage:"); eprintln!(" upkr [-l level(0-9)] [config options] []"); @@ -207,6 +250,11 @@ fn print_help(exit_code: i32) -> ! { eprintln!(" --heatmap calculate heatmap from compressed file"); eprintln!(" --margin calculate margin for overlapped unpacking of a packed file"); eprintln!(); + eprintln!("When no infile is given, or the infile is '-', read from stdin."); + eprintln!( + "When no outfile is given and reading from stdin, or when outfile is '-', write to stdout." + ); + eprintln!(); eprintln!("Version: {}", env!("CARGO_PKG_VERSION")); eprintln!(); eprintln!("Config presets for specific unpackers:");