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