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