use crate::{Archive, Error, FileEntry, Result};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
#[derive(Debug)]
pub struct PatchChain {
archives: Vec<ChainEntry>,
file_map: HashMap<String, usize>,
}
#[derive(Debug)]
struct ChainEntry {
archive: Archive,
priority: i32,
path: PathBuf,
}
impl PatchChain {
pub fn new() -> Self {
Self {
archives: Vec::new(),
file_map: HashMap::new(),
}
}
pub fn add_archive<P: AsRef<Path>>(&mut self, path: P, priority: i32) -> Result<()> {
let path = path.as_ref();
let archive = Archive::open(path)?;
let entry = ChainEntry {
archive,
priority,
path: path.to_path_buf(),
};
let insert_pos = self
.archives
.iter()
.position(|e| e.priority < priority)
.unwrap_or(self.archives.len());
self.archives.insert(insert_pos, entry);
self.rebuild_file_map()?;
Ok(())
}
pub fn remove_archive<P: AsRef<Path>>(&mut self, path: P) -> Result<bool> {
let path = path.as_ref();
if let Some(pos) = self.archives.iter().position(|e| e.path == path) {
self.archives.remove(pos);
self.rebuild_file_map()?;
Ok(true)
} else {
Ok(false)
}
}
pub fn clear(&mut self) {
self.archives.clear();
self.file_map.clear();
}
pub fn archive_count(&self) -> usize {
self.archives.len()
}
pub fn read_file(&mut self, filename: &str) -> Result<Vec<u8>> {
let lookup_key = crate::path::normalize_mpq_path(filename).to_uppercase();
if let Some(&archive_idx) = self.file_map.get(&lookup_key) {
let file_info = self.archives[archive_idx]
.archive
.find_file(filename)?
.ok_or_else(|| Error::FileNotFound(filename.to_string()))?;
if file_info.is_patch_file() {
self.read_patched_file(filename, archive_idx)
} else {
self.archives[archive_idx].archive.read_file(filename)
}
} else {
Err(Error::FileNotFound(filename.to_string()))
}
}
fn read_patched_file(&mut self, filename: &str, _patch_idx: usize) -> Result<Vec<u8>> {
use crate::patch::{PatchFile, apply_patch};
let mut base_data: Option<Vec<u8>> = None;
let mut patches = Vec::new();
for (idx, entry) in self.archives.iter_mut().enumerate() {
if let Ok(Some(file_info)) = entry.archive.find_file(filename) {
if file_info.is_patch_file() {
match entry.archive.read_patch_file_raw(filename) {
Ok(patch_data) => {
match PatchFile::parse(&patch_data) {
Ok(patch) => patches.push((idx, patch)),
Err(e) => {
log::warn!(
"Failed to parse patch file '{}' in archive {} (priority {}): {}",
filename,
entry.path.display(),
entry.priority,
e
);
}
}
}
Err(e) => {
log::warn!(
"Failed to read patch file '{}' in archive {} (priority {}): {}",
filename,
entry.path.display(),
entry.priority,
e
);
}
}
} else if base_data.is_none() {
match entry.archive.read_file(filename) {
Ok(data) => {
log::debug!(
"Found base file '{}' in archive {} (priority {})",
filename,
entry.path.display(),
entry.priority
);
base_data = Some(data);
}
Err(e) => {
log::warn!(
"Failed to read base file '{}' in archive {} (priority {}): {}",
filename,
entry.path.display(),
entry.priority,
e
);
}
}
}
}
}
let mut current_data = base_data.ok_or_else(|| {
Error::FileNotFound(format!(
"No base file found for patch file '{filename}' in patch chain"
))
})?;
patches.reverse();
for (idx, patch) in patches {
log::debug!(
"Applying patch '{}' from archive {} (priority {})",
filename,
self.archives[idx].path.display(),
self.archives[idx].priority
);
current_data = apply_patch(&patch, ¤t_data)?;
}
Ok(current_data)
}
pub fn contains_file(&self, filename: &str) -> bool {
let lookup_key = crate::path::normalize_mpq_path(filename).to_uppercase();
self.file_map.contains_key(&lookup_key)
}
pub fn find_file_archive(&self, filename: &str) -> Option<&Path> {
let lookup_key = crate::path::normalize_mpq_path(filename).to_uppercase();
self.file_map
.get(&lookup_key)
.map(|&idx| self.archives[idx].path.as_path())
}
pub fn list(&mut self) -> Result<Vec<FileEntry>> {
let mut seen = HashMap::new();
let mut result = Vec::new();
for (idx, entry) in self.archives.iter_mut().enumerate() {
match entry.archive.list() {
Ok(files) => {
for file in files {
if !seen.contains_key(&file.name) {
seen.insert(file.name.clone(), idx);
result.push(file);
}
}
}
Err(_) => {
if let Ok(files) = entry.archive.list_all() {
for file in files {
if !seen.contains_key(&file.name) {
seen.insert(file.name.clone(), idx);
result.push(file);
}
}
}
}
}
}
result.sort_by(|a, b| a.name.cmp(&b.name));
Ok(result)
}
pub fn get_chain_info(&mut self) -> Vec<ChainInfo> {
self.archives
.iter_mut()
.filter_map(|entry| {
entry.archive.get_info().ok().map(|info| ChainInfo {
path: entry.path.clone(),
priority: entry.priority,
file_count: info.file_count,
archive_size: info.file_size,
format_version: info.format_version,
})
})
.collect()
}
fn rebuild_file_map(&mut self) -> Result<()> {
self.file_map.clear();
for (idx, entry) in self.archives.iter_mut().enumerate() {
let files = match entry.archive.list() {
Ok(files) => files,
Err(_) => {
match entry.archive.list_all() {
Ok(files) => files,
Err(_) => continue, }
}
};
for file in files {
let normalized_key = crate::path::normalize_mpq_path(&file.name).to_uppercase();
self.file_map.entry(normalized_key).or_insert(idx);
}
}
Ok(())
}
pub fn extract_files(&mut self, filenames: &[&str]) -> Vec<(String, Result<Vec<u8>>)> {
filenames
.iter()
.map(|&filename| {
let result = self.read_file(filename);
(filename.to_string(), result)
})
.collect()
}
pub fn get_archive<P: AsRef<Path>>(&self, path: P) -> Option<&Archive> {
let path = path.as_ref();
self.archives
.iter()
.find(|e| e.path == path)
.map(|e| &e.archive)
}
pub fn get_priority<P: AsRef<Path>>(&self, path: P) -> Option<i32> {
let path = path.as_ref();
self.archives
.iter()
.find(|e| e.path == path)
.map(|e| e.priority)
}
pub fn from_archives_parallel<P: AsRef<Path> + Sync>(archives: Vec<(P, i32)>) -> Result<Self> {
use rayon::prelude::*;
let loaded_archives: Result<Vec<_>> = archives
.par_iter()
.map(|(path, priority)| {
let path_ref = path.as_ref();
Archive::open(path_ref).map(|archive| ChainEntry {
archive,
priority: *priority,
path: path_ref.to_path_buf(),
})
})
.collect();
let mut loaded_archives = loaded_archives?;
loaded_archives.sort_by(|a, b| b.priority.cmp(&a.priority));
let mut chain = Self {
archives: loaded_archives,
file_map: HashMap::new(),
};
chain.rebuild_file_map()?;
Ok(chain)
}
pub fn add_archives_parallel<P: AsRef<Path> + Sync>(
&mut self,
archives: Vec<(P, i32)>,
) -> Result<()> {
use rayon::prelude::*;
let new_archives: Result<Vec<_>> = archives
.par_iter()
.map(|(path, priority)| {
let path_ref = path.as_ref();
Archive::open(path_ref).map(|archive| ChainEntry {
archive,
priority: *priority,
path: path_ref.to_path_buf(),
})
})
.collect();
let new_archives = new_archives?;
for entry in new_archives {
let insert_pos = self
.archives
.iter()
.position(|e| e.priority < entry.priority)
.unwrap_or(self.archives.len());
self.archives.insert(insert_pos, entry);
}
self.rebuild_file_map()?;
Ok(())
}
pub fn set_priority<P: AsRef<Path>>(&mut self, path: P, new_priority: i32) -> Result<()> {
let path = path.as_ref();
let archive_idx = self
.archives
.iter()
.position(|e| e.path == path)
.ok_or_else(|| Error::InvalidFormat("Archive not found in chain".to_string()))?;
let mut entry = self.archives.remove(archive_idx);
entry.priority = new_priority;
let insert_pos = self
.archives
.iter()
.position(|e| e.priority < new_priority)
.unwrap_or(self.archives.len());
self.archives.insert(insert_pos, entry);
self.rebuild_file_map()?;
Ok(())
}
}
impl Default for PatchChain {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct ChainInfo {
pub path: PathBuf,
pub priority: i32,
pub file_count: usize,
pub archive_size: u64,
pub format_version: crate::FormatVersion,
}
#[cfg(test)]
mod tests;
#[cfg(test)]
mod integration_tests {
use super::*;
use crate::{ArchiveBuilder, ListfileOption};
use tempfile::TempDir;
fn create_test_archive(dir: &Path, name: &str, files: &[(&str, &[u8])]) -> PathBuf {
let path = dir.join(name);
let mut builder = ArchiveBuilder::new().listfile_option(ListfileOption::Generate);
for (filename, data) in files {
builder = builder.add_file_data(data.to_vec(), filename);
}
builder.build(&path).unwrap();
path
}
#[test]
fn test_patch_chain_priority() {
let temp = TempDir::new().unwrap();
let base_files: Vec<(&str, &[u8])> = vec![
("file1.txt", b"base file1"),
("file2.txt", b"base file2"),
("file3.txt", b"base file3"),
];
let base_path = create_test_archive(temp.path(), "base.mpq", &base_files);
let patch_files: Vec<(&str, &[u8])> =
vec![("file2.txt", b"patched file2"), ("file4.txt", b"new file4")];
let patch_path = create_test_archive(temp.path(), "patch.mpq", &patch_files);
let mut chain = PatchChain::new();
chain.add_archive(&base_path, 0).unwrap();
chain.add_archive(&patch_path, 100).unwrap();
assert_eq!(chain.read_file("file1.txt").unwrap(), b"base file1");
assert_eq!(chain.read_file("file2.txt").unwrap(), b"patched file2"); assert_eq!(chain.read_file("file3.txt").unwrap(), b"base file3");
assert_eq!(chain.read_file("file4.txt").unwrap(), b"new file4");
}
#[test]
fn test_patch_chain_listing() {
let temp = TempDir::new().unwrap();
let base_files: Vec<(&str, &[u8])> = vec![("file1.txt", b"data1"), ("file2.txt", b"data2")];
let patch_files: Vec<(&str, &[u8])> =
vec![("file2.txt", b"patch2"), ("file3.txt", b"data3")];
let base_path = create_test_archive(temp.path(), "base.mpq", &base_files);
let patch_path = create_test_archive(temp.path(), "patch.mpq", &patch_files);
let mut chain = PatchChain::new();
chain.add_archive(&base_path, 0).unwrap();
chain.add_archive(&patch_path, 100).unwrap();
let files = chain.list().unwrap();
let files: Vec<_> = files
.into_iter()
.filter(|f| f.name != "(listfile)")
.collect();
assert_eq!(files.len(), 3);
let names: Vec<_> = files.iter().map(|f| f.name.as_str()).collect();
assert!(names.contains(&"file1.txt"));
assert!(names.contains(&"file2.txt"));
assert!(names.contains(&"file3.txt"));
}
#[test]
fn test_find_file_archive() {
let temp = TempDir::new().unwrap();
let base_files: Vec<(&str, &[u8])> = vec![("file1.txt", b"data")];
let patch_files: Vec<(&str, &[u8])> = vec![("file2.txt", b"data")];
let base_path = create_test_archive(temp.path(), "base.mpq", &base_files);
let patch_path = create_test_archive(temp.path(), "patch.mpq", &patch_files);
let mut chain = PatchChain::new();
chain.add_archive(&base_path, 0).unwrap();
chain.add_archive(&patch_path, 100).unwrap();
assert_eq!(
chain.find_file_archive("file1.txt"),
Some(base_path.as_path())
);
assert_eq!(
chain.find_file_archive("file2.txt"),
Some(patch_path.as_path())
);
assert_eq!(chain.find_file_archive("nonexistent.txt"), None);
}
#[test]
fn test_remove_archive() {
let temp = TempDir::new().unwrap();
let files: Vec<(&str, &[u8])> = vec![("file.txt", b"data")];
let path = create_test_archive(temp.path(), "test.mpq", &files);
let mut chain = PatchChain::new();
chain.add_archive(&path, 0).unwrap();
assert!(chain.contains_file("file.txt"));
assert!(chain.remove_archive(&path).unwrap());
assert!(!chain.contains_file("file.txt"));
assert!(!chain.remove_archive(&path).unwrap()); }
#[test]
fn test_priority_reordering() {
let temp = TempDir::new().unwrap();
let files: Vec<(&str, &[u8])> = vec![("file.txt", b"data")];
let path1 = create_test_archive(temp.path(), "test1.mpq", &files);
let path2 = create_test_archive(temp.path(), "test2.mpq", &files);
let mut chain = PatchChain::new();
chain.add_archive(&path1, 100).unwrap();
chain.add_archive(&path2, 50).unwrap();
assert_eq!(chain.archives[0].priority, 100);
chain.set_priority(&path2, 150).unwrap();
assert_eq!(chain.archives[0].priority, 150);
}
#[test]
fn test_parallel_patch_chain_loading() {
let temp = TempDir::new().unwrap();
let mut archive_paths = Vec::new();
for i in 0..5 {
let common_content = format!("Common content v{i}");
let unique_name = format!("unique_{i}.txt");
let unique_content = format!("Unique to archive {i}");
let files: Vec<(&str, &[u8])> = vec![
("common.txt", common_content.as_bytes()),
(&unique_name, unique_content.as_bytes()),
];
let path = create_test_archive(temp.path(), &format!("archive_{i}.mpq"), &files);
archive_paths.push((path, i * 100)); }
let start = std::time::Instant::now();
let mut chain_seq = PatchChain::new();
for (path, priority) in &archive_paths {
chain_seq.add_archive(path, *priority).unwrap();
}
let seq_duration = start.elapsed();
let start = std::time::Instant::now();
let mut chain_par = PatchChain::from_archives_parallel(archive_paths.clone()).unwrap();
let par_duration = start.elapsed();
assert_eq!(
chain_seq.list().unwrap().len(),
chain_par.list().unwrap().len()
);
let common_content = chain_par.read_file("common.txt").unwrap();
assert_eq!(common_content, b"Common content v4");
for i in 0..5 {
let unique_file = format!("unique_{i}.txt");
let content = chain_par.read_file(&unique_file).unwrap();
assert_eq!(content, format!("Unique to archive {i}").as_bytes());
}
println!("Sequential loading: {seq_duration:?}");
println!("Parallel loading: {par_duration:?}");
}
#[test]
fn test_add_archives_parallel() {
let temp = TempDir::new().unwrap();
let base_files: Vec<(&str, &[u8])> = vec![("base.txt", b"base content")];
let base_path = create_test_archive(temp.path(), "base.mpq", &base_files);
let mut chain = PatchChain::new();
chain.add_archive(&base_path, 0).unwrap();
let mut patch_archives = Vec::new();
for i in 1..=3 {
let patch_name = format!("patch_{i}.txt");
let patch_content = format!("Patch {i} content");
let common_content = format!("Common from patch {i}");
let files: Vec<(&str, &[u8])> = vec![
(&patch_name, patch_content.as_bytes()),
("common.txt", common_content.as_bytes()),
];
let path = create_test_archive(temp.path(), &format!("patch_{i}.mpq"), &files);
patch_archives.push((path, i * 100));
}
chain.add_archives_parallel(patch_archives).unwrap();
assert_eq!(chain.read_file("base.txt").unwrap(), b"base content");
assert_eq!(chain.read_file("patch_1.txt").unwrap(), b"Patch 1 content");
assert_eq!(chain.read_file("patch_2.txt").unwrap(), b"Patch 2 content");
assert_eq!(chain.read_file("patch_3.txt").unwrap(), b"Patch 3 content");
assert_eq!(
chain.read_file("common.txt").unwrap(),
b"Common from patch 3"
);
let info = chain.get_chain_info();
assert_eq!(info.len(), 4); }
#[test]
fn test_parallel_loading_with_invalid_archive() {
let temp = TempDir::new().unwrap();
let mut archives = Vec::new();
for i in 0..2 {
let file_name = format!("file_{i}.txt");
let files: Vec<(&str, &[u8])> = vec![(&file_name, b"content")];
let path = create_test_archive(temp.path(), &format!("valid_{i}.mpq"), &files);
archives.push((path, i * 100));
}
archives.push((temp.path().join("nonexistent.mpq"), 200));
let result = PatchChain::from_archives_parallel(archives);
assert!(result.is_err());
}
}