use rayon::prelude::*;
use std::fs::{self, File};
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use super::compinit::{CompFileDef, CompInitResult};
pub fn compdump(
result: &CompInitResult,
dump_path: &Path,
zsh_version: &str,
) -> std::io::Result<PathBuf> {
let host = hostname();
let pid = std::process::id();
let tmp = dump_path
.with_file_name(format!(
"{}.{}.{}",
dump_path
.file_name()
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_default(),
host,
pid
));
{
let mut file = File::create(&tmp)?;
writeln!(
file,
"#files: {}\tversion: {}",
result.files_scanned, zsh_version
)?;
write_assoc_dump(&mut file, "_comps", &result.comps)?;
write_assoc_dump(&mut file, "_services", &result.services)?;
write_assoc_dump(&mut file, "_patcomps", &result.patcomps)?;
write_assoc_dump(&mut file, "_postpatcomps", &result.postpatcomps)?;
write_assoc_dump(&mut file, "_compautos", &result.compautos)?;
let mut autoload_names: Vec<String> = result
.files
.iter()
.filter_map(|f| match &f.def {
CompFileDef::CompDef(_) => Some(f.name.clone()),
_ => None,
})
.collect();
autoload_names.sort();
autoload_names.dedup();
if !autoload_names.is_empty() {
writeln!(file, "autoload -Uz \\")?;
for (i, name) in autoload_names.iter().enumerate() {
let cont = if i + 1 < autoload_names.len() { " \\" } else { "" };
writeln!(file, " {}{}", name, cont)?;
}
}
let mut compautos_sorted: Vec<(&String, &String)> =
result.compautos.iter().collect();
compautos_sorted.sort_by(|a, b| a.0.cmp(b.0));
for (name, opts) in &compautos_sorted {
let opt_str = if opts.is_empty() {
"-Uz".to_string()
} else {
opts.to_string()
};
writeln!(file, "autoload {} {}", opt_str, name)?;
}
file.sync_all()?;
}
fs::rename(&tmp, dump_path)?;
Ok(dump_path.to_path_buf())
}
fn write_assoc_dump<W: Write>(
w: &mut W,
name: &str,
entries: &std::collections::HashMap<String, String>,
) -> std::io::Result<()> {
writeln!(w, "typeset -gHA {}", name)?;
if entries.is_empty() {
writeln!(w, "{}=(\n)", name)?;
return Ok(());
}
let mut sorted: Vec<(&String, &String)> = entries.iter().collect();
sorted.sort_by(|a, b| a.0.cmp(b.0));
writeln!(w, "{}=(", name)?;
for (k, v) in &sorted {
writeln!(w, " {} {}", qq(k), qq(v))?;
}
writeln!(w, ")")?;
Ok(())
}
pub(super) fn qq(s: &str) -> String {
let mut out = String::with_capacity(s.len() + 2);
out.push('\'');
out.push_str(&s.replace('\'', "'\\''"));
out.push('\'');
out
}
fn hostname() -> String {
std::env::var("HOST")
.or_else(|_| std::env::var("HOSTNAME"))
.ok()
.filter(|s| !s.is_empty())
.unwrap_or_else(|| "localhost".to_string())
}
pub fn check_dump(dump_path: &Path, fpath: &[PathBuf], zsh_version: &str) -> bool {
let file = match File::open(dump_path) {
Ok(f) => f,
Err(_) => return false,
};
let mut reader = BufReader::new(file);
let mut first_line = String::new();
if reader.read_line(&mut first_line).is_err() {
return false;
}
let line = first_line.trim_end_matches('\n');
let stripped = match line.strip_prefix("#files:") {
Some(s) => s.trim_start(),
None => return false,
};
let mut parts = stripped.splitn(2, '\t');
let n_str = parts.next().unwrap_or("").trim();
let version_part = parts.next().unwrap_or("");
let stored_version = match version_part.strip_prefix("version:") {
Some(s) => s.trim(),
None => return false,
};
let stored_count: usize = match n_str.parse() {
Ok(n) => n,
Err(_) => return false,
};
let current_count: usize = fpath
.par_iter()
.filter(|dir| dir.as_os_str() != "." && dir.exists())
.map(|dir| {
fs::read_dir(dir)
.map(|entries| {
entries
.filter_map(|e| e.ok())
.filter(|e| e.file_name().to_string_lossy().starts_with('_'))
.count()
})
.unwrap_or(0)
})
.sum();
stored_count == current_count && stored_version == zsh_version
}
pub(super) fn escape_zsh_string(s: &str) -> String {
s.replace('\'', "'\\''")
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
#[test]
fn test_escape_zsh_string() {
assert_eq!(escape_zsh_string("hello"), "hello");
assert_eq!(escape_zsh_string("it's"), "it'\\''s");
}
#[test]
fn qq_wraps_in_single_quotes_and_escapes() {
assert_eq!(qq("plain"), "'plain'");
assert_eq!(qq("it's"), "'it'\\''s'");
assert_eq!(qq(""), "''");
}
fn empty_result() -> CompInitResult {
CompInitResult {
files_scanned: 3,
dirs_scanned: 0,
scan_time_ms: 0,
files: Vec::new(),
comps: HashMap::new(),
services: HashMap::new(),
patcomps: HashMap::new(),
postpatcomps: HashMap::new(),
compautos: HashMap::new(),
}
}
#[test]
fn header_matches_upstream_format() {
let mut r = empty_result();
r.files_scanned = 42;
let tmp = std::env::temp_dir().join("zshrs_compdump_header_test");
let _ = std::fs::remove_file(&tmp);
let _ = compdump(&r, &tmp, "5.9").unwrap();
let content = std::fs::read_to_string(&tmp).unwrap();
let first_line = content.lines().next().unwrap();
assert_eq!(first_line, "#files: 42\tversion: 5.9");
let _ = std::fs::remove_file(&tmp);
}
#[test]
fn check_dump_accepts_upstream_header() {
let tmp = std::env::temp_dir().join("zshrs_check_dump_upstream");
std::fs::write(
&tmp,
"#files: 0\tversion: 5.9\ntypeset -gHA _comps\n_comps=(\n)\n",
)
.unwrap();
assert!(check_dump(&tmp, &[], "5.9"));
assert!(!check_dump(&tmp, &[], "5.10"));
let _ = std::fs::remove_file(&tmp);
}
#[test]
fn check_dump_rejects_old_zshrs_format() {
let tmp = std::env::temp_dir().join("zshrs_check_dump_old_format");
std::fs::write(&tmp, "#compdump 0 . 5.9\n").unwrap();
assert!(!check_dump(&tmp, &[], "5.9"));
let _ = std::fs::remove_file(&tmp);
}
#[test]
fn assoc_dump_sorts_keys_deterministically() {
let mut r = empty_result();
r.comps.insert("zfoo".to_string(), "_z".to_string());
r.comps.insert("alpha".to_string(), "_a".to_string());
r.comps.insert("mike".to_string(), "_m".to_string());
let tmp = std::env::temp_dir().join("zshrs_compdump_sort_test");
let _ = std::fs::remove_file(&tmp);
let _ = compdump(&r, &tmp, "5.9").unwrap();
let content = std::fs::read_to_string(&tmp).unwrap();
let comps_pos = content.find("_comps=(").unwrap();
let after = &content[comps_pos..];
let alpha_pos = after.find("'alpha'").unwrap();
let mike_pos = after.find("'mike'").unwrap();
let zfoo_pos = after.find("'zfoo'").unwrap();
assert!(alpha_pos < mike_pos, "alpha must precede mike");
assert!(mike_pos < zfoo_pos, "mike must precede zfoo");
let _ = std::fs::remove_file(&tmp);
}
#[test]
fn rename_replaces_existing_dump_atomically() {
let tmp = std::env::temp_dir().join("zshrs_compdump_atomic_test");
let _ = std::fs::remove_file(&tmp);
let _ = compdump(&empty_result(), &tmp, "5.9").unwrap();
assert!(tmp.exists());
let parent = tmp.parent().unwrap();
let stray: Vec<_> = std::fs::read_dir(parent)
.unwrap()
.filter_map(|e| e.ok())
.filter(|e| {
let n = e.file_name().to_string_lossy().into_owned();
n.starts_with("zshrs_compdump_atomic_test.") && n != "zshrs_compdump_atomic_test"
})
.collect();
assert!(stray.is_empty(), "temp file leaked: {:?}", stray);
let _ = std::fs::remove_file(&tmp);
}
}