Add CRX import support

src/args.rs

@@ -16,6 +16,17 @@ fn parse_compression_level(level: &str) -> Result<u32, String> {
     clap_num::number_range(level, 0, 9)
 }
 
+#[cfg(feature = "zstd")]
+fn parse_zstd_compression_level(level: &str) -> Result<i32, String> {
+    let lower = level.to_ascii_lowercase();
+    if lower == "default" {
+        return Ok(3);
+    } else if lower == "best" {
+        return Ok(22);
+    }
+    clap_num::number_range(level, 0, 22)
+}
+
 /// Tools for export and import scripts
 #[derive(Parser, Debug)]
 #[clap(group = ArgGroup::new("encodingg").multiple(false), group = ArgGroup::new("output_encodingg").multiple(false), group = ArgGroup::new("archive_encodingg").multiple(false), group = ArgGroup::new("artemis_indentg").multiple(false))]
@@ -200,12 +211,24 @@ pub struct Arg {
     /// If not specified, the first language will be used.
     pub cat_system_cstl_lang: Option<String>,
     #[cfg(feature = "flate2")]
-    #[arg(short = 'z', long, global = true, value_name = "LEVEL", value_parser = parse_compression_level)]
+    #[arg(short = 'z', long, global = true, value_name = "LEVEL", value_parser = parse_compression_level, default_value_t = 6)]
-    /// Zlib compression level. Default is 6. 0 means no compression, 9 means best compression.
+    /// Zlib compression level. 0 means no compression, 9 means best compression.
-    pub zlib_compression_level: Option<u32>,
+    pub zlib_compression_level: u32,
     #[cfg(feature = "image")]
     #[arg(short = 'g', long, global = true, value_enum, default_value_t = PngCompressionLevel::Fast)]
     pub png_compression_level: PngCompressionLevel,
+    #[cfg(feature = "circus-img")]
+    #[arg(long, global = true, action = ArgAction::SetTrue)]
+    /// Keep original BPP when importing Circus CRX images.
+    pub circus_crx_keep_original_bpp: bool,
+    #[cfg(feature = "circus-img")]
+    #[arg(long, global = true, action = ArgAction::SetTrue)]
+    /// Use zstd compression for Circus CRX images. (CIRCUS Engine don't support this. Hook is required.)
+    pub circus_crx_zstd: bool,
+    #[cfg(feature = "zstd")]
+    #[arg(short = 'Z', long, global = true, value_name = "LEVEL", value_parser = parse_zstd_compression_level, default_value_t = 3)]
+    /// Zstd compression level. 0 means default compression level (3), 22 means best compression.
+    pub zstd_compression_level: i32,
     #[command(subcommand)]
     /// Command
     pub command: Command,

src/main.rs

@@ -1459,6 +1459,12 @@ fn main() {
         zlib_compression_level: arg.zlib_compression_level,
         #[cfg(feature = "image")]
         png_compression_level: arg.png_compression_level,
+        #[cfg(feature = "circus-img")]
+        circus_crx_keep_original_bpp: arg.circus_crx_keep_original_bpp,
+        #[cfg(feature = "circus-img")]
+        circus_crx_zstd: arg.circus_crx_zstd,
+        #[cfg(feature = "zstd")]
+        zstd_compression_level: arg.zstd_compression_level,
     };
     match &arg.command {
         args::Command::Export { input, output } => {

src/scripts/cat_system/cst.rs

@@ -169,7 +169,7 @@ impl CstScript {
             data: file,
             compressed: compressed_size != 0,
             strings,
-            compress_level: config.zlib_compression_level.unwrap_or(6),
+            compress_level: config.zlib_compression_level,
         })
     }
 }

src/scripts/circus/image/crx.rs

@@ -1,6 +1,7 @@
 use crate::ext::io::*;
 use crate::scripts::base::*;
 use crate::types::*;
+use crate::utils::img::*;
 use crate::utils::struct_pack::*;
 use anyhow::Result;
 use msg_tool_macro::*;
