Fix dict decompress usage for uk-editor

NiceneNerd · Jul 4, 2024 · 09444e8 · 09444e8
1 parent 78ad2cc
commit 09444e8
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 61 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/uk-editor/src/project.rs b/crates/uk-editor/src/project.rs
@@ -8,20 +8,20 @@ use fs_err as fs;
 use rayon::prelude::*;
 use uk_content::{resource::ResourceData, util::IndexMap};
 use uk_manager::{core::Manager, settings::Platform};
-use uk_mod::{pack::sanitise, unpack::ParallelZipReader, Meta};
+use uk_mod::{pack::sanitise, unpack::ParallelZipReader, zstd::zstd_safe::WriteBuf, Meta};
 
 #[derive(Debug, Clone)]
 pub struct Project {
-    pub path:  PathBuf,
-    pub meta:  Meta,
+    pub path: PathBuf,
+    pub meta: Meta,
     pub files: BTreeSet<PathBuf>,
 }
 
 impl Project {
     pub fn new(name: &str, path: &Path, platform: Platform) -> Self {
         Project {
-            path:  path.join(name),
-            meta:  Meta {
+            path: path.join(name),
+            meta: Meta {
                 api: env!("CARGO_PKG_VERSION").into(),
                 name: name.into(),
                 author: Default::default(),
@@ -75,6 +75,7 @@ impl Project {
         )
         .context("Failed to parse mod meta")?;
         let path = core.settings().projects_dir().join(sanitise(&meta.name));
+        let decomp = uk_mod::unpack::init_decompressor();
         let files = zip
             .iter()
             .par_bridge()
@@ -85,31 +86,34 @@ impl Project {
                     let data = zip.get_file(file).with_context(|| {
                         format!("Failed to read file {} from ZIP", file.display())
                     })?;
-                    if matches!(
-                        file.file_name()
-                            .unwrap_or_default()
-                            .to_str()
-                            .unwrap_or_default(),
-                        "meta.yml" | "manifest.yml"
-                    ) {
+                    let file_name = file
+                        .file_name()
+                        .unwrap_or_default()
+                        .to_str()
+                        .unwrap_or_default();
+                    if file_name.ends_with(".yml") || file_name.starts_with("thumb.") {
                         fs::write(dest, data).with_context(|| {
                             format!("Failed to extract file {}", file.display())
                         })?;
                         return Ok(None);
                     }
-                    let resource: ResourceData = minicbor_ser::from_slice(
-                        &uk_mod::zstd::decode_all(data.as_slice()).with_context(|| {
+                    let decomp_size =
+                        uk_mod::zstd::bulk::Decompressor::upper_bound(data.as_slice())
+                            .unwrap_or(data.len() * 1024);
+                    let decomp_data = decomp
+                        .lock()
+                        .decompress(data.as_slice(), decomp_size)
+                        .or_else(|_| uk_mod::zstd::decode_all(data.as_slice()))
+                        .with_context(|| {
                             format!("Failed to decompress contents of {} in ZIP", file.display())
-                        })?,
-                    )
-                    .with_context(|| format!("Failed to parse resource {}", file.display()))?;
+                        })?;
+                    let resource: ResourceData = minicbor_ser::from_slice(&decomp_data)
+                        .with_context(|| format!("Failed to parse resource {}", file.display()))?;
                     let data = match resource {
                         ResourceData::Binary(bin) => bin,
-                        res => {
-                            ron::ser::to_string_pretty(&res, Default::default())
-                                .expect("Failed to serialize resource to RON")
-                                .into()
-                        }
+                        res => ron::ser::to_string_pretty(&res, Default::default())
+                            .expect("Failed to serialize resource to RON")
+                            .into(),
                     };
                     fs::write(dest, data)
                         .with_context(|| format!("Failed to extract file {}", file.display()))?;

diff --git a/crates/uk-mod/src/lib.rs b/crates/uk-mod/src/lib.rs
@@ -22,7 +22,7 @@ pub struct Manifest {
     #[serde(rename = "content")]
     pub content_files: BTreeSet<String>,
     #[serde(rename = "aoc")]
-    pub aoc_files:     BTreeSet<String>,
+    pub aoc_files: BTreeSet<String>,
 }
 
 impl Manifest {

diff --git a/crates/uk-mod/src/unpack.rs b/crates/uk-mod/src/unpack.rs
@@ -77,10 +77,10 @@ impl std::fmt::Debug for ZipData {
 
 #[self_referencing]
 pub struct ParallelZipReader {
-    data:  ZipData,
+    data: ZipData,
     #[borrows(data)]
     #[covariant]
-    zip:   piz::ZipArchive<'this>,
+    zip: piz::ZipArchive<'this>,
     #[borrows(zip)]
     #[covariant]
     files: HashMap<&'this Path, &'this piz::read::FileMetadata<'this>>,
@@ -153,7 +153,7 @@ impl ParallelZipReader {
 }
 
 #[inline]
-fn init_decompressor() -> Arc<Mutex<zstd::bulk::Decompressor<'static>>> {
+pub fn init_decompressor() -> Arc<Mutex<zstd::bulk::Decompressor<'static>>> {
     Arc::new(Mutex::new(
         zstd::bulk::Decompressor::with_dictionary(super::DICTIONARY).unwrap(),
     ))
@@ -477,14 +477,14 @@ static RSTB_EXCLUDE_NAMES: &[&str] = &["ActorInfo.product.byml"];
 
 // #[derive(Debug)]
 pub struct ModUnpacker {
-    dump:     Arc<ResourceReader>,
+    dump: Arc<ResourceReader>,
     manifest: Option<Manifest>,
-    mods:     Vec<ModReader>,
-    endian:   Endian,
-    lang:     Language,
-    rstb:     DashMap<String, Option<u32>>,
-    hashes:   StockHashTable,
-    out_dir:  PathBuf,
+    mods: Vec<ModReader>,
+    endian: Endian,
+    lang: Language,
+    rstb: DashMap<String, Option<u32>>,
+    hashes: StockHashTable,
+    out_dir: PathBuf,
 }
 
 impl ModUnpacker {
@@ -566,19 +566,18 @@ impl ModUnpacker {
                     Ok(Err(e)) => anyhow_ext::bail!(e),
                     Ok(Ok(_)) => (),
                     Err(e) => {
-                        anyhow::bail!(
-                            e.downcast::<std::string::String>()
-                                .or_else(|e| {
-                                    e.downcast::<&'static str>().map(|s| Box::new((*s).into()))
-                                })
-                                .unwrap_or_else(|_| {
-                                    Box::new(
-                                        "An unknown error occured, check the log for possible \
+                        anyhow::bail!(e
+                            .downcast::<std::string::String>()
+                            .or_else(|e| {
+                                e.downcast::<&'static str>().map(|s| Box::new((*s).into()))
+                            })
+                            .unwrap_or_else(|_| {
+                                Box::new(
+                                    "An unknown error occured, check the log for possible \
                                          details."
-                                            .to_string(),
-                                    )
-                                })
-                        )
+                                        .to_string(),
+                                )
+                            }))
                     }
                 }
             }