use std::collections::HashSet;
use std::fs::{self, File};
use std::io::{self, BufWriter};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use rayon::prelude::*;
use crate::archive::{open_archive, sanitize_relative_path};
use crate::codec::EncodingKind;
use crate::error::XzipError;
use crate::filter::PathFilter;
use crate::verbosity::Verbosity;
const WRITE_BUFFER_SIZE: usize = 128 * 1024;
const PROGRESS_INTERVAL: usize = 250;
const PARALLEL_MIN_FILES: usize = 256;
const PARALLEL_MIN_BYTES: u64 = 4 * 1024 * 1024;
#[derive(Debug, Clone)]
struct ExtractPlan {
index: usize,
outpath: PathBuf,
display_name: String,
is_dir: bool,
uncompressed_size: u64,
compressed_size: u64,
}
fn ensure_dir(path: &Path, created: &mut HashSet<PathBuf>) -> Result<bool, XzipError> {
let path = path.to_path_buf();
if created.contains(&path) {
return Ok(false);
}
fs::create_dir_all(&path)?;
created.insert(path);
Ok(true)
}
fn available_cpus() -> usize {
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
}
fn total_uncompressed_bytes(plans: &[ExtractPlan]) -> u64 {
plans
.iter()
.filter(|plan| !plan.is_dir)
.map(|plan| plan.uncompressed_size)
.sum()
}
fn resolve_jobs(jobs: usize, file_count: usize, total_bytes: u64) -> usize {
match jobs {
0 => {
if file_count >= PARALLEL_MIN_FILES && total_bytes >= PARALLEL_MIN_BYTES {
available_cpus()
} else {
1
}
}
1 => 1,
n => n,
}
}
fn scan_entries(
archive: &mut zip::read::ZipArchive<File>,
output_dir: &Path,
encoding: EncodingKind,
filter: &PathFilter,
) -> Result<Vec<ExtractPlan>, XzipError> {
let mut plans = Vec::new();
for i in 0..archive.len() {
let entry = archive.by_index(i)?;
let raw_name = entry.name_raw();
let decoded_name = encoding.decode(raw_name)?;
if !filter.allows(&decoded_name) {
continue;
}
let safe_rel = sanitize_relative_path(&decoded_name)?;
let outpath = output_dir.join(&safe_rel);
plans.push(ExtractPlan {
index: i,
outpath,
display_name: decoded_name,
is_dir: entry.is_dir(),
uncompressed_size: entry.size(),
compressed_size: entry.compressed_size(),
});
}
Ok(plans)
}
fn create_output_dirs(output_dir: &Path, plans: &[ExtractPlan]) -> Result<usize, XzipError> {
let mut created = HashSet::new();
ensure_dir(output_dir, &mut created)?;
let mut dirs_created = 0usize;
for plan in plans {
if plan.is_dir {
if ensure_dir(&plan.outpath, &mut created)? {
dirs_created += 1;
}
} else if let Some(parent) = plan.outpath.parent() {
if ensure_dir(parent, &mut created)? {
dirs_created += 1;
}
}
}
Ok(dirs_created)
}
fn print_file_progress(verbosity: Verbosity, plan: &ExtractPlan) {
if verbosity.is_extra() {
eprintln!(
"extracting: {} ({} -> {} bytes)",
plan.display_name, plan.uncompressed_size, plan.compressed_size
);
}
}
fn maybe_print_throttled_progress(processed: usize, total: usize) {
if processed == total || processed % PROGRESS_INTERVAL == 0 {
eprintln!("extracting: {processed}/{total} files...");
}
}
fn extract_file(
archive: &mut zip::read::ZipArchive<File>,
plan: &ExtractPlan,
) -> Result<(), XzipError> {
let mut entry = archive.by_index(plan.index)?;
let mut outfile = BufWriter::with_capacity(WRITE_BUFFER_SIZE, File::create(&plan.outpath)?);
io::copy(&mut entry, &mut outfile)?;
Ok(())
}
fn extract_sequential(
input_zip: &Path,
file_plans: &[ExtractPlan],
verbosity: Verbosity,
) -> Result<(), XzipError> {
let total = file_plans.len();
let mut archive = open_archive(input_zip)?;
for (processed, plan) in file_plans.iter().enumerate() {
if verbosity.is_verbose() {
if verbosity.is_extra() {
print_file_progress(verbosity, plan);
} else {
maybe_print_throttled_progress(processed + 1, total);
}
}
extract_file(&mut archive, plan)?;
}
Ok(())
}
fn extract_parallel(
input_zip: &Path,
file_plans: &[ExtractPlan],
jobs: usize,
verbosity: Verbosity,
) -> Result<(), XzipError> {
let input_zip = input_zip.to_path_buf();
let total = file_plans.len();
let processed = AtomicUsize::new(0);
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(jobs)
.build()
.map_err(|err| XzipError::Io(io::Error::other(err.to_string())))?;
pool.install(|| {
file_plans.par_iter().try_for_each_init(
|| {
open_archive(&input_zip)
.expect("archive was readable during scan and should remain readable")
},
|archive, plan| {
extract_file(archive, plan)?;
if verbosity.is_verbose() && !verbosity.is_extra() {
let count = processed.fetch_add(1, Ordering::Relaxed) + 1;
maybe_print_throttled_progress(count, total);
}
Ok::<(), XzipError>(())
},
)
})?;
Ok(())
}
pub fn unpack_archive(
input_zip: &Path,
output_dir: &Path,
encoding: EncodingKind,
include: &[String],
exclude: &[String],
verbosity: Verbosity,
jobs: usize,
) -> Result<(), XzipError> {
let started = Instant::now();
let mut archive = open_archive(input_zip)?;
let filter = PathFilter::new(include, exclude)?;
let plans = scan_entries(&mut archive, output_dir, encoding, &filter)?;
let dirs_created = create_output_dirs(output_dir, &plans)?;
let file_plans: Vec<_> = plans.iter().filter(|plan| !plan.is_dir).cloned().collect();
let files_extracted = file_plans.len();
let total_bytes = total_uncompressed_bytes(&plans);
let effective_jobs = if verbosity.is_extra() {
1
} else {
resolve_jobs(jobs, files_extracted, total_bytes)
};
if effective_jobs == 1 {
extract_sequential(input_zip, &file_plans, verbosity)?;
} else {
extract_parallel(input_zip, &file_plans, effective_jobs, verbosity)?;
}
if verbosity.is_verbose() {
let elapsed = started.elapsed();
let file_label = if files_extracted == 1 {
"file"
} else {
"files"
};
eprintln!(
"{files_extracted} {file_label} extracted ({dirs_created} directories) in {:.2}s",
elapsed.as_secs_f64()
);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{resolve_jobs, PARALLEL_MIN_BYTES, PARALLEL_MIN_FILES};
#[test]
fn auto_jobs_stays_sequential_for_small_archives() {
assert_eq!(resolve_jobs(0, PARALLEL_MIN_FILES - 1, PARALLEL_MIN_BYTES), 1);
assert_eq!(resolve_jobs(0, PARALLEL_MIN_FILES, PARALLEL_MIN_BYTES - 1), 1);
}
#[test]
fn auto_jobs_parallelizes_large_archives() {
assert!(resolve_jobs(0, PARALLEL_MIN_FILES, PARALLEL_MIN_BYTES) > 1);
}
#[test]
fn explicit_jobs_are_honored() {
assert_eq!(resolve_jobs(1, 10_000, PARALLEL_MIN_BYTES * 10), 1);
assert_eq!(resolve_jobs(4, 10_000, PARALLEL_MIN_BYTES * 10), 4);
}
}