@@ -82,6 +83,10 @@ pub struct CrxImage {
     header: Header,
     color_type: ImageColorType,
     data: Vec<u8>,
+    compress_level: u32,
+    keep_original_bpp: bool,
+    zstd: bool,
+    zstd_compression_level: i32,
 }
 
 impl std::fmt::Debug for CrxImage {
@@ -95,7 +100,7 @@ impl std::fmt::Debug for CrxImage {
 }
 
 impl CrxImage {
-    pub fn new<T: Read + Seek>(data: T, _config: &ExtraConfig) -> Result<Self> {
+    pub fn new<T: Read + Seek>(data: T, config: &ExtraConfig) -> Result<Self> {
         let mut reader = data;
         let mut magic = [0; 4];
         reader.read_exact(&mut magic)?;
@@ -138,6 +143,10 @@ impl CrxImage {
             header,
             color_type,
             data: uncompessed,
+            compress_level: config.zlib_compression_level,
+            keep_original_bpp: config.circus_crx_keep_original_bpp,
+            zstd: config.circus_crx_zstd,
+            zstd_compression_level: config.zstd_compression_level,
         })
     }
 
@@ -336,6 +345,201 @@ impl CrxImage {
         }
         Ok(())
     }
+
+    fn encode_row0(
+        dst: &mut Vec<u8>,
+        mut dst_p: usize,
+        src: &[u8],
+        width: u16,
+        pixel_size: u8,
+        y: u16,
+    ) -> Result<usize> {
+        let pixel_size = pixel_size as usize;
+        let mut src_p = y as usize * width as usize * pixel_size;
+        for _ in 0..pixel_size {
+            dst[dst_p] = src[src_p];
+            dst_p += 1;
+            src_p += 1;
+        }
+        for _ in 1..width {
+            for _ in 0..pixel_size {
+                dst[dst_p] = src[src_p].wrapping_sub(src[src_p - pixel_size]);
+                dst_p += 1;
+                src_p += 1;
+            }
+        }
+        Ok(dst_p)
+    }
+
+    fn encode_row1(
+        dst: &mut Vec<u8>,
+        mut dst_p: usize,
+        src: &[u8],
+        width: u16,
+        pixel_size: u8,
+        y: u16,
+    ) -> Result<usize> {
+        let pixel_size = pixel_size as usize;
+        let mut src_p = y as usize * width as usize * pixel_size;
+        let mut prev_row_p = (y as usize - 1) * width as usize * pixel_size;
+        for _ in 0..width {
+            for _ in 0..pixel_size {
+                dst[dst_p] = src[src_p].wrapping_sub(src[prev_row_p]);
+                dst_p += 1;
+                src_p += 1;
+                prev_row_p += 1;
+            }
+        }
+        Ok(dst_p)
+    }
+
+    fn encode_row2(
+        dst: &mut Vec<u8>,
+        mut dst_p: usize,
+        src: &[u8],
+        width: u16,
+        pixel_size: u8,
+        y: u16,
+    ) -> Result<usize> {
+        let pixel_size = pixel_size as usize;
+        let mut src_p = y as usize * width as usize * pixel_size;
+        let mut prev_row_p = (y as usize - 1) * width as usize * pixel_size;
+        for _ in 0..pixel_size {
+            dst[dst_p] = src[src_p];
+            dst_p += 1;
+            src_p += 1;
+        }
+        for _ in 1..width {
+            for _ in 0..pixel_size {
+                dst[dst_p] = src[src_p].wrapping_sub(src[prev_row_p]);
+                dst_p += 1;
+                src_p += 1;
+                prev_row_p += 1;
+            }
+        }
+        Ok(dst_p)
+    }
+
+    fn encode_row3(
+        dst: &mut Vec<u8>,
+        mut dst_p: usize,
+        src: &[u8],
+        width: u16,
+        pixel_size: u8,
+        y: u16,
+    ) -> Result<usize> {
+        let pixel_size = pixel_size as usize;
+        let mut src_p = y as usize * width as usize * pixel_size;
+        let mut prev_row_p = (y as usize - 1) * width as usize * pixel_size + pixel_size;
+        for _ in 0..width - 1 {
+            for _ in 0..pixel_size {
+                dst[dst_p] = src[src_p].wrapping_sub(src[prev_row_p]);
+                dst_p += 1;
+                src_p += 1;
+                prev_row_p += 1;
+            }
+        }
+        for _ in 0..pixel_size {
+            dst[dst_p] = src[src_p];
+            dst_p += 1;
+            src_p += 1;
+        }
+        Ok(dst_p)
+    }
+
+    fn encode_row4(
+        dst: &mut Vec<u8>,
+        mut dst_p: usize,
+        src: &[u8],
+        width: u16,
+        pixel_size: u8,
+        y: u16,
+    ) -> Result<usize> {
+        let pixel_size = pixel_size as usize;
+        let src_p = y as usize * width as usize * pixel_size;
+        for offset in 0..pixel_size {
+            let mut src_c = src_p + offset;
+            let mut remaining = width;
+            let value = src[src_c];
+            src_c += pixel_size;
+            dst[dst_p] = value;
+            dst_p += 1;
+            remaining -= 1;
+            if remaining == 0 {
+                continue;
+            }
+            let mut count = 0;
+            loop {
+                if count as u16 >= remaining || count >= 255 || src[src_c] != value {
+                    break;
+                }
+                src_c += pixel_size;
+                count += 1;
+            }
+            if count > 0 {
+                dst[dst_p] = value;
+                dst_p += 1;
+                dst[dst_p] = count;
+                dst_p += 1;
+                remaining -= count as u16;
+            }
+            while remaining > 0 {
+                let value = src[src_c];
+                src_c += pixel_size;
+                dst[dst_p] = value;
+                dst_p += 1;
+                remaining -= 1;
+                if remaining == 0 {
+                    break;
+                }
+                let mut count = 0;
+                loop {
+                    if count as u16 >= remaining || count >= 255 || src[src_c] != value {
+                        break;
+                    }
+                    src_c += pixel_size;
+                    count += 1;
+                }
+                if count > 0 {
+                    dst[dst_p] = value;
+                    dst_p += 1;
+                    dst[dst_p] = count;
+                    dst_p += 1;
+                    remaining -= count as u16;
+                }
+            }
+        }
+        Ok(dst_p)
+    }
+
+    fn encode_image_origin(
+        src: &[u8],
+        width: u16,
+        height: u16,
+        pixel_size: u8,
+        row_type: &[u8],
+    ) -> Result<Vec<u8>> {
+        if row_type.len() != height as usize {
+            return Err(anyhow::anyhow!("Row type length does not match height"));
+        }
+        let size = width as usize * height as usize * pixel_size as usize + height as usize;
+        let mut dst = vec![0; size];
+        let mut dst_p = 0;
+        for y in 0..height {
+            let data = row_type[y as usize];
+            dst[dst_p] = data;
+            dst_p += 1;
+            dst_p = match data {
+                0 => Self::encode_row0(&mut dst, dst_p, src, width, pixel_size, y)?,
+                1 => Self::encode_row1(&mut dst, dst_p, src, width, pixel_size, y)?,
+                2 => Self::encode_row2(&mut dst, dst_p, src, width, pixel_size, y)?,
+                3 => Self::encode_row3(&mut dst, dst_p, src, width, pixel_size, y)?,
+                4 => Self::encode_row4(&mut dst, dst_p, src, width, pixel_size, y)?,
+                _ => return Err(anyhow::anyhow!("Invalid row type: {}", data)),
+            };
+        }
+        Ok(dst)
+    }
 }
 
 impl Script for CrxImage {
@@ -386,4 +590,121 @@ impl Script for CrxImage {
             data,
         })
     }
