1use crate::{Archive, Error, FileEntry, Result};
8use std::collections::HashMap;
9use std::path::{Path, PathBuf};
10
11#[derive(Debug)]
52pub struct PatchChain {
53 archives: Vec<ChainEntry>,
55 file_map: HashMap<String, usize>,
57}
58
59#[derive(Debug)]
60struct ChainEntry {
61 archive: Archive,
63 priority: i32,
65 path: PathBuf,
67}
68
69impl PatchChain {
70 pub fn new() -> Self {
72 Self {
73 archives: Vec::new(),
74 file_map: HashMap::new(),
75 }
76 }
77
78 pub fn add_archive<P: AsRef<Path>>(&mut self, path: P, priority: i32) -> Result<()> {
92 let path = path.as_ref();
93 let archive = Archive::open(path)?;
94
95 let entry = ChainEntry {
97 archive,
98 priority,
99 path: path.to_path_buf(),
100 };
101
102 let insert_pos = self
103 .archives
104 .iter()
105 .position(|e| e.priority < priority)
106 .unwrap_or(self.archives.len());
107
108 self.archives.insert(insert_pos, entry);
109
110 self.rebuild_file_map()?;
112
113 Ok(())
114 }
115
116 pub fn remove_archive<P: AsRef<Path>>(&mut self, path: P) -> Result<bool> {
118 let path = path.as_ref();
119
120 if let Some(pos) = self.archives.iter().position(|e| e.path == path) {
121 self.archives.remove(pos);
122 self.rebuild_file_map()?;
123 Ok(true)
124 } else {
125 Ok(false)
126 }
127 }
128
129 pub fn clear(&mut self) {
131 self.archives.clear();
132 self.file_map.clear();
133 }
134
135 pub fn archive_count(&self) -> usize {
137 self.archives.len()
138 }
139
140 pub fn read_file(&mut self, filename: &str) -> Result<Vec<u8>> {
148 let lookup_key = crate::path::normalize_mpq_path(filename).to_uppercase();
151
152 if let Some(&archive_idx) = self.file_map.get(&lookup_key) {
153 let file_info = self.archives[archive_idx]
155 .archive
156 .find_file(filename)?
157 .ok_or_else(|| Error::FileNotFound(filename.to_string()))?;
158
159 if file_info.is_patch_file() {
160 self.read_patched_file(filename, archive_idx)
162 } else {
163 self.archives[archive_idx].archive.read_file(filename)
165 }
166 } else {
167 Err(Error::FileNotFound(filename.to_string()))
168 }
169 }
170
171 fn read_patched_file(&mut self, filename: &str, _patch_idx: usize) -> Result<Vec<u8>> {
178 use crate::patch::{PatchFile, apply_patch};
179
180 let mut base_data: Option<Vec<u8>> = None;
182 let mut patches = Vec::new();
183
184 for (idx, entry) in self.archives.iter_mut().enumerate() {
186 if let Ok(Some(file_info)) = entry.archive.find_file(filename) {
187 if file_info.is_patch_file() {
188 match entry.archive.read_patch_file_raw(filename) {
190 Ok(patch_data) => {
191 match PatchFile::parse(&patch_data) {
193 Ok(patch) => patches.push((idx, patch)),
194 Err(e) => {
195 log::warn!(
196 "Failed to parse patch file '{}' in archive {} (priority {}): {}",
197 filename,
198 entry.path.display(),
199 entry.priority,
200 e
201 );
202 }
203 }
204 }
205 Err(e) => {
206 log::warn!(
207 "Failed to read patch file '{}' in archive {} (priority {}): {}",
208 filename,
209 entry.path.display(),
210 entry.priority,
211 e
212 );
213 }
214 }
215 } else if base_data.is_none() {
216 match entry.archive.read_file(filename) {
218 Ok(data) => {
219 log::debug!(
220 "Found base file '{}' in archive {} (priority {})",
221 filename,
222 entry.path.display(),
223 entry.priority
224 );
225 base_data = Some(data);
226 }
227 Err(e) => {
228 log::warn!(
229 "Failed to read base file '{}' in archive {} (priority {}): {}",
230 filename,
231 entry.path.display(),
232 entry.priority,
233 e
234 );
235 }
236 }
237 }
238 }
239 }
240
241 let mut current_data = base_data.ok_or_else(|| {
243 Error::FileNotFound(format!(
244 "No base file found for patch file '{filename}' in patch chain"
245 ))
246 })?;
247
248 patches.reverse();
251 for (idx, patch) in patches {
252 log::debug!(
253 "Applying patch '{}' from archive {} (priority {})",
254 filename,
255 self.archives[idx].path.display(),
256 self.archives[idx].priority
257 );
258
259 current_data = apply_patch(&patch, ¤t_data)?;
260 }
261
262 Ok(current_data)
263 }
264
265 pub fn contains_file(&self, filename: &str) -> bool {
267 let lookup_key = crate::path::normalize_mpq_path(filename).to_uppercase();
268 self.file_map.contains_key(&lookup_key)
269 }
270
271 pub fn find_file_archive(&self, filename: &str) -> Option<&Path> {
275 let lookup_key = crate::path::normalize_mpq_path(filename).to_uppercase();
276 self.file_map
277 .get(&lookup_key)
278 .map(|&idx| self.archives[idx].path.as_path())
279 }
280
281 pub fn list(&mut self) -> Result<Vec<FileEntry>> {
286 let mut seen = HashMap::new();
287 let mut result = Vec::new();
288
289 for (idx, entry) in self.archives.iter_mut().enumerate() {
291 match entry.archive.list() {
292 Ok(files) => {
293 for file in files {
294 if !seen.contains_key(&file.name) {
296 seen.insert(file.name.clone(), idx);
297 result.push(file);
298 }
299 }
300 }
301 Err(_) => {
302 if let Ok(files) = entry.archive.list_all() {
304 for file in files {
305 if !seen.contains_key(&file.name) {
306 seen.insert(file.name.clone(), idx);
307 result.push(file);
308 }
309 }
310 }
311 }
312 }
313 }
314
315 result.sort_by(|a, b| a.name.cmp(&b.name));
317
318 Ok(result)
319 }
320
321 pub fn get_chain_info(&mut self) -> Vec<ChainInfo> {
323 self.archives
324 .iter_mut()
325 .filter_map(|entry| {
326 entry.archive.get_info().ok().map(|info| ChainInfo {
327 path: entry.path.clone(),
328 priority: entry.priority,
329 file_count: info.file_count,
330 archive_size: info.file_size,
331 format_version: info.format_version,
332 })
333 })
334 .collect()
335 }
336
337 fn rebuild_file_map(&mut self) -> Result<()> {
339 self.file_map.clear();
340
341 for (idx, entry) in self.archives.iter_mut().enumerate() {
343 let files = match entry.archive.list() {
345 Ok(files) => files,
346 Err(_) => {
347 match entry.archive.list_all() {
349 Ok(files) => files,
350 Err(_) => continue, }
352 }
353 };
354
355 for file in files {
358 let normalized_key = crate::path::normalize_mpq_path(&file.name).to_uppercase();
359 self.file_map.entry(normalized_key).or_insert(idx);
360 }
361 }
362
363 Ok(())
364 }
365
366 pub fn extract_files(&mut self, filenames: &[&str]) -> Vec<(String, Result<Vec<u8>>)> {
371 filenames
372 .iter()
373 .map(|&filename| {
374 let result = self.read_file(filename);
375 (filename.to_string(), result)
376 })
377 .collect()
378 }
379
380 pub fn get_archive<P: AsRef<Path>>(&self, path: P) -> Option<&Archive> {
382 let path = path.as_ref();
383 self.archives
384 .iter()
385 .find(|e| e.path == path)
386 .map(|e| &e.archive)
387 }
388
389 pub fn get_priority<P: AsRef<Path>>(&self, path: P) -> Option<i32> {
391 let path = path.as_ref();
392 self.archives
393 .iter()
394 .find(|e| e.path == path)
395 .map(|e| e.priority)
396 }
397
398 pub fn from_archives_parallel<P: AsRef<Path> + Sync>(archives: Vec<(P, i32)>) -> Result<Self> {
422 use rayon::prelude::*;
423
424 let loaded_archives: Result<Vec<_>> = archives
426 .par_iter()
427 .map(|(path, priority)| {
428 let path_ref = path.as_ref();
429 Archive::open(path_ref).map(|archive| ChainEntry {
430 archive,
431 priority: *priority,
432 path: path_ref.to_path_buf(),
433 })
434 })
435 .collect();
436
437 let mut loaded_archives = loaded_archives?;
438
439 loaded_archives.sort_by(|a, b| b.priority.cmp(&a.priority));
441
442 let mut chain = Self {
443 archives: loaded_archives,
444 file_map: HashMap::new(),
445 };
446
447 chain.rebuild_file_map()?;
449
450 Ok(chain)
451 }
452
453 pub fn add_archives_parallel<P: AsRef<Path> + Sync>(
458 &mut self,
459 archives: Vec<(P, i32)>,
460 ) -> Result<()> {
461 use rayon::prelude::*;
462
463 let new_archives: Result<Vec<_>> = archives
465 .par_iter()
466 .map(|(path, priority)| {
467 let path_ref = path.as_ref();
468 Archive::open(path_ref).map(|archive| ChainEntry {
469 archive,
470 priority: *priority,
471 path: path_ref.to_path_buf(),
472 })
473 })
474 .collect();
475
476 let new_archives = new_archives?;
477
478 for entry in new_archives {
480 let insert_pos = self
481 .archives
482 .iter()
483 .position(|e| e.priority < entry.priority)
484 .unwrap_or(self.archives.len());
485
486 self.archives.insert(insert_pos, entry);
487 }
488
489 self.rebuild_file_map()?;
491
492 Ok(())
493 }
494
495 pub fn set_priority<P: AsRef<Path>>(&mut self, path: P, new_priority: i32) -> Result<()> {
497 let path = path.as_ref();
498
499 let archive_idx = self
501 .archives
502 .iter()
503 .position(|e| e.path == path)
504 .ok_or_else(|| Error::InvalidFormat("Archive not found in chain".to_string()))?;
505
506 let mut entry = self.archives.remove(archive_idx);
507 entry.priority = new_priority;
508
509 let insert_pos = self
511 .archives
512 .iter()
513 .position(|e| e.priority < new_priority)
514 .unwrap_or(self.archives.len());
515
516 self.archives.insert(insert_pos, entry);
517
518 self.rebuild_file_map()?;
520
521 Ok(())
522 }
523}
524
525impl Default for PatchChain {
526 fn default() -> Self {
527 Self::new()
528 }
529}
530
531#[derive(Debug, Clone)]
533pub struct ChainInfo {
534 pub path: PathBuf,
536 pub priority: i32,
538 pub file_count: usize,
540 pub archive_size: u64,
542 pub format_version: crate::FormatVersion,
544}
545
546#[cfg(test)]
547mod tests;
548
549#[cfg(test)]
550mod integration_tests {
551 use super::*;
552 use crate::{ArchiveBuilder, ListfileOption};
553 use tempfile::TempDir;
554
555 fn create_test_archive(dir: &Path, name: &str, files: &[(&str, &[u8])]) -> PathBuf {
556 let path = dir.join(name);
557 let mut builder = ArchiveBuilder::new().listfile_option(ListfileOption::Generate);
558
559 for (filename, data) in files {
560 builder = builder.add_file_data(data.to_vec(), filename);
561 }
562
563 builder.build(&path).unwrap();
564 path
565 }
566
567 #[test]
568 fn test_patch_chain_priority() {
569 let temp = TempDir::new().unwrap();
570
571 let base_files: Vec<(&str, &[u8])> = vec![
573 ("file1.txt", b"base file1"),
574 ("file2.txt", b"base file2"),
575 ("file3.txt", b"base file3"),
576 ];
577 let base_path = create_test_archive(temp.path(), "base.mpq", &base_files);
578
579 let patch_files: Vec<(&str, &[u8])> =
581 vec![("file2.txt", b"patched file2"), ("file4.txt", b"new file4")];
582 let patch_path = create_test_archive(temp.path(), "patch.mpq", &patch_files);
583
584 let mut chain = PatchChain::new();
586 chain.add_archive(&base_path, 0).unwrap();
587 chain.add_archive(&patch_path, 100).unwrap();
588
589 assert_eq!(chain.read_file("file1.txt").unwrap(), b"base file1");
591 assert_eq!(chain.read_file("file2.txt").unwrap(), b"patched file2"); assert_eq!(chain.read_file("file3.txt").unwrap(), b"base file3");
593 assert_eq!(chain.read_file("file4.txt").unwrap(), b"new file4");
594 }
595
596 #[test]
597 fn test_patch_chain_listing() {
598 let temp = TempDir::new().unwrap();
599
600 let base_files: Vec<(&str, &[u8])> = vec![("file1.txt", b"data1"), ("file2.txt", b"data2")];
602 let patch_files: Vec<(&str, &[u8])> =
603 vec![("file2.txt", b"patch2"), ("file3.txt", b"data3")];
604
605 let base_path = create_test_archive(temp.path(), "base.mpq", &base_files);
606 let patch_path = create_test_archive(temp.path(), "patch.mpq", &patch_files);
607
608 let mut chain = PatchChain::new();
610 chain.add_archive(&base_path, 0).unwrap();
611 chain.add_archive(&patch_path, 100).unwrap();
612
613 let files = chain.list().unwrap();
615
616 let files: Vec<_> = files
618 .into_iter()
619 .filter(|f| f.name != "(listfile)")
620 .collect();
621
622 assert_eq!(files.len(), 3);
623
624 let names: Vec<_> = files.iter().map(|f| f.name.as_str()).collect();
625 assert!(names.contains(&"file1.txt"));
626 assert!(names.contains(&"file2.txt"));
627 assert!(names.contains(&"file3.txt"));
628 }
629
630 #[test]
631 fn test_find_file_archive() {
632 let temp = TempDir::new().unwrap();
633
634 let base_files: Vec<(&str, &[u8])> = vec![("file1.txt", b"data")];
635 let patch_files: Vec<(&str, &[u8])> = vec![("file2.txt", b"data")];
636
637 let base_path = create_test_archive(temp.path(), "base.mpq", &base_files);
638 let patch_path = create_test_archive(temp.path(), "patch.mpq", &patch_files);
639
640 let mut chain = PatchChain::new();
641 chain.add_archive(&base_path, 0).unwrap();
642 chain.add_archive(&patch_path, 100).unwrap();
643
644 assert_eq!(
645 chain.find_file_archive("file1.txt"),
646 Some(base_path.as_path())
647 );
648 assert_eq!(
649 chain.find_file_archive("file2.txt"),
650 Some(patch_path.as_path())
651 );
652 assert_eq!(chain.find_file_archive("nonexistent.txt"), None);
653 }
654
655 #[test]
656 fn test_remove_archive() {
657 let temp = TempDir::new().unwrap();
658
659 let files: Vec<(&str, &[u8])> = vec![("file.txt", b"data")];
660 let path = create_test_archive(temp.path(), "test.mpq", &files);
661
662 let mut chain = PatchChain::new();
663 chain.add_archive(&path, 0).unwrap();
664
665 assert!(chain.contains_file("file.txt"));
666 assert!(chain.remove_archive(&path).unwrap());
667 assert!(!chain.contains_file("file.txt"));
668 assert!(!chain.remove_archive(&path).unwrap()); }
670
671 #[test]
672 fn test_priority_reordering() {
673 let temp = TempDir::new().unwrap();
674
675 let files: Vec<(&str, &[u8])> = vec![("file.txt", b"data")];
676 let path1 = create_test_archive(temp.path(), "test1.mpq", &files);
677 let path2 = create_test_archive(temp.path(), "test2.mpq", &files);
678
679 let mut chain = PatchChain::new();
680 chain.add_archive(&path1, 100).unwrap();
681 chain.add_archive(&path2, 50).unwrap();
682
683 assert_eq!(chain.archives[0].priority, 100);
685
686 chain.set_priority(&path2, 150).unwrap();
688
689 assert_eq!(chain.archives[0].priority, 150);
691 }
692
693 #[test]
694 fn test_parallel_patch_chain_loading() {
695 let temp = TempDir::new().unwrap();
696
697 let mut archive_paths = Vec::new();
699 for i in 0..5 {
700 let common_content = format!("Common content v{i}");
701 let unique_name = format!("unique_{i}.txt");
702 let unique_content = format!("Unique to archive {i}");
703 let files: Vec<(&str, &[u8])> = vec![
704 ("common.txt", common_content.as_bytes()),
705 (&unique_name, unique_content.as_bytes()),
706 ];
707 let path = create_test_archive(temp.path(), &format!("archive_{i}.mpq"), &files);
708 archive_paths.push((path, i * 100)); }
710
711 let start = std::time::Instant::now();
713 let mut chain_seq = PatchChain::new();
714 for (path, priority) in &archive_paths {
715 chain_seq.add_archive(path, *priority).unwrap();
716 }
717 let seq_duration = start.elapsed();
718
719 let start = std::time::Instant::now();
721 let mut chain_par = PatchChain::from_archives_parallel(archive_paths.clone()).unwrap();
722 let par_duration = start.elapsed();
723
724 assert_eq!(
726 chain_seq.list().unwrap().len(),
727 chain_par.list().unwrap().len()
728 );
729
730 let common_content = chain_par.read_file("common.txt").unwrap();
732 assert_eq!(common_content, b"Common content v4"); for i in 0..5 {
736 let unique_file = format!("unique_{i}.txt");
737 let content = chain_par.read_file(&unique_file).unwrap();
738 assert_eq!(content, format!("Unique to archive {i}").as_bytes());
739 }
740
741 println!("Sequential loading: {seq_duration:?}");
743 println!("Parallel loading: {par_duration:?}");
744 }
745
746 #[test]
747 fn test_add_archives_parallel() {
748 let temp = TempDir::new().unwrap();
749
750 let base_files: Vec<(&str, &[u8])> = vec![("base.txt", b"base content")];
752 let base_path = create_test_archive(temp.path(), "base.mpq", &base_files);
753
754 let mut chain = PatchChain::new();
756 chain.add_archive(&base_path, 0).unwrap();
757
758 let mut patch_archives = Vec::new();
760 for i in 1..=3 {
761 let patch_name = format!("patch_{i}.txt");
762 let patch_content = format!("Patch {i} content");
763 let common_content = format!("Common from patch {i}");
764 let files: Vec<(&str, &[u8])> = vec![
765 (&patch_name, patch_content.as_bytes()),
766 ("common.txt", common_content.as_bytes()),
767 ];
768 let path = create_test_archive(temp.path(), &format!("patch_{i}.mpq"), &files);
769 patch_archives.push((path, i * 100));
770 }
771
772 chain.add_archives_parallel(patch_archives).unwrap();
774
775 assert_eq!(chain.read_file("base.txt").unwrap(), b"base content");
777 assert_eq!(chain.read_file("patch_1.txt").unwrap(), b"Patch 1 content");
778 assert_eq!(chain.read_file("patch_2.txt").unwrap(), b"Patch 2 content");
779 assert_eq!(chain.read_file("patch_3.txt").unwrap(), b"Patch 3 content");
780
781 assert_eq!(
783 chain.read_file("common.txt").unwrap(),
784 b"Common from patch 3"
785 );
786
787 let info = chain.get_chain_info();
789 assert_eq!(info.len(), 4); }
791
792 #[test]
793 fn test_parallel_loading_with_invalid_archive() {
794 let temp = TempDir::new().unwrap();
795
796 let mut archives = Vec::new();
798 for i in 0..2 {
799 let file_name = format!("file_{i}.txt");
800 let files: Vec<(&str, &[u8])> = vec![(&file_name, b"content")];
801 let path = create_test_archive(temp.path(), &format!("valid_{i}.mpq"), &files);
802 archives.push((path, i * 100));
803 }
804
805 archives.push((temp.path().join("nonexistent.mpq"), 200));
807
808 let result = PatchChain::from_archives_parallel(archives);
810 assert!(result.is_err());
811 }
812}