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 pub lookup: TagLookupIndex,
52
53 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 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 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 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 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 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 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 pub platform: String,
323 pub name: String,
325
326 pub language: Option<String>,
328
329 pub id: String,
331 pub patch: u8,
332
333 pub path: String,
335 pub filename: String,
336}
337
338impl PackagePath {
339 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 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}