+
+    fn import_image<'a>(
+        &'a self,
+        mut data: ImageData,
+        mut file: Box<dyn WriteSeek + 'a>,
+    ) -> Result<()> {
+        let mut color_type = match data.color_type {
+            ImageColorType::Bgr => ImageColorType::Bgr,
+            ImageColorType::Bgra => ImageColorType::Bgra,
+            ImageColorType::Rgb => {
+                convert_rgb_to_bgr(&mut data)?;
+                ImageColorType::Bgr
+            }
+            ImageColorType::Rgba => {
+                convert_rgba_to_bgra(&mut data)?;
+                ImageColorType::Bgra
+            }
+            _ => {
+                return Err(anyhow::anyhow!(
+                    "Unsupported color type: {:?}",
+                    data.color_type
+                ));
+            }
+        };
+        if data.width != self.header.width as u32 {
+            return Err(anyhow::anyhow!(
+                "Image width does not match: expected {}, got {}",
+                self.header.width,
+                data.width
+            ));
+        }
+        if data.height != self.header.height as u32 {
+            return Err(anyhow::anyhow!(
+                "Image height does not match: expected {}, got {}",
+                self.header.height,
+                data.height
+            ));
+        }
+        if data.depth != 8 {
+            return Err(anyhow::anyhow!("Image depth must be 8, got {}", data.depth));
+        }
+        if data.color_type != self.color_type && self.keep_original_bpp {
+            if self.color_type == ImageColorType::Bgr {
+                convert_bgra_to_bgr(&mut data)?;
+            } else if self.color_type == ImageColorType::Bgra {
+                convert_bgr_to_bgra(&mut data)?;
+            } else {
+                return Err(anyhow::anyhow!(
+                    "Unsupported color type for import: {:?}",
+                    self.color_type
+                ));
+            }
+            color_type = self.color_type;
+        }
+        let mut new_header = self.header.clone();
+        new_header.bpp = match color_type {
+            ImageColorType::Bgr => 0,
+            ImageColorType::Bgra => 1,
+            _ => return Err(anyhow::anyhow!("Unsupported color type: {:?}", color_type)),
+        };
+        new_header.flags |= 0x10; // Force add compressed data length
+        let pixel_size = color_type.bpp(1) as u8;
+        if color_type == ImageColorType::Bgra && self.header.mode != 1 {
+            let alpha_flip = if self.header.mode == 2 { 0 } else { 0xFF };
+            for i in (0..data.data.len()).step_by(4) {
+                let b = data.data[i];
+                let g = data.data[i + 1];
+                let r = data.data[i + 2];
+                let a = data.data[i + 3];
+                data.data[i] = a ^ alpha_flip;
+                data.data[i + 1] = b;
+                data.data[i + 2] = g;
+                data.data[i + 3] = r;
+            }
+        }
+        let mut row_type = Vec::with_capacity(self.header.height as usize);
+        let mut dst = vec![
+            0;
+            self.header.width as usize
+                * self.header.height as usize
+                * self.color_type.bpp(1) as usize
+        ];
+        Self::decode_image(
+            &mut dst,
+            &self.data,
+            self.header.width,
+            self.header.height,
+            self.color_type.bpp(1) as u8,
+            &mut row_type,
+        )?;
+        let encoded = Self::encode_image_origin(
+            &data.data,
+            new_header.width,
+            new_header.height,
+            pixel_size,
+            &row_type,
+        )?;
+        let compressed = if self.zstd {
+            let mut encoder = zstd::Encoder::new(MemWriter::new(), self.zstd_compression_level)?;
+            encoder.write_all(&encoded)?;
+            let compressed_data = encoder.finish()?;
+            compressed_data.into_inner()
+        } else {
+            let mut encoder = flate2::write::ZlibEncoder::new(
+                MemWriter::new(),
+                flate2::Compression::new(self.compress_level),
+            );
+            encoder.write_all(&encoded)?;
+            let compressed_data = encoder.finish()?;
+            compressed_data.into_inner()
+        };
+        file.write_all(b"CRXG")?;
+        new_header.pack(&mut file, false, Encoding::Utf8)?;
+        file.write_u32(compressed.len() as u32)?;
+        file.write_all(&compressed)?;
+        Ok(())
+    }
 }

