tiger_pkg/manager/
mod.rs

1pub mod lookup_cache;
2pub mod path_cache;
3
4use std::{
5    fmt::Display,
6    fs,
7    io::Cursor,
8    path::{Path, PathBuf},
9    str::FromStr,
10    sync::Arc,
11};
12
13use ahash::HashMap;
14use anyhow::Context;
15use binrw::{BinRead, BinReaderExt};
16use parking_lot::RwLock;
17use rayon::prelude::*;
18use tracing::{debug_span, info, warn};
19
20use crate::{
21    d2_shared::PackageNamedTagEntry,
22    oodle,
23    package::{Package, PackagePlatform, Redaction, UEntryHeader},
24    tag::TagHash64,
25    GameVersion, TagHash, Version,
26};
27
28#[derive(Clone, bincode::Decode, bincode::Encode)]
29pub struct HashTableEntryShort {
30    pub hash32: TagHash,
31    pub reference: TagHash,
32}
33
34#[derive(Default, bincode::Decode, bincode::Encode)]
35pub struct TagLookupIndex {
36    pub tag32_entries_by_pkg: HashMap<u16, Vec<UEntryHeader>>,
37    pub tag64_entries: HashMap<u64, HashTableEntryShort>,
38    pub tag32_to_tag64: HashMap<TagHash, TagHash64>,
39
40    pub named_tags: Vec<PackageNamedTagEntry>,
41    pub redaction_levels: HashMap<u16, Redaction>,
42}
43
44pub struct PackageManager {
45    pub package_dir: PathBuf,
46    pub package_paths: HashMap<u16, PackagePath>,
47    pub version: GameVersion,
48    pub platform: PackagePlatform,
49
50    /// Tag Lookup Index (TLI)
51    pub lookup: TagLookupIndex,
52
53    /// Packages that are currently open for reading
54    pkgs: RwLock<HashMap<u16, Arc<dyn Package>>>,
55}
56
57impl PackageManager {
58    pub fn new<P: AsRef<Path>>(
59        packages_dir: P,
60        version: GameVersion,
61        platform: Option<PackagePlatform>,
62    ) -> anyhow::Result<PackageManager> {
63        // All the latest packages
64        let mut packages: HashMap<u16, String> = Default::default();
65
66        let oo2core_3_path = packages_dir.as_ref().join("../bin/x64/oo2core_3_win64.dll");
67        let oo2core_9_path = packages_dir.as_ref().join("../bin/x64/oo2core_9_win64.dll");
68
69        if oo2core_3_path.exists() {
70            let mut o = oodle::OODLE_3.write();
71            if o.is_err() {
72                if let Ok(oodle) = oodle::Oodle::from_path(oo2core_3_path) {
73                    *o = Ok(oodle);
74                }
75            }
76        }
77
78        if oo2core_9_path.exists() {
79            let mut o = oodle::OODLE_9.write();
80            if o.is_err() {
81                if let Ok(oodle) = oodle::Oodle::from_path(oo2core_9_path) {
82                    *o = Ok(oodle);
83                }
84            }
85        }
86
87        let build_new_cache = match Self::validate_cache(version, platform, packages_dir.as_ref()) {
88            Ok(paths) => {
89                packages = paths;
90                false
91            }
92            Err(e) => {
93                warn!("Caches need to be rebuilt: {e}");
94                true
95            }
96        };
97
98        if build_new_cache {
99            info!("Creating new package cache for {}", version.id());
100            let path = packages_dir.as_ref();
101            // Every package in the given directory, including every patch
102            let mut packages_all = vec![];
103            debug_span!("Discover packages in directory").in_scope(|| -> anyhow::Result<()> {
104                for entry in fs::read_dir(path)? {
105                    let entry = entry?;
106                    let path = entry.path();
107                    if path.is_file() && path.to_string_lossy().to_lowercase().ends_with(".pkg") {
108                        packages_all.push(path.to_string_lossy().to_string());
109                    }
110                }
111
112                Ok(())
113            })?;
114
115            packages_all.sort_by_cached_key(|p| {
116                let p = PackagePath::parse_with_defaults(p);
117                (p.id, p.patch)
118            });
119
120            debug_span!("Filter latest packages").in_scope(|| {
121                for p in packages_all {
122                    let parts: Vec<&str> = p.split('_').collect();
123                    if let Some(Ok(pkg_id)) = parts
124                        .get(parts.len() - 2)
125                        .map(|s| u16::from_str_radix(s, 16))
126                    {
127                        packages.insert(pkg_id, p);
128                    } else {
129                        let _span = debug_span!("Open package to find package ID").entered();
130                        // Take the long route and extract the package ID from the header
131                        if let Ok(pkg) = version.open(&p) {
132                            if pkg.language().english_or_none() {
133                                packages.insert(pkg.pkg_id(), p);
134                            }
135                        }
136                    }
137                }
138            });
139        }
140
141        let package_paths: HashMap<u16, PackagePath> = packages
142            .into_iter()
143            .map(|(id, p)| (id, PackagePath::parse_with_defaults(&p)))
144            .collect();
145
146        let first_path = package_paths.values().next().context("No packages found")?;
147
148        let platform = if let Ok(pkg) = version.open(&first_path.path) {
149            pkg.platform()
150        } else {
151            PackagePlatform::from_str(first_path.platform.as_str())?
152        };
153
154        let mut s = Self {
155            package_dir: packages_dir.as_ref().to_path_buf(),
156            platform,
157            package_paths,
158            version,
159            lookup: Default::default(),
160            pkgs: Default::default(),
161        };
162
163        if build_new_cache {
164            s.build_lookup_tables();
165            s.write_package_cache().ok();
166            s.write_lookup_cache().ok();
167        } else if let Some(lookup_cache) = s.read_lookup_cache() {
168            s.lookup = lookup_cache;
169        } else {
170            info!("No valid index cache found, rebuilding");
171            s.build_lookup_tables();
172            s.write_lookup_cache().ok();
173        }
174
175        Ok(s)
176    }
177}
178
179impl PackageManager {
180    pub fn get_all_by_reference(&self, reference: u32) -> Vec<(TagHash, UEntryHeader)> {
181        self.lookup
182            .tag32_entries_by_pkg
183            .par_iter()
184            .map(|(p, e)| {
185                e.iter()
186                    .enumerate()
187                    .filter(|(_, e)| e.reference == reference)
188                    .map(|(i, e)| (TagHash::new(*p, i as _), e.clone()))
189                    .collect::<Vec<(TagHash, UEntryHeader)>>()
190            })
191            .flatten()
192            .collect()
193    }
194
195    pub fn get_all_by_type(&self, etype: u8, esubtype: Option<u8>) -> Vec<(TagHash, UEntryHeader)> {
196        self.lookup
197            .tag32_entries_by_pkg
198            .par_iter()
199            .map(|(p, e)| {
200                e.iter()
201                    .enumerate()
202                    .filter(|(_, e)| {
203                        e.file_type == etype
204                            && esubtype.map(|t| t == e.file_subtype).unwrap_or(true)
205                    })
206                    .map(|(i, e)| (TagHash::new(*p, i as _), e.clone()))
207                    .collect::<Vec<(TagHash, UEntryHeader)>>()
208            })
209            .flatten()
210            .collect()
211    }
212
213    fn get_or_load_pkg(&self, pkg_id: u16) -> anyhow::Result<Arc<dyn Package>> {
214        let v = self.pkgs.read();
215        if let Some(pkg) = v.get(&pkg_id) {
216            Ok(Arc::clone(pkg))
217        } else {
218            drop(v);
219            let package_path = self
220                .package_paths
221                .get(&pkg_id)
222                .with_context(|| format!("Couldn't get a path for package id {pkg_id:04x}"))?;
223
224            let package = self
225                .version
226                .open(&package_path.path)
227                .with_context(|| format!("Failed to open package '{}'", package_path.filename))?;
228
229            self.pkgs.write().insert(pkg_id, Arc::clone(&package));
230            Ok(package)
231        }
232    }
233
234    pub fn read_tag(&self, tag: impl Into<TagHash>) -> anyhow::Result<Vec<u8>> {
235        let tag = tag.into();
236        self.get_or_load_pkg(tag.pkg_id())?
237            .read_entry(tag.entry_index() as _)
238    }
239
240    pub fn read_tag64(&self, hash: impl Into<TagHash64>) -> anyhow::Result<Vec<u8>> {
241        let hash = hash.into();
242        let tag = self
243            .lookup
244            .tag64_entries
245            .get(&hash.0)
246            .context("Hash not found")?
247            .hash32;
248        self.read_tag(tag)
249    }
250
251    pub fn get_entry(&self, tag: impl Into<TagHash>) -> Option<UEntryHeader> {
252        let tag: TagHash = tag.into();
253
254        self.lookup
255            .tag32_entries_by_pkg
256            .get(&tag.pkg_id())?
257            .get(tag.entry_index() as usize)
258            .cloned()
259    }
260
261    pub fn get_named_tag(&self, name: &str, class_hash: u32) -> Option<TagHash> {
262        self.lookup
263            .named_tags
264            .iter()
265            .find(|n| n.name == name && n.class_hash == class_hash)
266            .map(|n| n.hash)
267    }
268
269    pub fn get_named_tags_by_class(&self, class_hash: u32) -> Vec<(String, TagHash)> {
270        self.lookup
271            .named_tags
272            .iter()
273            .filter(|n| n.class_hash == class_hash)
274            .map(|n| (n.name.clone(), n.hash))
275            .collect()
276    }
277
278    /// Find the name of a tag by its hash, if it has one.
279    pub fn get_tag_name(&self, tag: impl Into<TagHash>) -> Option<String> {
280        let tag: TagHash = tag.into();
281        self.lookup
282            .named_tags
283            .iter()
284            .find(|n| n.hash == tag)
285            .map(|n| n.name.clone())
286    }
287
288    pub fn get_tag64_for_tag32(&self, tag: impl Into<TagHash>) -> Option<TagHash64> {
289        let tag: TagHash = tag.into();
290        self.lookup.tag32_to_tag64.get(&tag).copied()
291    }
292
293    /// Read any BinRead type
294    pub fn read_tag_binrw<'a, T: BinRead>(&self, tag: impl Into<TagHash>) -> anyhow::Result<T>
295    where
296        T::Args<'a>: Default + Clone,
297    {
298        let tag = tag.into();
299        let data = self.read_tag(tag)?;
300        let mut cursor = Cursor::new(&data);
301        Ok(cursor.read_type(self.version.endian())?)
302    }
303
304    /// Read any BinRead type
305    pub fn read_tag64_binrw<'a, T: BinRead>(&self, hash: impl Into<TagHash64>) -> anyhow::Result<T>
306    where
307        T::Args<'a>: Default + Clone,
308    {
309        let data = self.read_tag64(hash)?;
310        let mut cursor = Cursor::new(&data);
311        Ok(cursor.read_type(self.version.endian())?)
312    }
313
314    pub fn package_redaction_level(&self, pkg_id: u16) -> Option<Redaction> {
315        self.lookup.redaction_levels.get(&pkg_id).copied()
316    }
317}
318
319#[derive(Debug, Clone)]
320pub struct PackagePath {
321    /// eg. ps3, w64
322    pub platform: String,
323    /// eg. arch_fallen, dungeon_prophecy, europa
324    pub name: String,
325
326    /// 2-letter language code (en, fr, de, etc.)
327    pub language: Option<String>,
328
329    /// eg. 0059, 043c, unp1, unp2
330    pub id: String,
331    pub patch: u8,
332
333    /// Full path to the package
334    pub path: String,
335    pub filename: String,
336}
337
338impl PackagePath {
339    /// Example path: ps3_arch_fallen_0059_0.pkg
340    pub fn parse(path: &str) -> Option<Self> {
341        let path_filename = Path::new(path).file_name()?.to_string_lossy();
342        let parts: Vec<&str> = path_filename.split('_').collect();
343        if parts.len() < 4 {
344            return None;
345        }
346
347        let platform = parts[0].to_string();
348        let mut name = parts[1..parts.len() - 2].join("_");
349        let mut id = parts[parts.len() - 2].to_string();
350        let mut language = None;
351        if id.len() == 2 {
352            // ID is actually language code
353            language = Some(id.clone());
354            name = parts[1..parts.len() - 3].join("_");
355            id = parts[parts.len() - 3].to_string();
356        }
357
358        let patch = parts[parts.len() - 1].split('.').next()?.parse().ok()?;
359
360        Some(Self {
361            platform,
362            name,
363            language,
364            id,
365            patch,
366            path: path.to_string(),
367            filename: path_filename.to_string(),
368        })
369    }
370
371    pub fn parse_with_defaults(path: &str) -> Self {
372        let path_filename = Path::new(path)
373            .file_name()
374            .map_or(path.to_string(), |p| p.to_string_lossy().to_string());
375        Self::parse(path).unwrap_or_else(|| Self {
376            platform: "unknown".to_string(),
377            name: "unknown".to_string(),
378            id: "unknown".to_string(),
379            language: None,
380            patch: 0,
381            path: path.to_string(),
382            filename: path_filename,
383        })
384    }
385}
386
387impl Display for PackagePath {
388    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
389        write!(f, "{}", self.filename)
390    }
391}