Fix huffman encode problem

src/scripts/artemis/asb.rs

@@ -4,6 +4,7 @@ use crate::types::*;
 use crate::utils::encoding::*;
 use crate::utils::escape::*;
 use anyhow::Result;
+use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::io::{Read, Write};
 use std::ops::Index;
@@ -49,6 +50,20 @@ impl ScriptBuilder for ArtemisAsbBuilder {
         }
         None
     }
+
+    fn can_create_file(&self) -> bool {
+        true
+    }
+
+    fn create_file<'a>(
+        &'a self,
+        filename: &'a str,
+        writer: Box<dyn WriteSeek + 'a>,
+        encoding: Encoding,
+        file_encoding: Encoding,
+    ) -> Result<()> {
+        create_file(filename, writer, encoding, file_encoding)
+    }
 }
 
 fn escape_text(s: &str) -> String {
@@ -63,7 +78,7 @@ fn escape_text(s: &str) -> String {
     escaped
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 struct Command {
     pub name: String,
     pub line_number: u32,
@@ -110,7 +125,8 @@ impl Index<String> for Command {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
 enum Item {
     Command(Command),
     Label(String),
@@ -212,10 +228,11 @@ struct TextParser<'a> {
     text: Vec<&'a str>,
     pos: usize,
     len: usize,
+    hcls_index: usize,
 }
 
 impl<'a> TextParser<'a> {
-    pub fn new(str: &'a str) -> Self {
+    pub fn new(str: &'a str, hcls_index: usize) -> Self {
         let text: Vec<&'a str> = UnicodeSegmentation::graphemes(str, true).collect();
         let len = text.len();
         TextParser {
@@ -223,6 +240,7 @@ impl<'a> TextParser<'a> {
             text,
             pos: 0,
             len,
+            hcls_index,
         }
     }
 
@@ -253,8 +271,10 @@ impl<'a> TextParser<'a> {
                 }
             }
         }
-        self.items
-            .push(Item::Command(Command::new("hcls".to_string(), 0)));
+        let mut hcls = Command::new("hcls".to_string(), 0);
+        hcls.attributes
+            .insert("0".to_string(), self.hcls_index.to_string());
+        self.items.push(Item::Command(hcls));
         Ok(self.items)
     }
 
@@ -413,10 +433,18 @@ impl Script for Asb {
         OutputScriptType::Json
     }
 
+    fn is_output_supported(&self, _: OutputScriptType) -> bool {
+        true
+    }
+
     fn default_format_type(&self) -> FormatOptions {
         FormatOptions::None
     }
 
+    fn custom_output_extension<'a>(&'a self) -> &'a str {
+        "json"
+    }
+
     fn extract_messages(&self) -> Result<Vec<Message>> {
         let mut messages = Vec::new();
         let mut name = None;
@@ -501,6 +529,7 @@ impl Script for Asb {
         let mut mes_index = 0;
         let mut item_index = 0;
         let mut print_index = None;
+        let mut hcls_index = 1;
         while item_index < items.len() {
             if let Some(print_ind) = print_index.clone() {
                 if items[item_index].is_command_name("hcls") {
@@ -538,7 +567,8 @@ impl Script for Asb {
                             m = m.replace(k, v);
                         }
                     }
-                    let new_cmds = TextParser::new(&m.replace("\n", "<rt>")).parse()?;
+                    let new_cmds = TextParser::new(&m.replace("\n", "<rt>"), hcls_index).parse()?;
+                    hcls_index += 1;
                     let new_cmds_len = new_cmds.len();
                     items.splice(print_ind..=item_index, new_cmds);
                     print_index = None;
@@ -603,12 +633,47 @@ impl Script for Asb {
         file.flush()?;
         Ok(())
     }
+
+    fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
+        let s = serde_json::to_string_pretty(&self.items)?;
+        let s = encode_string(encoding, &s, false)?;
+        let mut file = std::fs::File::create(filename)?;
+        file.write_all(&s)?;
+        Ok(())
+    }
+
+    fn custom_import<'a>(
+        &'a self,
+        custom_filename: &'a str,
+        file: Box<dyn WriteSeek + 'a>,
+        encoding: Encoding,
+        output_encoding: Encoding,
+    ) -> Result<()> {
+        create_file(custom_filename, file, encoding, output_encoding)
+    }
+}
+
+pub fn create_file<'a>(
+    custom_filename: &'a str,
+    mut writer: Box<dyn WriteSeek + 'a>,
+    encoding: Encoding,
+    output_encoding: Encoding,
+) -> Result<()> {
+    let f = crate::utils::files::read_file(custom_filename)?;
+    let s = decode_to_string(output_encoding, &f, true)?;
+    let items: Vec<Item> = serde_json::from_str(&s)?;
+    writer.write_all(b"ASB\0\0")?;
+    writer.write_u32(items.len() as u32)?;
+    for item in items {
+        writer.write_item(&item, encoding)?;
+    }
+    Ok(())
 }
 
 #[test]
 fn test_parse() {
     let text = "Hello &lt; &amp; World!<tag><tags x=\"123\"><name 0=\"Ok\">Test";
-    let parser = TextParser::new(text);
+    let parser = TextParser::new(text, 1);
     let items = parser.parse().unwrap();
     assert_eq!(
         items,
@@ -641,7 +706,7 @@ fn test_parse() {
             Item::Command(Command {
                 name: "hcls".to_string(),
                 line_number: 0,
-                attributes: BTreeMap::new(),
+                attributes: BTreeMap::from([("0".to_string(), "1".to_string())]),
             }),
         ]
     )

src/scripts/bgi/archive/dsc.rs

@@ -241,26 +241,70 @@ impl Ord for FreqNode {
 }
 
 fn calculate_huffman_depths(freqs: &[u32]) -> Vec<u8> {
-    let mut heap = BinaryHeap::new();
-    for (symbol, &freq) in freqs.iter().enumerate() {
-        if freq > 0 {
-            heap.push(FreqNode {
-                freq,
-                symbol: Some(symbol as u16),
-                left: None,
-                right: None,
-            });
-        }
-    }
+    const MAX_DEPTH: u8 = 9;
 
-    if heap.len() <= 1 {
-        let mut depths = vec![0; 512];
-        if let Some(node) = heap.pop() {
-            depths[node.symbol.unwrap() as usize] = 1;
-        }
+    // 收集所有非零频率的符号
+    let mut symbols_with_freq: Vec<(u16, u32)> = freqs
+        .iter()
+        .enumerate()
+        .filter_map(|(symbol, &freq)| {
+            if freq > 0 {
+                Some((symbol as u16, freq))
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    let mut depths = vec![0u8; 512];
+
+    if symbols_with_freq.is_empty() {
         return depths;
     }
 
+    if symbols_with_freq.len() == 1 {
+        depths[symbols_with_freq[0].0 as usize] = 1;
+        return depths;
+    }
+
+    // 使用受限Huffman算法
+    loop {
+        let current_depths = build_huffman_tree(&symbols_with_freq);
+        let max_depth = current_depths.iter().max().copied().unwrap_or(0);
+
+        if max_depth <= MAX_DEPTH {
+            // 将深度映射回原始数组
+            for &(symbol, _) in &symbols_with_freq {
+                let symbol_index = symbols_with_freq
+                    .iter()
+                    .position(|(s, _)| *s == symbol)
+                    .unwrap();
+                depths[symbol as usize] = current_depths[symbol_index];
+            }
+            break;
+        }
+
+        // 如果深度超限，调整频率
+        adjust_frequencies_for_depth_limit(&mut symbols_with_freq);
+    }
+
+    depths
+}
+
+fn build_huffman_tree(symbols_with_freq: &[(u16, u32)]) -> Vec<u8> {
+    let mut heap = BinaryHeap::new();
+
+    // 添加所有叶子节点
+    for &(symbol, freq) in symbols_with_freq {
+        heap.push(FreqNode {
+            freq,
+            symbol: Some(symbol),
+            left: None,
+            right: None,
+        });
+    }
+
+    // 构建Huffman树
     while heap.len() > 1 {
         let node1 = heap.pop().unwrap();
         let node2 = heap.pop().unwrap();
@@ -273,29 +317,53 @@ fn calculate_huffman_depths(freqs: &[u32]) -> Vec<u8> {
         heap.push(new_node);
     }
 
-    let mut depths = vec![0; 512];
+    // 计算深度
+    let mut depths = vec![0u8; symbols_with_freq.len()];
     if let Some(root) = heap.pop() {
-        fn traverse(node: &FreqNode, depth: u8, depths: &mut [u8]) {
-            if let Some(symbol) = node.symbol {
-                if depth == 0 {
-                    depths[symbol as usize] = 1;
-                } else {
-                    depths[symbol as usize] = depth;
-                }
-            } else {
-                if let Some(ref left) = node.left {
-                    traverse(left, depth + 1, depths);
-                }
-                if let Some(ref right) = node.right {
-                    traverse(right, depth + 1, depths);
-                }
-            }
-        }
-        traverse(&root, 0, &mut depths);
+        calculate_depths(&root, 0, symbols_with_freq, &mut depths);
     }
+
     depths
 }
 
+fn calculate_depths(
+    node: &FreqNode,
+    depth: u8,
+    symbols_with_freq: &[(u16, u32)],
+    depths: &mut [u8],
+) {
+    if let Some(symbol) = node.symbol {
+        let symbol_index = symbols_with_freq
+            .iter()
+            .position(|(s, _)| *s == symbol)
+            .unwrap();
+        depths[symbol_index] = if depth == 0 { 1 } else { depth };
+    } else {
+        if let Some(ref left) = node.left {
+            calculate_depths(left, depth + 1, symbols_with_freq, depths);
+        }
+        if let Some(ref right) = node.right {
+            calculate_depths(right, depth + 1, symbols_with_freq, depths);
+        }
+    }
+}
+
+fn adjust_frequencies_for_depth_limit(symbols_with_freq: &mut [(u16, u32)]) {
+    // 按频率排序
+    symbols_with_freq.sort_by(|a, b| a.1.cmp(&b.1));
+
+    // 使用Package-Merge算法的简化版本
+    // 这里使用一个启发式方法：增加低频符号的频率
+    let min_freq = symbols_with_freq[0].1;
+    let adjustment = (min_freq as f64 * 0.1).max(1.0) as u32;
+
+    // 找到频率最低的几个符号并调整它们的频率
+    let num_to_adjust = (symbols_with_freq.len() / 4).max(1);
+    for i in 0..num_to_adjust.min(symbols_with_freq.len()) {
+        symbols_with_freq[i].1 += adjustment;
+    }
+}
+
 fn generate_canonical_codes(depths: &[u8]) -> Vec<Option<(u16, u8)>> {
     let mut codes_with_depths = vec![];
     for (symbol, &depth) in depths.iter().enumerate() {
@@ -575,7 +643,7 @@ impl Script for Dsc {
     }
 
     fn custom_output_extension(&self) -> &'static str {
-        "unk"
+        ""
     }
 
     fn custom_export(&self, filename: &std::path::Path, _encoding: Encoding) -> Result<()> {