src/types.rs

@@ -236,9 +236,15 @@ pub struct ExtraConfig {
     #[cfg(feature = "cat-system")]
     pub cat_system_cstl_lang: Option<String>,
     #[cfg(feature = "flate2")]
-    pub zlib_compression_level: Option<u32>,
+    pub zlib_compression_level: u32,
     #[cfg(feature = "image")]
     pub png_compression_level: PngCompressionLevel,
+    #[cfg(feature = "circus-img")]
+    pub circus_crx_keep_original_bpp: bool,
+    #[cfg(feature = "circus-img")]
+    pub circus_crx_zstd: bool,
+    #[cfg(feature = "zstd")]
+    pub zstd_compression_level: i32,
 }
 
 #[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq, PartialOrd, Ord)]

src/utils/img.rs

@@ -16,6 +16,27 @@ pub fn reverse_alpha_values(data: &mut ImageData) -> Result<()> {
     Ok(())
 }
 
+pub fn convert_bgr_to_bgra(data: &mut ImageData) -> Result<()> {
+    if data.color_type != ImageColorType::Bgr {
+        return Err(anyhow::anyhow!("Image is not BGR"));
+    }
+    if data.depth != 8 {
+        return Err(anyhow::anyhow!(
+            "BGR to BGRA conversion only supports 8-bit depth"
+        ));
+    }
+    let mut new_data = Vec::with_capacity(data.data.len() / 3 * 4);
+    for chunk in data.data.chunks_exact(3) {
+        new_data.push(chunk[0]); // B
+        new_data.push(chunk[1]); // G
+        new_data.push(chunk[2]); // R
+        new_data.push(255); // A
+    }
+    data.data = new_data;
+    data.color_type = ImageColorType::Bgra;
+    Ok(())
+}
+
 pub fn convert_bgr_to_rgb(data: &mut ImageData) -> Result<()> {
     if data.color_type != ImageColorType::Bgr {
         return Err(anyhow::anyhow!("Image is not BGR"));
@@ -34,6 +55,26 @@ pub fn convert_bgr_to_rgb(data: &mut ImageData) -> Result<()> {
     Ok(())
 }
 
+pub fn convert_bgra_to_bgr(data: &mut ImageData) -> Result<()> {
+    if data.color_type != ImageColorType::Bgra {
+        return Err(anyhow::anyhow!("Image is not BGRA"));
+    }
+    if data.depth != 8 {
+        return Err(anyhow::anyhow!(
+            "BGRA to BGR conversion only supports 8-bit depth"
+        ));
+    }
+    let mut new_data = Vec::with_capacity(data.data.len() / 4 * 3);
+    for chunk in data.data.chunks_exact(4) {
+        new_data.push(chunk[0]); // B
+        new_data.push(chunk[1]); // G
+        new_data.push(chunk[2]); // R
+    }
+    data.data = new_data;
+    data.color_type = ImageColorType::Bgr;
+    Ok(())
+}
+
 pub fn convert_bgra_to_rgba(data: &mut ImageData) -> Result<()> {
     if data.color_type != ImageColorType::Bgra {
         return Err(anyhow::anyhow!("Image is not BGRA"));