mirror of
https://github.com/lifegpc/msg-tool.git
synced 2026-06-13 16:38:52 +08:00
Compare commits
10 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| fed1f2d14b | |||
| ff526d600e | |||
| 0860ceb8a0 | |||
| 3dbfaa866b | |||
| 5fb7a8b601 | |||
| 199442ac6d | |||
| 2e7cd8119e | |||
| 4391ad6de5 | |||
| 01477b31a0 | |||
| e5e437b1e3 |
44
Cargo.lock
generated
44
Cargo.lock
generated
@@ -226,15 +226,6 @@ dependencies = [
|
||||
"clap_derive 4.5.47",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "clap-num"
|
||||
version = "1.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "822c4000301ac390e65995c62207501e3ef800a1fc441df913a5e8e4dc374816"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "clap_builder"
|
||||
version = "4.5.47"
|
||||
@@ -1145,14 +1136,32 @@ dependencies = [
|
||||
"nasm-rs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "msg-tool-jpegxl-src"
|
||||
version = "0.11.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d5b6b09d7b013a95614b62ef2bf668b268853d9a3d90b93ada715f2007c815b8"
|
||||
dependencies = [
|
||||
"cmake",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "msg-tool-jpegxl-sys"
|
||||
version = "0.11.2+libjxl-0.11.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "eb3d4c72bd5bb5a7d6d9e902365bae80088479c05e651947bc285cdfd384f201"
|
||||
dependencies = [
|
||||
"msg-tool-jpegxl-src",
|
||||
"pkg-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "msg_tool"
|
||||
version = "0.2.3"
|
||||
version = "0.2.4"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"byteorder",
|
||||
"clap 4.5.47",
|
||||
"clap-num",
|
||||
"csv",
|
||||
"ctrlc",
|
||||
"emote-psb",
|
||||
@@ -1168,7 +1177,9 @@ dependencies = [
|
||||
"markup5ever_rcdom",
|
||||
"memchr",
|
||||
"mozjpeg",
|
||||
"msg-tool-jpegxl-sys",
|
||||
"msg_tool_macro",
|
||||
"num_cpus",
|
||||
"overf",
|
||||
"parse-size",
|
||||
"pelite",
|
||||
@@ -1184,7 +1195,7 @@ dependencies = [
|
||||
"url",
|
||||
"utf16string",
|
||||
"webp",
|
||||
"windows-sys 0.61.0",
|
||||
"windows-sys 0.59.0",
|
||||
"xml5ever",
|
||||
"zstd",
|
||||
]
|
||||
@@ -1232,15 +1243,6 @@ version = "0.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c"
|
||||
|
||||
[[package]]
|
||||
name = "num-traits"
|
||||
version = "0.2.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num_cpus"
|
||||
version = "1.17.0"
|
||||
|
||||
13
Cargo.toml
13
Cargo.toml
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "msg_tool"
|
||||
version = "0.2.3"
|
||||
version = "0.2.4"
|
||||
edition = "2024"
|
||||
repository = "https://github.com/lifegpc/msg-tool"
|
||||
description = "A command-line tool for exporting, importing, packing, and unpacking script files."
|
||||
@@ -11,7 +11,6 @@ exclude = [".github", "*.py"]
|
||||
anyhow = "1"
|
||||
byteorder = { version = "1.5", default-features = false, optional = true}
|
||||
clap = { version = "4.5", features = ["derive"] }
|
||||
clap-num = "1.2"
|
||||
csv = "1.3"
|
||||
ctrlc = "3.4"
|
||||
emote-psb = { version = "0.5", optional = true , features = ["serde"] }
|
||||
@@ -20,6 +19,7 @@ fancy-regex = { version = "0.16", optional = true }
|
||||
flate2 = { version = "1.1", optional = true }
|
||||
int-enum = { version = "1.2", optional = true }
|
||||
json = { version = "0.12", optional = true }
|
||||
jpegxl-sys = { package = "msg-tool-jpegxl-sys", version = "0.11", optional = true, features = ["vendored"] }
|
||||
lazy_static = "1.5.0"
|
||||
libflac-sys = { version = "0.3", optional = true }
|
||||
libtlg-rs = { version = "0.2", optional = true, features = ["encode"] }
|
||||
@@ -28,6 +28,7 @@ markup5ever_rcdom = { version = "0.35", optional = true }
|
||||
memchr = { version = "2.7", optional = true }
|
||||
mozjpeg = { version = "0.10", optional = true }
|
||||
msg_tool_macro = { version = "0.2.1" }
|
||||
num_cpus = { version = "1.17", optional = true }
|
||||
overf = "0.1"
|
||||
pelite = { version = "0.10", optional = true }
|
||||
png = { version = "0.18", optional = true }
|
||||
@@ -46,7 +47,7 @@ xml5ever = { version = "0.35", optional = true }
|
||||
zstd = { version = "0.13", optional = true }
|
||||
|
||||
[features]
|
||||
default = ["all-fmt", "image-jpg", "image-webp", "audio-flac"]
|
||||
default = ["all-fmt", "image-jpg", "image-jxl", "image-webp", "audio-flac"]
|
||||
all-fmt = ["all-script", "all-img", "all-arc", "all-audio"]
|
||||
all-script = ["artemis", "artemis-panmimisoft", "bgi", "cat-system", "circus", "entis-gls", "escude", "ex-hibit", "favorite", "hexen-haus", "kirikiri", "silky", "softpal", "will-plus", "yaneurao", "yaneurao-itufuru"]
|
||||
all-img = ["bgi-img", "cat-system-img", "circus-img", "emote-img", "kirikiri-img"]
|
||||
@@ -58,7 +59,7 @@ artemis-arc = ["artemis", "msg_tool_macro/artemis-arc", "sha1"]
|
||||
bgi = ["fancy-regex"]
|
||||
bgi-arc = ["bgi", "rand", "utils-bit-stream"]
|
||||
bgi-audio = ["bgi"]
|
||||
bgi-img = ["bgi", "image", "rand", "utils-bit-stream"]
|
||||
bgi-img = ["bgi", "image", "rand", "utils-threadpool", "utils-bit-stream"]
|
||||
cat-system = ["fancy-regex", "flate2", "int-enum"]
|
||||
cat-system-arc = ["cat-system", "pelite", "utils-blowfish", "utils-crc32"]
|
||||
cat-system-img = ["cat-system", "flate2", "image", "mozjpeg", "utils-bit-stream"]
|
||||
@@ -81,8 +82,9 @@ will-plus = ["utils-str"]
|
||||
yaneurao = []
|
||||
yaneurao-itufuru = ["yaneurao"]
|
||||
# basic feature
|
||||
image = ["png"]
|
||||
image = ["png", "utils-threadpool"]
|
||||
image-jpg = ["mozjpeg"]
|
||||
image-jxl = ["image", "jpegxl-sys", "utils-threadpool"]
|
||||
image-webp = ["webp"]
|
||||
lossless-audio = ["utils-pcm"]
|
||||
audio-flac = ["libflac-sys", "utils-pcm"]
|
||||
@@ -94,6 +96,7 @@ utils-crc32 = []
|
||||
utils-escape = ["fancy-regex"]
|
||||
utils-pcm = []
|
||||
utils-str = []
|
||||
utils-threadpool = ["num_cpus"]
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
windows-sys = { version = "0", features = ["Win32_Globalization", "Win32_System_Diagnostics_Debug"] }
|
||||
|
||||
@@ -53,6 +53,7 @@ msg-tool create -t <script-type> <input> <output>
|
||||
- `json` - [GalTransl](https://github.com/GalTransl/GalTransl)'s JSON format
|
||||
- `m3t` / `m3ta` - A simple text format that supports both original/llm/translated messages.
|
||||
- `yaml` - Same as `json`, but in YAML format.
|
||||
- `po`/`pot` - Gettext PO/POT format.
|
||||
|
||||
## Supported Image Types
|
||||
| Image Type | Feature Name |
|
||||
@@ -60,6 +61,7 @@ msg-tool create -t <script-type> <input> <output>
|
||||
| `png` | `image` (enabled automatically if any image script types are enabled) |
|
||||
| `jpg` | `image-jpg` |
|
||||
| `webp` | `image-webp` |
|
||||
| `jxl` | `image-jxl` |
|
||||
|
||||
## Supported Script Types
|
||||
### Artemis Engine
|
||||
|
||||
67
src/args.rs
67
src/args.rs
@@ -1,4 +1,6 @@
|
||||
use crate::types::*;
|
||||
#[allow(unused)]
|
||||
use crate::utils::num_range::*;
|
||||
use clap::{ArgAction, ArgGroup, Parser, Subcommand};
|
||||
|
||||
#[cfg(feature = "flate2")]
|
||||
@@ -13,7 +15,7 @@ fn parse_compression_level(level: &str) -> Result<u32, String> {
|
||||
} else if lower == "fast" {
|
||||
return Ok(1);
|
||||
}
|
||||
clap_num::number_range(level, 0, 9)
|
||||
number_range(level, 0, 9)
|
||||
}
|
||||
|
||||
#[cfg(feature = "mozjpeg")]
|
||||
@@ -22,7 +24,7 @@ fn parse_jpeg_quality(quality: &str) -> Result<u8, String> {
|
||||
if lower == "best" {
|
||||
return Ok(100);
|
||||
}
|
||||
clap_num::number_range(quality, 0, 100)
|
||||
number_range(quality, 0, 100)
|
||||
}
|
||||
|
||||
#[cfg(feature = "zstd")]
|
||||
@@ -33,7 +35,7 @@ fn parse_zstd_compression_level(level: &str) -> Result<i32, String> {
|
||||
} else if lower == "best" {
|
||||
return Ok(22);
|
||||
}
|
||||
clap_num::number_range(level, 0, 22)
|
||||
number_range(level, 0, 22)
|
||||
}
|
||||
|
||||
#[cfg(feature = "webp")]
|
||||
@@ -42,7 +44,7 @@ fn parse_webp_quality(quality: &str) -> Result<u8, String> {
|
||||
if lower == "best" {
|
||||
return Ok(100);
|
||||
}
|
||||
clap_num::number_range(quality, 0, 100)
|
||||
number_range(quality, 0, 100)
|
||||
}
|
||||
|
||||
#[cfg(feature = "audio-flac")]
|
||||
@@ -55,7 +57,18 @@ fn parse_flac_compression_level(level: &str) -> Result<u32, String> {
|
||||
} else if lower == "default" {
|
||||
return Ok(5);
|
||||
}
|
||||
clap_num::number_range(level, 0, 8)
|
||||
number_range(level, 0, 8)
|
||||
}
|
||||
|
||||
#[cfg(feature = "image-jxl")]
|
||||
fn parse_jxl_distance(s: &str) -> Result<f32, String> {
|
||||
let lower = s.to_ascii_lowercase();
|
||||
if lower == "lossless" {
|
||||
return Ok(0.0);
|
||||
} else if lower == "visually-lossless" {
|
||||
return Ok(1.0);
|
||||
}
|
||||
number_range(s, 0.0, 25.0)
|
||||
}
|
||||
|
||||
/// Tools for export and import scripts
|
||||
@@ -159,7 +172,7 @@ pub struct Arg {
|
||||
/// Enable this will cause BGI scripts to become very large.
|
||||
pub bgi_import_duplicate: bool,
|
||||
#[cfg(feature = "bgi")]
|
||||
#[arg(long, action = ArgAction::SetTrue, global = true, alias = "bgi-no-append")]
|
||||
#[arg(long, action = ArgAction::SetTrue, global = true, visible_alias = "bgi-no-append")]
|
||||
/// Disable appending new strings to the end of BGI scripts.
|
||||
/// Disable may cause BGI scripts broken.
|
||||
pub bgi_disable_append: bool,
|
||||
@@ -172,6 +185,11 @@ pub struct Arg {
|
||||
/// Whether to create scrambled SysGrp images. When in import mode, the default value depends on the original image.
|
||||
/// When in creation mode, it is not enabled by default.
|
||||
pub bgi_img_scramble: Option<bool>,
|
||||
#[cfg(feature = "bgi-img")]
|
||||
#[arg(long, global = true, default_value_t = crate::types::get_default_threads())]
|
||||
/// Workers count for decode BGI compressed images v2 in parallel. Default is half of CPU cores.
|
||||
/// Set this to 1 to disable parallel decoding. 0 means same as 1.
|
||||
pub bgi_img_workers: usize,
|
||||
#[cfg(feature = "cat-system-arc")]
|
||||
#[arg(long, global = true, group = "cat_system_int_encrypt_passwordg")]
|
||||
/// CatSystem2 engine int archive password
|
||||
@@ -211,11 +229,11 @@ pub struct Arg {
|
||||
/// Kirikiri language list. First language code is code for language index 1.
|
||||
pub kirikiri_languages: Option<Vec<String>>,
|
||||
#[cfg(feature = "kirikiri")]
|
||||
#[arg(long, global = true, action = ArgAction::SetTrue, alias = "kr-title")]
|
||||
#[arg(long, global = true, action = ArgAction::SetTrue, visible_alias = "kr-title")]
|
||||
/// Whether to handle title in Kirikiri SCN script.
|
||||
pub kirikiri_title: bool,
|
||||
#[cfg(feature = "kirikiri")]
|
||||
#[arg(long, global = true, action = ArgAction::SetTrue, alias = "kr-no-empty-lines", alias = "kirikiri-no-empty-lines")]
|
||||
#[arg(long, global = true, action = ArgAction::SetTrue, visible_alias = "kr-no-empty-lines", visible_alias = "kirikiri-no-empty-lines")]
|
||||
/// Remove empty lines in Kirikiri KS script.
|
||||
pub kirikiri_remove_empty_lines: bool,
|
||||
#[cfg(feature = "kirikiri")]
|
||||
@@ -423,6 +441,25 @@ pub struct Arg {
|
||||
#[arg(long, global = true, action = ArgAction::SetTrue)]
|
||||
/// Do not filter ascii strings in Favorite HCB script.
|
||||
pub favorite_hcb_no_filter_ascii: bool,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
#[arg(long, global = true, action = ArgAction::SetTrue, visible_alias = "jxl-no-lossless")]
|
||||
/// Disable JXL lossless compression for output images
|
||||
pub jxl_lossy: bool,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
#[arg(long, global = true, default_value_t = 1.0, value_parser = parse_jxl_distance)]
|
||||
/// JXL distance for output images. 0 means mathematically lossless compression. 1.0 means visually lossless compression.
|
||||
/// Allowed range is 0.0-25.0. Recommended range is 0.5-3.0. Default value is 1
|
||||
pub jxl_distance: f32,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
#[arg(long, global = true, default_value_t = 1, visible_alias = "jxl-jobs")]
|
||||
/// Workers count for encode JXL images in parallel. Default is 1.
|
||||
/// Set this to 1 to disable parallel encoding. 0 means same as 1
|
||||
pub jxl_workers: usize,
|
||||
#[cfg(feature = "image")]
|
||||
#[arg(short = 'J', long, global = true, default_value_t = crate::types::get_default_threads(), visible_alias = "img-jobs", visible_alias = "img-workers", visible_alias = "image-jobs")]
|
||||
/// Workers count for encode images in parallel. Default is half of CPU cores.
|
||||
/// Set this to 1 to disable parallel encoding. 0 means same as 1.
|
||||
pub image_workers: usize,
|
||||
#[command(subcommand)]
|
||||
/// Command
|
||||
pub command: Command,
|
||||
@@ -444,11 +481,21 @@ pub struct ImportArgs {
|
||||
#[arg(short = 'P', long, group = "patched_encodingg")]
|
||||
/// Patched script code page
|
||||
pub patched_code_page: Option<u32>,
|
||||
#[arg(long, value_enum, group = "patched_archive_encodingg", alias = "pa")]
|
||||
#[arg(
|
||||
long,
|
||||
value_enum,
|
||||
group = "patched_archive_encodingg",
|
||||
visible_alias = "pa"
|
||||
)]
|
||||
/// Patched archive filename encoding
|
||||
pub patched_archive_encoding: Option<TextEncoding>,
|
||||
#[cfg(windows)]
|
||||
#[arg(long, value_enum, group = "patched_archive_encodingg", alias = "PA")]
|
||||
#[arg(
|
||||
long,
|
||||
value_enum,
|
||||
group = "patched_archive_encodingg",
|
||||
visible_alias = "PA"
|
||||
)]
|
||||
/// Patched archive code page
|
||||
pub patched_archive_code_page: Option<u32>,
|
||||
#[arg(long)]
|
||||
|
||||
@@ -3,6 +3,7 @@ pub mod atomic;
|
||||
#[cfg(feature = "fancy-regex")]
|
||||
pub mod fancy_regex;
|
||||
pub mod io;
|
||||
pub mod mutex;
|
||||
pub mod path;
|
||||
#[cfg(feature = "emote-psb")]
|
||||
pub mod psb;
|
||||
|
||||
11
src/ext/mutex.rs
Normal file
11
src/ext/mutex.rs
Normal file
@@ -0,0 +1,11 @@
|
||||
//! Extension for [std::sync::Mutex].
|
||||
pub trait MutexExt<T> {
|
||||
/// Lock the mutex, blocking the current thread until it can be acquired.
|
||||
fn lock_blocking(&self) -> std::sync::MutexGuard<'_, T>;
|
||||
}
|
||||
|
||||
impl<T> MutexExt<T> for std::sync::Mutex<T> {
|
||||
fn lock_blocking(&self) -> std::sync::MutexGuard<'_, T> {
|
||||
self.lock().unwrap_or_else(|err| err.into_inner())
|
||||
}
|
||||
}
|
||||
326
src/main.rs
326
src/main.rs
@@ -146,7 +146,7 @@ fn get_patched_archive_encoding(
|
||||
pub fn parse_script(
|
||||
filename: &str,
|
||||
arg: &args::Arg,
|
||||
config: &types::ExtraConfig,
|
||||
config: std::sync::Arc<types::ExtraConfig>,
|
||||
) -> anyhow::Result<(
|
||||
Box<dyn scripts::Script>,
|
||||
&'static Box<dyn scripts::ScriptBuilder + Send + Sync>,
|
||||
@@ -162,7 +162,7 @@ pub fn parse_script(
|
||||
filename,
|
||||
encoding,
|
||||
archive_encoding,
|
||||
config,
|
||||
&config,
|
||||
None,
|
||||
)?,
|
||||
builder,
|
||||
@@ -192,7 +192,7 @@ pub fn parse_script(
|
||||
let encoding = get_encoding(arg, builder);
|
||||
let archive_encoding = get_archived_encoding(arg, builder, encoding);
|
||||
return Ok((
|
||||
builder.build_script_from_file(filename, encoding, archive_encoding, config, None)?,
|
||||
builder.build_script_from_file(filename, encoding, archive_encoding, &config, None)?,
|
||||
builder,
|
||||
));
|
||||
}
|
||||
@@ -223,7 +223,7 @@ pub fn parse_script(
|
||||
let encoding = get_encoding(arg, builder);
|
||||
let archive_encoding = get_archived_encoding(arg, builder, encoding);
|
||||
return Ok((
|
||||
builder.build_script_from_file(filename, encoding, archive_encoding, config, None)?,
|
||||
builder.build_script_from_file(filename, encoding, archive_encoding, &config, None)?,
|
||||
builder,
|
||||
));
|
||||
}
|
||||
@@ -240,7 +240,7 @@ pub fn parse_script(
|
||||
pub fn parse_script_from_archive<'a>(
|
||||
file: &mut Box<dyn ArchiveContent + 'a>,
|
||||
arg: &args::Arg,
|
||||
config: &types::ExtraConfig,
|
||||
config: std::sync::Arc<types::ExtraConfig>,
|
||||
archive: &Box<dyn scripts::Script>,
|
||||
) -> anyhow::Result<(
|
||||
Box<dyn scripts::Script>,
|
||||
@@ -258,7 +258,7 @@ pub fn parse_script_from_archive<'a>(
|
||||
file.name(),
|
||||
encoding,
|
||||
archive_encoding,
|
||||
config,
|
||||
&config,
|
||||
Some(archive),
|
||||
)?,
|
||||
builder,
|
||||
@@ -293,7 +293,7 @@ pub fn parse_script_from_archive<'a>(
|
||||
file.name(),
|
||||
encoding,
|
||||
archive_encoding,
|
||||
config,
|
||||
&config,
|
||||
Some(archive),
|
||||
)?,
|
||||
builder,
|
||||
@@ -326,7 +326,7 @@ pub fn parse_script_from_archive<'a>(
|
||||
file.name(),
|
||||
encoding,
|
||||
archive_encoding,
|
||||
config,
|
||||
&config,
|
||||
Some(archive),
|
||||
)?,
|
||||
builder,
|
||||
@@ -346,12 +346,15 @@ pub fn parse_script_from_archive<'a>(
|
||||
pub fn export_script(
|
||||
filename: &str,
|
||||
arg: &args::Arg,
|
||||
config: &types::ExtraConfig,
|
||||
config: std::sync::Arc<types::ExtraConfig>,
|
||||
output: &Option<String>,
|
||||
root_dir: Option<&std::path::Path>,
|
||||
#[cfg(feature = "image")] img_threadpool: Option<
|
||||
&utils::threadpool::ThreadPool<Result<(), anyhow::Error>>,
|
||||
>,
|
||||
) -> anyhow::Result<types::ScriptResult> {
|
||||
eprintln!("Exporting {}", filename);
|
||||
let script = parse_script(filename, arg, config)?.0;
|
||||
let script = parse_script(filename, arg, config.clone())?.0;
|
||||
if script.is_archive() {
|
||||
let odir = match output.as_ref() {
|
||||
Some(output) => {
|
||||
@@ -408,18 +411,18 @@ pub fn export_script(
|
||||
}
|
||||
};
|
||||
if arg.force_script || f.is_script() {
|
||||
let (script_file, _) = match parse_script_from_archive(&mut f, arg, config, &script)
|
||||
{
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
eprintln!("Error parsing script '{}' from archive: {}", filename, e);
|
||||
COUNTER.inc_error();
|
||||
if arg.backtrace {
|
||||
eprintln!("Backtrace: {}", e.backtrace());
|
||||
let (script_file, _) =
|
||||
match parse_script_from_archive(&mut f, arg, config.clone(), &script) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
eprintln!("Error parsing script '{}' from archive: {}", filename, e);
|
||||
COUNTER.inc_error();
|
||||
if arg.backtrace {
|
||||
eprintln!("Backtrace: {}", e.backtrace());
|
||||
}
|
||||
continue;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
};
|
||||
#[cfg(feature = "image")]
|
||||
if script_file.is_image() {
|
||||
if script_file.is_multi_image() {
|
||||
@@ -462,13 +465,44 @@ pub fn export_script(
|
||||
continue;
|
||||
}
|
||||
}
|
||||
utils::img::encode_img(
|
||||
img_data.data,
|
||||
out_type,
|
||||
&out_path.to_string_lossy(),
|
||||
config,
|
||||
)?;
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
if let Some(threadpool) = img_threadpool {
|
||||
let outpath = out_path.to_string_lossy().into_owned();
|
||||
let config = config.clone();
|
||||
threadpool.execute(
|
||||
move |_| {
|
||||
utils::img::encode_img(
|
||||
img_data.data,
|
||||
out_type,
|
||||
&outpath,
|
||||
&config,
|
||||
)
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to encode image {}: {}",
|
||||
outpath,
|
||||
e
|
||||
)
|
||||
})
|
||||
},
|
||||
true,
|
||||
)?;
|
||||
continue;
|
||||
} else {
|
||||
match utils::img::encode_img(
|
||||
img_data.data,
|
||||
out_type,
|
||||
&out_path.to_string_lossy(),
|
||||
&config,
|
||||
) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
eprintln!("Error encoding image: {}", e);
|
||||
COUNTER.inc_error();
|
||||
continue;
|
||||
}
|
||||
}
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
}
|
||||
}
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
continue;
|
||||
@@ -499,20 +533,35 @@ pub fn export_script(
|
||||
continue;
|
||||
}
|
||||
}
|
||||
match utils::img::encode_img(
|
||||
img_data,
|
||||
out_type,
|
||||
&out_path.to_string_lossy(),
|
||||
config,
|
||||
) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
eprintln!("Error encoding image: {}", e);
|
||||
COUNTER.inc_error();
|
||||
continue;
|
||||
if let Some(threadpool) = img_threadpool {
|
||||
let outpath = out_path.to_string_lossy().into_owned();
|
||||
let config = config.clone();
|
||||
threadpool.execute(
|
||||
move |_| {
|
||||
utils::img::encode_img(img_data, out_type, &outpath, &config)
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!("Failed to encode image {}: {}", outpath, e)
|
||||
})
|
||||
},
|
||||
true,
|
||||
)?;
|
||||
continue;
|
||||
} else {
|
||||
match utils::img::encode_img(
|
||||
img_data,
|
||||
out_type,
|
||||
&out_path.to_string_lossy(),
|
||||
&config,
|
||||
) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
eprintln!("Error encoding image: {}", e);
|
||||
COUNTER.inc_error();
|
||||
continue;
|
||||
}
|
||||
}
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
}
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
continue;
|
||||
}
|
||||
let mut of = match &arg.output_type {
|
||||
@@ -836,8 +885,30 @@ pub fn export_script(
|
||||
continue;
|
||||
}
|
||||
}
|
||||
utils::img::encode_img(img_data.data, out_type, &f, config)?;
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
if let Some(threadpool) = img_threadpool {
|
||||
let outpath = f.clone();
|
||||
let config = config.clone();
|
||||
threadpool.execute(
|
||||
move |_| {
|
||||
utils::img::encode_img(img_data.data, out_type, &outpath, &config)
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!("Failed to encode image {}: {}", outpath, e)
|
||||
})
|
||||
},
|
||||
true,
|
||||
)?;
|
||||
continue;
|
||||
} else {
|
||||
match utils::img::encode_img(img_data.data, out_type, &f, &config) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
eprintln!("Error encoding image: {}", e);
|
||||
COUNTER.inc_error();
|
||||
continue;
|
||||
}
|
||||
}
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
}
|
||||
}
|
||||
return Ok(types::ScriptResult::Ok);
|
||||
}
|
||||
@@ -887,7 +958,20 @@ pub fn export_script(
|
||||
}
|
||||
};
|
||||
utils::files::make_sure_dir_exists(&f)?;
|
||||
utils::img::encode_img(img_data, out_type, &f, config)?;
|
||||
if let Some(threadpool) = img_threadpool {
|
||||
let outpath = f.clone();
|
||||
let config = config.clone();
|
||||
threadpool.execute(
|
||||
move |_| {
|
||||
utils::img::encode_img(img_data, out_type, &outpath, &config)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to encode image {}: {}", outpath, e))
|
||||
},
|
||||
true,
|
||||
)?;
|
||||
return Ok(types::ScriptResult::Uncount);
|
||||
} else {
|
||||
utils::img::encode_img(img_data, out_type, &f, &config)?;
|
||||
}
|
||||
return Ok(types::ScriptResult::Ok);
|
||||
}
|
||||
let mut of = match &arg.output_type {
|
||||
@@ -988,14 +1072,14 @@ pub fn export_script(
|
||||
pub fn import_script(
|
||||
filename: &str,
|
||||
arg: &args::Arg,
|
||||
config: &types::ExtraConfig,
|
||||
config: std::sync::Arc<types::ExtraConfig>,
|
||||
imp_cfg: &args::ImportArgs,
|
||||
root_dir: Option<&std::path::Path>,
|
||||
name_csv: Option<&std::collections::HashMap<String, String>>,
|
||||
repl: Option<&types::ReplacementTable>,
|
||||
) -> anyhow::Result<types::ScriptResult> {
|
||||
eprintln!("Importing {}", filename);
|
||||
let (script, builder) = parse_script(filename, arg, config)?;
|
||||
let (script, builder) = parse_script(filename, arg, config.clone())?;
|
||||
if script.is_archive() {
|
||||
let odir = {
|
||||
let mut pb = std::path::PathBuf::from(&imp_cfg.output);
|
||||
@@ -1036,7 +1120,7 @@ pub fn import_script(
|
||||
let pencoding = get_patched_encoding(imp_cfg, builder);
|
||||
let enc = get_patched_archive_encoding(imp_cfg, builder, pencoding);
|
||||
utils::files::make_sure_dir_exists(&patched_f)?;
|
||||
let mut arch = builder.create_archive(&patched_f, &files, enc, config)?;
|
||||
let mut arch = builder.create_archive(&patched_f, &files, enc, &config)?;
|
||||
for (index, filename) in script.iter_archive_filename()?.enumerate() {
|
||||
let filename = match filename {
|
||||
Ok(f) => f,
|
||||
@@ -1062,18 +1146,18 @@ pub fn import_script(
|
||||
};
|
||||
let mut writer = arch.new_file(f.name())?;
|
||||
if arg.force_script || f.is_script() {
|
||||
let (script_file, _) = match parse_script_from_archive(&mut f, arg, config, &script)
|
||||
{
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
eprintln!("Error parsing script '{}' from archive: {}", filename, e);
|
||||
COUNTER.inc_error();
|
||||
if arg.backtrace {
|
||||
eprintln!("Backtrace: {}", e.backtrace());
|
||||
let (script_file, _) =
|
||||
match parse_script_from_archive(&mut f, arg, config.clone(), &script) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
eprintln!("Error parsing script '{}' from archive: {}", filename, e);
|
||||
COUNTER.inc_error();
|
||||
if arg.backtrace {
|
||||
eprintln!("Backtrace: {}", e.backtrace());
|
||||
}
|
||||
continue;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
};
|
||||
let mut of = match &arg.output_type {
|
||||
Some(t) => t.clone(),
|
||||
None => script_file.default_output_script_type(),
|
||||
@@ -1523,7 +1607,7 @@ pub fn pack_archive(
|
||||
input: &str,
|
||||
output: Option<&str>,
|
||||
arg: &args::Arg,
|
||||
config: &types::ExtraConfig,
|
||||
config: std::sync::Arc<types::ExtraConfig>,
|
||||
) -> anyhow::Result<()> {
|
||||
let typ = match &arg.script_type {
|
||||
Some(t) => t,
|
||||
@@ -1569,7 +1653,7 @@ pub fn pack_archive(
|
||||
&output,
|
||||
&reff,
|
||||
get_archived_encoding(arg, builder, get_encoding(arg, builder)),
|
||||
config,
|
||||
&config,
|
||||
)?;
|
||||
for (file, name) in files.iter().zip(reff) {
|
||||
let mut f = match std::fs::File::open(file) {
|
||||
@@ -1606,7 +1690,7 @@ pub fn pack_archive(
|
||||
pub fn unpack_archive(
|
||||
filename: &str,
|
||||
arg: &args::Arg,
|
||||
config: &types::ExtraConfig,
|
||||
config: std::sync::Arc<types::ExtraConfig>,
|
||||
output: &Option<String>,
|
||||
root_dir: Option<&std::path::Path>,
|
||||
) -> anyhow::Result<types::ScriptResult> {
|
||||
@@ -1706,7 +1790,7 @@ pub fn create_file(
|
||||
input: &str,
|
||||
output: Option<&str>,
|
||||
arg: &args::Arg,
|
||||
config: &types::ExtraConfig,
|
||||
config: std::sync::Arc<types::ExtraConfig>,
|
||||
) -> anyhow::Result<()> {
|
||||
let typ = match &arg.script_type {
|
||||
Some(t) => t,
|
||||
@@ -1750,7 +1834,7 @@ pub fn create_file(
|
||||
pb.to_string_lossy().into_owned()
|
||||
}
|
||||
};
|
||||
builder.create_image_file_filename(data, &output, config)?;
|
||||
builder.create_image_file_filename(data, &output, &config)?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -1783,17 +1867,42 @@ pub fn create_file(
|
||||
&output,
|
||||
get_encoding(arg, builder),
|
||||
get_output_encoding(arg),
|
||||
config,
|
||||
&config,
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
static ref COUNTER: utils::counter::Counter = utils::counter::Counter::new();
|
||||
static ref EXIT_LISTENER: std::sync::Mutex<std::collections::BTreeMap<usize, Box<dyn Fn() + Send + Sync>>> = std::sync::Mutex::new(std::collections::BTreeMap::new());
|
||||
#[allow(unused)]
|
||||
static ref EXIT_LISTENER_ID: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn add_exit_listener<F: Fn() + Send + Sync + 'static>(f: F) -> usize {
|
||||
let id = EXIT_LISTENER_ID.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
EXIT_LISTENER
|
||||
.lock()
|
||||
.unwrap_or_else(|err| err.into_inner())
|
||||
.insert(id, Box::new(f));
|
||||
id
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn remove_exit_listener(id: usize) {
|
||||
EXIT_LISTENER
|
||||
.lock()
|
||||
.unwrap_or_else(|err| err.into_inner())
|
||||
.remove(&id);
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let _ = ctrlc::try_set_handler(|| {
|
||||
let listeners = EXIT_LISTENER.lock().unwrap_or_else(|err| err.into_inner());
|
||||
for (_, f) in listeners.iter() {
|
||||
f();
|
||||
}
|
||||
eprintln!("Aborted.");
|
||||
eprintln!("{}", std::ops::Deref::deref(&COUNTER));
|
||||
std::process::exit(1);
|
||||
@@ -1802,7 +1911,7 @@ fn main() {
|
||||
if arg.backtrace {
|
||||
unsafe { std::env::set_var("RUST_LIB_BACKTRACE", "1") };
|
||||
}
|
||||
let cfg = types::ExtraConfig {
|
||||
let cfg = std::sync::Arc::new(types::ExtraConfig {
|
||||
#[cfg(feature = "circus")]
|
||||
circus_mes_type: arg.circus_mes_type.clone(),
|
||||
#[cfg(feature = "escude-arc")]
|
||||
@@ -1925,7 +2034,15 @@ fn main() {
|
||||
kirikiri_title: arg.kirikiri_title,
|
||||
#[cfg(feature = "favorite")]
|
||||
favorite_hcb_filter_ascii: !arg.favorite_hcb_no_filter_ascii,
|
||||
};
|
||||
#[cfg(feature = "bgi-img")]
|
||||
bgi_img_workers: arg.bgi_img_workers,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
jxl_lossless: !arg.jxl_lossy,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
jxl_distance: arg.jxl_distance,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
jxl_workers: arg.jxl_workers,
|
||||
});
|
||||
match &arg.command {
|
||||
args::Command::Export { input, output } => {
|
||||
let (scripts, is_dir) =
|
||||
@@ -1951,8 +2068,43 @@ fn main() {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
#[cfg(feature = "image")]
|
||||
let img_threadpool = if arg.image_workers > 1 {
|
||||
let tp = std::sync::Arc::new(
|
||||
utils::threadpool::ThreadPool::<Result<(), anyhow::Error>>::new(
|
||||
arg.image_workers,
|
||||
Some("img-output-worker-"),
|
||||
false,
|
||||
)
|
||||
.expect("Failed to create image thread pool"),
|
||||
);
|
||||
let tp2 = tp.clone();
|
||||
let id = add_exit_listener(move || {
|
||||
for r in tp2.take_results() {
|
||||
if let Err(e) = r {
|
||||
eprintln!("{}", e);
|
||||
COUNTER.inc_error();
|
||||
} else {
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
}
|
||||
}
|
||||
});
|
||||
Some((tp, id))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
for script in scripts.iter() {
|
||||
let re = export_script(&script, &arg, &cfg, output, root_dir);
|
||||
#[cfg(feature = "image")]
|
||||
let re = export_script(
|
||||
&script,
|
||||
&arg,
|
||||
cfg.clone(),
|
||||
output,
|
||||
root_dir,
|
||||
img_threadpool.as_ref().map(|(t, _)| &**t),
|
||||
);
|
||||
#[cfg(not(feature = "image"))]
|
||||
let re = export_script(&script, &arg, cfg.clone(), output, root_dir);
|
||||
match re {
|
||||
Ok(s) => {
|
||||
COUNTER.inc(s);
|
||||
@@ -1965,7 +2117,31 @@ fn main() {
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(feature = "image")]
|
||||
img_threadpool.as_ref().map(|(t, _)| {
|
||||
for r in t.take_results() {
|
||||
if let Err(e) = r {
|
||||
COUNTER.inc_error();
|
||||
eprintln!("{}", e);
|
||||
} else {
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
#[cfg(feature = "image")]
|
||||
img_threadpool.map(|(t, id)| {
|
||||
t.join();
|
||||
remove_exit_listener(id);
|
||||
for r in t.take_results() {
|
||||
if let Err(e) = r {
|
||||
COUNTER.inc_error();
|
||||
eprintln!("{}", e);
|
||||
} else {
|
||||
COUNTER.inc(types::ScriptResult::Ok);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
args::Command::Import(args) => {
|
||||
let name_csv = match &args.name_csv {
|
||||
@@ -2006,7 +2182,7 @@ fn main() {
|
||||
let re = import_script(
|
||||
&script,
|
||||
&arg,
|
||||
&cfg,
|
||||
cfg.clone(),
|
||||
args,
|
||||
root_dir,
|
||||
name_csv.as_ref(),
|
||||
@@ -2027,7 +2203,12 @@ fn main() {
|
||||
}
|
||||
}
|
||||
args::Command::Pack { input, output } => {
|
||||
let re = pack_archive(input, output.as_ref().map(|s| s.as_str()), &arg, &cfg);
|
||||
let re = pack_archive(
|
||||
input,
|
||||
output.as_ref().map(|s| s.as_str()),
|
||||
&arg,
|
||||
cfg.clone(),
|
||||
);
|
||||
if let Err(e) = re {
|
||||
COUNTER.inc_error();
|
||||
eprintln!("Error packing archive: {}", e);
|
||||
@@ -2057,7 +2238,7 @@ fn main() {
|
||||
None
|
||||
};
|
||||
for script in scripts.iter() {
|
||||
let re = unpack_archive(&script, &arg, &cfg, output, root_dir);
|
||||
let re = unpack_archive(&script, &arg, cfg.clone(), output, root_dir);
|
||||
match re {
|
||||
Ok(s) => {
|
||||
COUNTER.inc(s);
|
||||
@@ -2073,7 +2254,12 @@ fn main() {
|
||||
}
|
||||
}
|
||||
args::Command::Create { input, output } => {
|
||||
let re = create_file(input, output.as_ref().map(|s| s.as_str()), &arg, &cfg);
|
||||
let re = create_file(
|
||||
input,
|
||||
output.as_ref().map(|s| s.as_str()),
|
||||
&arg,
|
||||
cfg.clone(),
|
||||
);
|
||||
if let Err(e) = re {
|
||||
COUNTER.inc_error();
|
||||
eprintln!("Error creating file: {}", e);
|
||||
|
||||
@@ -4,6 +4,7 @@ use crate::ext::vec::*;
|
||||
use crate::scripts::base::*;
|
||||
use crate::types::*;
|
||||
use crate::utils::bit_stream::*;
|
||||
use crate::utils::num_range::*;
|
||||
use anyhow::Result;
|
||||
use rand::Rng;
|
||||
use std::collections::BinaryHeap;
|
||||
@@ -688,5 +689,5 @@ impl Script for Dsc {
|
||||
|
||||
/// Parses the minimum length for LZSS compression from a string.
|
||||
pub fn parse_min_length(len: &str) -> Result<usize, String> {
|
||||
clap_num::number_range(len, 2, 256)
|
||||
number_range(len, 2, 256)
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ use crate::types::*;
|
||||
use crate::utils::bit_stream::*;
|
||||
use crate::utils::img::*;
|
||||
use crate::utils::struct_pack::*;
|
||||
use crate::utils::threadpool::*;
|
||||
use anyhow::Result;
|
||||
use msg_tool_macro::*;
|
||||
use std::io::{Read, Seek, Write};
|
||||
@@ -133,6 +134,7 @@ pub struct BgiCBG {
|
||||
header: BgiCBGHeader,
|
||||
data: MemReader,
|
||||
color_type: CbgColorType,
|
||||
decode_workers: usize,
|
||||
}
|
||||
|
||||
impl BgiCBG {
|
||||
@@ -140,7 +142,7 @@ impl BgiCBG {
|
||||
///
|
||||
/// * `data` - The buffer containing the script data.
|
||||
/// * `config` - Extra configuration options.
|
||||
pub fn new(data: Vec<u8>, _config: &ExtraConfig) -> Result<Self> {
|
||||
pub fn new(data: Vec<u8>, config: &ExtraConfig) -> Result<Self> {
|
||||
let mut reader = MemReader::new(data);
|
||||
let mut magic = [0u8; 16];
|
||||
reader.read_exact(&mut magic)?;
|
||||
@@ -167,6 +169,7 @@ impl BgiCBG {
|
||||
header,
|
||||
data: reader,
|
||||
color_type,
|
||||
decode_workers: config.bgi_img_workers.max(1),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -185,7 +188,12 @@ impl Script for BgiCBG {
|
||||
}
|
||||
|
||||
fn export_image(&self) -> Result<ImageData> {
|
||||
let decoder = CbgDecoder::new(self.data.to_ref(), &self.header, self.color_type)?;
|
||||
let decoder = CbgDecoder::new(
|
||||
self.data.to_ref(),
|
||||
&self.header,
|
||||
self.color_type,
|
||||
self.decode_workers,
|
||||
)?;
|
||||
Ok(decoder.unpack()?)
|
||||
}
|
||||
|
||||
@@ -209,6 +217,7 @@ struct CbgDecoder<'a> {
|
||||
magic: u32,
|
||||
pixel_size: u8,
|
||||
stride: usize,
|
||||
workers: usize,
|
||||
}
|
||||
|
||||
impl<'a> CbgDecoder<'a> {
|
||||
@@ -216,6 +225,7 @@ impl<'a> CbgDecoder<'a> {
|
||||
reader: MemReaderRef<'a>,
|
||||
info: &'a BgiCBGHeader,
|
||||
color_type: CbgColorType,
|
||||
workers: usize,
|
||||
) -> Result<Self> {
|
||||
let magic = 0;
|
||||
let key = info.key;
|
||||
@@ -230,6 +240,7 @@ impl<'a> CbgDecoder<'a> {
|
||||
color_type,
|
||||
pixel_size,
|
||||
stride,
|
||||
workers,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -334,7 +345,7 @@ impl<'a> CbgDecoder<'a> {
|
||||
has_alpha: AtomicBool::new(false),
|
||||
});
|
||||
|
||||
let mut tasks = Vec::new();
|
||||
let thread_pool = ThreadPool::new(self.workers, Some("cbg-decoder-worker-"), false)?;
|
||||
let mut dst = 0i32;
|
||||
|
||||
for i in 0..y_blocks {
|
||||
@@ -347,23 +358,27 @@ impl<'a> CbgDecoder<'a> {
|
||||
let closure_dst = dst;
|
||||
let decoder_ref = Arc::clone(&decoder);
|
||||
|
||||
let task = std::thread::spawn(move || {
|
||||
decoder_ref.unpack_block(block_offset, next_offset - block_offset, closure_dst)
|
||||
});
|
||||
tasks.push(task);
|
||||
thread_pool.execute(
|
||||
move |_| {
|
||||
decoder_ref.unpack_block(block_offset, next_offset - block_offset, closure_dst)
|
||||
},
|
||||
true,
|
||||
)?;
|
||||
dst += width * 32;
|
||||
}
|
||||
|
||||
if self.info.bpp == 32 {
|
||||
let decoder_ref = Arc::clone(&decoder);
|
||||
let task =
|
||||
std::thread::spawn(move || decoder_ref.unpack_alpha(offsets[y_blocks as usize]));
|
||||
tasks.push(task);
|
||||
thread_pool.execute(
|
||||
move |_| decoder_ref.unpack_alpha(offsets[y_blocks as usize]),
|
||||
true,
|
||||
)?;
|
||||
}
|
||||
|
||||
let tasks = thread_pool.into_results();
|
||||
|
||||
for task in tasks {
|
||||
task.join()
|
||||
.map_err(|e| anyhow::anyhow!("Thread join failed: {:?}", e))??;
|
||||
task?;
|
||||
}
|
||||
|
||||
let has_alpha = decoder.has_alpha.qload();
|
||||
@@ -394,13 +409,25 @@ impl<'a> CbgDecoder<'a> {
|
||||
ImageColorType::Bgr
|
||||
};
|
||||
|
||||
Ok(ImageData {
|
||||
let img = ImageData {
|
||||
width: decoder.width as u32,
|
||||
height: decoder.height as u32,
|
||||
color_type,
|
||||
depth: 8,
|
||||
data: output,
|
||||
})
|
||||
};
|
||||
|
||||
if decoder.width != self.info.width as i32 || decoder.height != self.info.height as i32 {
|
||||
return Ok(draw_on_canvas(
|
||||
img,
|
||||
self.info.width as u32,
|
||||
self.info.height as u32,
|
||||
0,
|
||||
0,
|
||||
)?);
|
||||
}
|
||||
|
||||
Ok(img)
|
||||
}
|
||||
|
||||
fn read_encoded(&mut self) -> Result<Vec<u8>> {
|
||||
|
||||
35
src/types.rs
35
src/types.rs
@@ -426,6 +426,25 @@ pub struct ExtraConfig {
|
||||
#[default(true)]
|
||||
/// Whether to filter ascii strings in Favorite HCB script.
|
||||
pub favorite_hcb_filter_ascii: bool,
|
||||
#[cfg(feature = "bgi-img")]
|
||||
#[default(get_default_threads())]
|
||||
/// Workers count for decode BGI compressed images v2 in parallel. Default is half of CPU cores.
|
||||
/// Set this to 1 to disable parallel decoding. 0 means same as 1.
|
||||
pub bgi_img_workers: usize,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
#[default(true)]
|
||||
/// Use JXL lossless compression for output images. Enabled by default.
|
||||
pub jxl_lossless: bool,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
#[default(1.0)]
|
||||
/// JXL distance for output images. 0 means mathematically lossless compression. 1.0 means visually lossless compression.
|
||||
/// Allowed range is 0.0-25.0. Recommended range is 0.5-3.0. Default value is 1.0.
|
||||
pub jxl_distance: f32,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
#[default(1)]
|
||||
/// Workers count for encode JXL images in parallel. Default is 1.
|
||||
/// Set this to 1 to disable parallel encoding. 0 means same as 1
|
||||
pub jxl_workers: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq, PartialOrd, Ord)]
|
||||
@@ -614,6 +633,9 @@ pub enum ScriptResult {
|
||||
/// Operation completed without any changes.
|
||||
/// For example, no messages found in the script.
|
||||
Ignored,
|
||||
/// Operation not completed.
|
||||
/// This will not count in statistics.
|
||||
Uncount,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq, PartialOrd, Ord)]
|
||||
@@ -706,6 +728,9 @@ pub enum ImageOutputType {
|
||||
#[cfg(feature = "image-webp")]
|
||||
/// WebP image
|
||||
Webp,
|
||||
#[cfg(feature = "image-jxl")]
|
||||
/// JPEG XL image
|
||||
Jxl,
|
||||
}
|
||||
|
||||
#[cfg(feature = "image")]
|
||||
@@ -723,6 +748,8 @@ impl TryFrom<&str> for ImageOutputType {
|
||||
"jpeg" => Ok(ImageOutputType::Jpg),
|
||||
#[cfg(feature = "image-webp")]
|
||||
"webp" => Ok(ImageOutputType::Webp),
|
||||
#[cfg(feature = "image-jxl")]
|
||||
"jxl" => Ok(ImageOutputType::Jxl),
|
||||
_ => Err(anyhow::anyhow!("Unsupported image output type: {}", value)),
|
||||
}
|
||||
}
|
||||
@@ -751,6 +778,8 @@ impl AsRef<str> for ImageOutputType {
|
||||
ImageOutputType::Jpg => "jpg",
|
||||
#[cfg(feature = "image-webp")]
|
||||
ImageOutputType::Webp => "webp",
|
||||
#[cfg(feature = "image-jxl")]
|
||||
ImageOutputType::Jxl => "jxl",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -915,3 +944,9 @@ impl AsRef<str> for LosslessAudioFormat {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "utils-threadpool")]
|
||||
#[allow(unused)]
|
||||
pub(crate) fn get_default_threads() -> usize {
|
||||
num_cpus::get().max(2) / 2
|
||||
}
|
||||
|
||||
@@ -35,9 +35,14 @@ impl Counter {
|
||||
/// Increments the count of script executions.
|
||||
pub fn inc(&self, result: ScriptResult) {
|
||||
match result {
|
||||
ScriptResult::Ok => self.ok.fetch_add(1, SeqCst),
|
||||
ScriptResult::Ignored => self.ignored.fetch_add(1, SeqCst),
|
||||
};
|
||||
ScriptResult::Ok => {
|
||||
self.ok.fetch_add(1, SeqCst);
|
||||
}
|
||||
ScriptResult::Ignored => {
|
||||
self.ignored.fetch_add(1, SeqCst);
|
||||
}
|
||||
ScriptResult::Uncount => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
//! Image Utilities
|
||||
#[cfg(feature = "image-jxl")]
|
||||
use super::jxl::*;
|
||||
use crate::ext::io::*;
|
||||
use crate::types::*;
|
||||
use anyhow::Result;
|
||||
@@ -285,6 +287,13 @@ pub fn encode_img(
|
||||
file.write_all(&re)?;
|
||||
Ok(())
|
||||
}
|
||||
#[cfg(feature = "image-jxl")]
|
||||
ImageOutputType::Jxl => {
|
||||
let mut file = crate::utils::files::write_file(filename)?;
|
||||
let data = encode_jxl(data, config)?;
|
||||
file.write_all(&data)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -401,6 +410,11 @@ pub fn decode_img(typ: ImageOutputType, filename: &str) -> Result<ImageData> {
|
||||
data,
|
||||
})
|
||||
}
|
||||
#[cfg(feature = "image-jxl")]
|
||||
ImageOutputType::Jxl => {
|
||||
let file = crate::utils::files::read_file(filename)?;
|
||||
decode_jxl(&file[..])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
310
src/utils/jxl.rs
Normal file
310
src/utils/jxl.rs
Normal file
@@ -0,0 +1,310 @@
|
||||
//! JPEG XL image support
|
||||
use super::img::*;
|
||||
use super::num_range::*;
|
||||
use super::threadpool::*;
|
||||
use crate::types::*;
|
||||
use anyhow::Result;
|
||||
use jpegxl_sys::common::types::*;
|
||||
use jpegxl_sys::decode::*;
|
||||
use jpegxl_sys::encoder::encode::*;
|
||||
use jpegxl_sys::metadata::codestream_header::*;
|
||||
use jpegxl_sys::threads::parallel_runner::*;
|
||||
use std::ffi::c_void;
|
||||
use std::io::Read;
|
||||
|
||||
struct JxlDecoderHandle {
|
||||
handle: *mut JxlDecoder,
|
||||
}
|
||||
|
||||
impl Drop for JxlDecoderHandle {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
JxlDecoderDestroy(self.handle);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct JxlEncoderHandle {
|
||||
handle: *mut JxlEncoder,
|
||||
}
|
||||
|
||||
impl Drop for JxlEncoderHandle {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
JxlEncoderDestroy(self.handle);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct ThreadPoolRunner {
|
||||
thread_pool: ThreadPool<()>,
|
||||
}
|
||||
|
||||
impl ThreadPoolRunner {
|
||||
fn new(workers: usize) -> Result<Self> {
|
||||
let thread_pool = ThreadPool::new(workers, Some("jxl-thread-runner-"), true)?;
|
||||
Ok(Self { thread_pool })
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct JpegxlPointer(*mut c_void);
|
||||
|
||||
unsafe impl Send for JpegxlPointer {}
|
||||
|
||||
unsafe extern "C-unwind" fn thread_pool_runner(
|
||||
runner_opaque: *mut c_void,
|
||||
jpegxl_opaque: *mut c_void,
|
||||
init: JxlParallelRunInit,
|
||||
func: JxlParallelRunFunction,
|
||||
start_range: u32,
|
||||
end_range: u32,
|
||||
) -> JxlParallelRetCode {
|
||||
if runner_opaque.is_null() || jpegxl_opaque.is_null() {
|
||||
return JXL_PARALLEL_RET_RUNNER_ERROR;
|
||||
}
|
||||
let runner = unsafe { &*(runner_opaque as *const ThreadPoolRunner) };
|
||||
let initre = unsafe { init(jpegxl_opaque, runner.thread_pool.size()) };
|
||||
if initre != JXL_PARALLEL_RET_SUCCESS {
|
||||
return initre;
|
||||
}
|
||||
let jpegxl = JpegxlPointer(jpegxl_opaque);
|
||||
for i in start_range..end_range {
|
||||
let jpegxl = jpegxl;
|
||||
let func = func;
|
||||
match runner.thread_pool.execute(
|
||||
move |thread_id| unsafe {
|
||||
let jpegxl = jpegxl;
|
||||
func(jpegxl.0, i, thread_id)
|
||||
},
|
||||
true,
|
||||
) {
|
||||
Ok(_) => {}
|
||||
Err(_) => return JXL_PARALLEL_RET_RUNNER_ERROR,
|
||||
}
|
||||
}
|
||||
runner.thread_pool.join();
|
||||
JXL_PARALLEL_RET_SUCCESS
|
||||
}
|
||||
|
||||
fn check_decoder_status(status: JxlDecoderStatus) -> Result<()> {
|
||||
match status {
|
||||
JxlDecoderStatus::Success => Ok(()),
|
||||
_ => Err(anyhow::anyhow!("JXL decoder error: {:?}", status)),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_encoder_status(status: JxlEncoderStatus) -> Result<()> {
|
||||
match status {
|
||||
JxlEncoderStatus::Success => Ok(()),
|
||||
_ => Err(anyhow::anyhow!("JXL encoder error: {:?}", status)),
|
||||
}
|
||||
}
|
||||
|
||||
fn default_basic_info() -> JxlBasicInfo {
|
||||
let basic_info = std::mem::MaybeUninit::<JxlBasicInfo>::zeroed();
|
||||
unsafe { basic_info.assume_init_read() }
|
||||
}
|
||||
|
||||
/// Decode JXL image from reader
|
||||
pub fn decode_jxl<R: Read>(mut r: R) -> Result<ImageData> {
|
||||
let decoder = unsafe { JxlDecoderCreate(std::ptr::null()) };
|
||||
if decoder.is_null() {
|
||||
return Err(anyhow::anyhow!("Failed to create JXL decoder"));
|
||||
}
|
||||
let dh = JxlDecoderHandle { handle: decoder };
|
||||
let events = JxlDecoderStatus::BasicInfo as i32
|
||||
| JxlDecoderStatus::FullImage as i32
|
||||
| JxlDecoderStatus::ColorEncoding as i32;
|
||||
check_decoder_status(unsafe { JxlDecoderSubscribeEvents(dh.handle, events) })?;
|
||||
let mut data = Vec::new();
|
||||
r.read_to_end(&mut data)?;
|
||||
check_decoder_status(unsafe { JxlDecoderSetInput(dh.handle, data.as_ptr(), data.len()) })?;
|
||||
unsafe {
|
||||
JxlDecoderCloseInput(dh.handle);
|
||||
};
|
||||
let mut basic_info = default_basic_info();
|
||||
let mut color_type = ImageColorType::Rgb;
|
||||
let mut buffer = Vec::new();
|
||||
loop {
|
||||
let status = unsafe { JxlDecoderProcessInput(dh.handle) };
|
||||
match status {
|
||||
JxlDecoderStatus::BasicInfo => {
|
||||
check_decoder_status(unsafe {
|
||||
JxlDecoderGetBasicInfo(dh.handle, &mut basic_info)
|
||||
})?;
|
||||
match basic_info.num_color_channels {
|
||||
1 => color_type = ImageColorType::Grayscale,
|
||||
3 => {
|
||||
if basic_info.alpha_bits > 0 {
|
||||
color_type = ImageColorType::Rgba;
|
||||
} else {
|
||||
color_type = ImageColorType::Rgb;
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Unsupported number of color channels: {}",
|
||||
basic_info.num_color_channels
|
||||
));
|
||||
}
|
||||
}
|
||||
if !matches!(basic_info.bits_per_sample, 8 | 16) {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Unsupported bits per sample: {}",
|
||||
basic_info.bits_per_sample
|
||||
));
|
||||
}
|
||||
}
|
||||
JxlDecoderStatus::NeedImageOutBuffer => {
|
||||
let format = JxlPixelFormat {
|
||||
num_channels: color_type.bpp(1) as u32,
|
||||
data_type: if basic_info.bits_per_sample <= 8 {
|
||||
JxlDataType::Uint8
|
||||
} else {
|
||||
JxlDataType::Uint16
|
||||
},
|
||||
endianness: JxlEndianness::Little,
|
||||
align: 0,
|
||||
};
|
||||
let mut buffer_size: usize = 0;
|
||||
check_decoder_status(unsafe {
|
||||
JxlDecoderImageOutBufferSize(dh.handle, &format, &mut buffer_size)
|
||||
})?;
|
||||
buffer.resize(buffer_size, 0);
|
||||
check_decoder_status(unsafe {
|
||||
JxlDecoderSetImageOutBuffer(
|
||||
dh.handle,
|
||||
&format,
|
||||
buffer.as_mut_ptr() as *mut _,
|
||||
buffer_size,
|
||||
)
|
||||
})?;
|
||||
}
|
||||
JxlDecoderStatus::Success => {
|
||||
break;
|
||||
}
|
||||
JxlDecoderStatus::Error => {
|
||||
return Err(anyhow::anyhow!("JXL decoding error"));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(ImageData {
|
||||
width: basic_info.xsize,
|
||||
height: basic_info.ysize,
|
||||
color_type,
|
||||
depth: basic_info.bits_per_sample as u8,
|
||||
data: buffer,
|
||||
})
|
||||
}
|
||||
|
||||
/// Encode image data to JXL format
|
||||
pub fn encode_jxl(mut img: ImageData, config: &ExtraConfig) -> Result<Vec<u8>> {
|
||||
let encoder = unsafe { JxlEncoderCreate(std::ptr::null()) };
|
||||
if encoder.is_null() {
|
||||
return Err(anyhow::anyhow!("Failed to create JXL encoder"));
|
||||
}
|
||||
let eh = JxlEncoderHandle { handle: encoder };
|
||||
let ph = if config.jxl_workers > 1 {
|
||||
let ph = ThreadPoolRunner::new(config.jxl_workers)?;
|
||||
Some(ph)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(ph) = &ph {
|
||||
check_encoder_status(unsafe {
|
||||
JxlEncoderSetParallelRunner(
|
||||
eh.handle,
|
||||
thread_pool_runner,
|
||||
ph as *const _ as *mut c_void,
|
||||
)
|
||||
})?;
|
||||
}
|
||||
let mut basic_info = default_basic_info();
|
||||
basic_info.xsize = img.width;
|
||||
basic_info.ysize = img.height;
|
||||
basic_info.bits_per_sample = match img.depth {
|
||||
8 => 8,
|
||||
16 => 16,
|
||||
_ => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Unsupported bits per sample: {}",
|
||||
img.depth
|
||||
));
|
||||
}
|
||||
};
|
||||
basic_info.alpha_bits = match img.color_type {
|
||||
ImageColorType::Rgba | ImageColorType::Bgra => img.depth as u32,
|
||||
_ => 0,
|
||||
};
|
||||
basic_info.num_color_channels = match img.color_type {
|
||||
ImageColorType::Bgr | ImageColorType::Rgb | ImageColorType::Bgra | ImageColorType::Rgba => {
|
||||
3
|
||||
}
|
||||
ImageColorType::Grayscale => 1,
|
||||
};
|
||||
basic_info.num_extra_channels = if basic_info.alpha_bits > 0 { 1 } else { 0 };
|
||||
basic_info.orientation = JxlOrientation::Identity;
|
||||
basic_info.uses_original_profile = JxlBool::True;
|
||||
check_encoder_status(unsafe { JxlEncoderSetBasicInfo(eh.handle, &basic_info) })?;
|
||||
let options = unsafe { JxlEncoderFrameSettingsCreate(eh.handle, std::ptr::null()) };
|
||||
if options.is_null() {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to create JXL encoder frame settings"
|
||||
));
|
||||
}
|
||||
check_encoder_status(unsafe {
|
||||
JxlEncoderSetFrameLossless(options, JxlBool::from(config.jxl_lossless))
|
||||
})?;
|
||||
if !config.jxl_lossless {
|
||||
let distance = check_range(config.jxl_distance, 0.0, 25.0)
|
||||
.map_err(|e| anyhow::anyhow!("Invalid JXL distance: {}", e))?;
|
||||
check_encoder_status(unsafe { JxlEncoderSetFrameDistance(options, distance) })?;
|
||||
}
|
||||
let format = JxlPixelFormat {
|
||||
num_channels: img.color_type.bpp(1) as u32,
|
||||
data_type: if img.depth <= 8 {
|
||||
JxlDataType::Uint8
|
||||
} else {
|
||||
JxlDataType::Uint16
|
||||
},
|
||||
endianness: JxlEndianness::Little,
|
||||
align: 0,
|
||||
};
|
||||
match img.color_type {
|
||||
ImageColorType::Bgr => {
|
||||
convert_bgr_to_rgb(&mut img)?;
|
||||
}
|
||||
ImageColorType::Bgra => {
|
||||
convert_bgra_to_rgba(&mut img)?;
|
||||
}
|
||||
_ => {}
|
||||
};
|
||||
check_encoder_status(unsafe {
|
||||
JxlEncoderAddImageFrame(
|
||||
options,
|
||||
&format,
|
||||
img.data.as_ptr() as *const _,
|
||||
img.data.len(),
|
||||
)
|
||||
})?;
|
||||
unsafe { JxlEncoderCloseInput(eh.handle) };
|
||||
let mut compressed_data = Vec::new();
|
||||
let mut buffer = [0u8; 4096];
|
||||
loop {
|
||||
let mut avail_out = buffer.len();
|
||||
let mut next_out = buffer.as_mut_ptr();
|
||||
let status = unsafe { JxlEncoderProcessOutput(eh.handle, &mut next_out, &mut avail_out) };
|
||||
let used = buffer.len() - avail_out;
|
||||
compressed_data.extend_from_slice(&buffer[..used]);
|
||||
match status {
|
||||
JxlEncoderStatus::Success => break,
|
||||
JxlEncoderStatus::NeedMoreOutput => {}
|
||||
_ => {
|
||||
return Err(anyhow::anyhow!("JXL encoding error: {:?}", status));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(compressed_data)
|
||||
}
|
||||
@@ -16,15 +16,20 @@ pub mod files;
|
||||
pub mod flac;
|
||||
#[cfg(feature = "image")]
|
||||
pub mod img;
|
||||
#[cfg(feature = "image-jxl")]
|
||||
pub mod jxl;
|
||||
#[cfg(feature = "lossless-audio")]
|
||||
pub mod lossless_audio;
|
||||
mod macros;
|
||||
pub mod name_replacement;
|
||||
pub mod num_range;
|
||||
#[cfg(feature = "utils-pcm")]
|
||||
pub mod pcm;
|
||||
#[cfg(feature = "utils-str")]
|
||||
pub mod str;
|
||||
pub mod struct_pack;
|
||||
#[cfg(feature = "utils-threadpool")]
|
||||
pub mod threadpool;
|
||||
|
||||
#[cfg(windows)]
|
||||
pub use encoding_win::WinError;
|
||||
|
||||
27
src/utils/num_range.rs
Normal file
27
src/utils/num_range.rs
Normal file
@@ -0,0 +1,27 @@
|
||||
//! Functions for parsing numbers within a range.
|
||||
/// Check if a value is within the specified range.
|
||||
pub fn check_range<T>(val: T, min: T, max: T) -> Result<T, String>
|
||||
where
|
||||
T: PartialOrd,
|
||||
T: std::fmt::Display,
|
||||
{
|
||||
if val < min {
|
||||
return Err(format!("Value {} is less than minimum {}", val, min));
|
||||
} else if val > max {
|
||||
return Err(format!("Value {} is greater than maximum {}", val, max));
|
||||
}
|
||||
Ok(val)
|
||||
}
|
||||
|
||||
/// Parse a number from a string and check if it is within the specified range.
|
||||
pub fn number_range<T>(s: &str, min: T, max: T) -> Result<T, String>
|
||||
where
|
||||
T: std::str::FromStr,
|
||||
<T as std::str::FromStr>::Err: std::fmt::Display,
|
||||
T: PartialOrd,
|
||||
T: std::fmt::Display,
|
||||
{
|
||||
debug_assert!(min <= max, "min should be less than or equal to max");
|
||||
let val = s.parse::<T>().map_err(|e| format!("{}", e))?;
|
||||
check_range(val, min, max)
|
||||
}
|
||||
227
src/utils/threadpool.rs
Normal file
227
src/utils/threadpool.rs
Normal file
@@ -0,0 +1,227 @@
|
||||
//! Thread pool utilities
|
||||
use crate::ext::mutex::*;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{
|
||||
Arc, Condvar, Mutex,
|
||||
mpsc::{Receiver, SyncSender, TrySendError, sync_channel},
|
||||
};
|
||||
use std::thread::{self, JoinHandle};
|
||||
|
||||
type Job<T> = Box<dyn FnOnce(usize) -> T + Send + 'static>;
|
||||
|
||||
/// A simple generic thread pool.
|
||||
///
|
||||
/// - T: the return type of tasks. Completed task results are stored in `results: Arc<Mutex<Vec<T>>>`.
|
||||
/// - execute accepts a task and a `block_if_full` flag:
|
||||
/// * if true, submission will block when the pool is saturated until a worker becomes available;
|
||||
/// * if false, submission will return an error when the pool is saturated.
|
||||
/// - join waits until all submitted tasks have completed (it does not shut down the pool).
|
||||
pub struct ThreadPool<T: Send + 'static> {
|
||||
sender: Option<SyncSender<Job<T>>>,
|
||||
#[allow(unused)]
|
||||
receiver: Arc<Mutex<Receiver<Job<T>>>>,
|
||||
workers: Vec<JoinHandle<()>>,
|
||||
/// Completed task results
|
||||
pub results: Arc<Mutex<Vec<T>>>,
|
||||
/// Number of pending tasks (queued + running)
|
||||
pending: Arc<AtomicUsize>,
|
||||
/// Pair for wait/notify in join
|
||||
pending_pair: Arc<(Mutex<()>, Condvar)>,
|
||||
size: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
/// Error type for [ThreadPool::execute]
|
||||
pub enum ExecuteError {
|
||||
/// Pool is full
|
||||
Full,
|
||||
/// Pool is closed
|
||||
Closed,
|
||||
}
|
||||
|
||||
impl std::error::Error for ExecuteError {}
|
||||
|
||||
impl std::fmt::Display for ExecuteError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
ExecuteError::Full => write!(f, "ThreadPool is full"),
|
||||
ExecuteError::Closed => write!(f, "ThreadPool is closed"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Send + 'static> ThreadPool<T> {
|
||||
/// Get the number of worker threads in the pool.
|
||||
pub fn size(&self) -> usize {
|
||||
self.size
|
||||
}
|
||||
|
||||
/// Create a new thread pool with `size` workers.
|
||||
/// The internal submission channel is bounded to `size`, so when all workers are busy and
|
||||
/// the channel is full, further submissions will block or return error depending on the flag.
|
||||
///
|
||||
/// * `name` - Optional base name for worker threads. If None, "threadpool-worker-" is used.
|
||||
/// * `no_result` - If true, results are not stored (saves some overhead if not needed).
|
||||
pub fn new<'a>(
|
||||
size: usize,
|
||||
name: Option<&'a str>,
|
||||
no_result: bool,
|
||||
) -> Result<Self, std::io::Error> {
|
||||
if size == 0 {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"worker size must be > 0",
|
||||
));
|
||||
}
|
||||
|
||||
let (tx, rx) = sync_channel::<Job<T>>(size);
|
||||
let receiver = Arc::new(Mutex::new(rx));
|
||||
let results = Arc::new(Mutex::new(Vec::new()));
|
||||
let pending = Arc::new(AtomicUsize::new(0));
|
||||
let pending_pair = Arc::new((Mutex::new(()), Condvar::new()));
|
||||
let thread_name = name.unwrap_or("threadpool-worker-");
|
||||
|
||||
let mut workers = Vec::with_capacity(size);
|
||||
for id in 0..size {
|
||||
let rx_clone = Arc::clone(&receiver);
|
||||
let results_clone = Arc::clone(&results);
|
||||
let pending_clone = Arc::clone(&pending);
|
||||
let pending_pair_clone = Arc::clone(&pending_pair);
|
||||
|
||||
let handle = thread::Builder::new()
|
||||
.name(format!("{}{}", thread_name, id))
|
||||
.spawn(move || {
|
||||
loop {
|
||||
// Lock receiver to call recv. Using a Mutex around Receiver serializes
|
||||
// the recv calls but is fine for this simple implementation.
|
||||
let job = {
|
||||
let guard = rx_clone.lock_blocking();
|
||||
// If recv returns Err, sender was dropped -> exit thread
|
||||
guard.recv()
|
||||
};
|
||||
|
||||
match job {
|
||||
Ok(job) => {
|
||||
// Execute the job and store result
|
||||
let res = job(id);
|
||||
if !no_result {
|
||||
let mut r = results_clone.lock_blocking();
|
||||
r.push(res);
|
||||
}
|
||||
|
||||
// Decrement pending count and notify join waiters
|
||||
pending_clone.fetch_sub(1, Ordering::SeqCst);
|
||||
let (lock, cvar) = &*pending_pair_clone;
|
||||
let _g = lock.lock_blocking();
|
||||
cvar.notify_all();
|
||||
}
|
||||
Err(_) => {
|
||||
// Channel closed -> shutdown worker
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
})?;
|
||||
|
||||
workers.push(handle);
|
||||
}
|
||||
|
||||
Ok(ThreadPool {
|
||||
sender: Some(tx),
|
||||
receiver,
|
||||
workers,
|
||||
results,
|
||||
pending,
|
||||
pending_pair,
|
||||
size,
|
||||
})
|
||||
}
|
||||
|
||||
/// Execute a task. If `block_if_full` is true, this call will block when the internal
|
||||
/// submission channel is full (i.e. all workers busy and buffer full) until space becomes available.
|
||||
/// If `block_if_full` is false, this returns Err(ExecuteError::Full) when the channel is full.
|
||||
///
|
||||
/// job: a closure that takes the worker id (0..size-1) and returns a T.
|
||||
pub fn execute<F>(&self, job: F, block_if_full: bool) -> Result<(), ExecuteError>
|
||||
where
|
||||
F: FnOnce(usize) -> T + Send + 'static,
|
||||
{
|
||||
let sender = match &self.sender {
|
||||
Some(s) => s,
|
||||
None => return Err(ExecuteError::Closed),
|
||||
};
|
||||
|
||||
// Increase pending count for this submission. If submission fails we will decrement.
|
||||
self.pending.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
let boxed: Job<T> = Box::new(job);
|
||||
|
||||
if block_if_full {
|
||||
// This will block until there is space in the bounded channel or the channel is closed.
|
||||
if sender.send(boxed).is_err() {
|
||||
// Channel closed
|
||||
self.pending.fetch_sub(1, Ordering::SeqCst);
|
||||
return Err(ExecuteError::Closed);
|
||||
}
|
||||
Ok(())
|
||||
} else {
|
||||
match sender.try_send(boxed) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(TrySendError::Full(_)) => {
|
||||
// revert pending increment
|
||||
self.pending.fetch_sub(1, Ordering::SeqCst);
|
||||
Err(ExecuteError::Full)
|
||||
}
|
||||
Err(TrySendError::Disconnected(_)) => {
|
||||
self.pending.fetch_sub(1, Ordering::SeqCst);
|
||||
Err(ExecuteError::Closed)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Wait until all submitted tasks have completed. This does not shut down the pool; new tasks
|
||||
/// can still be submitted after join returns.
|
||||
pub fn join(&self) {
|
||||
// Fast path
|
||||
if self.pending.load(Ordering::SeqCst) == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
let (lock, cvar) = &*self.pending_pair;
|
||||
let mut guard = lock.lock_blocking();
|
||||
while self.pending.load(Ordering::SeqCst) != 0 {
|
||||
guard = match cvar.wait(guard) {
|
||||
Ok(g) => g,
|
||||
Err(poisoned) => poisoned.into_inner(),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/// Take all results, leaving an empty results vector.
|
||||
pub fn take_results(&self) -> Vec<T> {
|
||||
let mut results = self.results.lock_blocking();
|
||||
results.split_off(0)
|
||||
}
|
||||
|
||||
/// Wait until all submitted tasks have completed, then return the results.
|
||||
pub fn into_results(self) -> Vec<T> {
|
||||
self.join();
|
||||
let mut results = self.results.lock_blocking();
|
||||
results.split_off(0)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Send + 'static> Drop for ThreadPool<T> {
|
||||
fn drop(&mut self) {
|
||||
// Close sender so worker threads exit recv loop
|
||||
self.sender.take();
|
||||
// Dropping the sender (SyncSender) happens above; but to ensure we close the channel we
|
||||
// explicitly drop any remaining clones by letting sender go out of scope.
|
||||
|
||||
// Join worker threads
|
||||
while let Some(handle) = self.workers.pop() {
|
||||
let _ = handle.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user