diskforge_core/

duplicate.rs

use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::SystemTime;

use anyhow::Result;
use ignore::WalkBuilder;
use rayon::prelude::*;

/// Options for duplicate file detection.
#[derive(Debug, Clone, Default)]
pub struct DupeOptions {
    /// Root directories to scan.
    pub root_paths: Vec<PathBuf>,
    /// Minimum file size in bytes (files smaller are skipped). Default: 1.
    pub min_size: Option<u64>,
    /// Minimum number of files to form a duplicate group. Default: 2.
    pub min_group_size: Option<usize>,
}

/// A single file that is part of a duplicate group.
#[derive(Debug, Clone)]
pub struct DuplicateFile {
    pub path: PathBuf,
    pub size: u64,
    pub mtime: Option<SystemTime>,
    pub inode: u64,
    pub device: u64,
}

/// A group of duplicate files sharing the same content.
#[derive(Debug, Clone)]
pub struct DuplicateGroup {
    /// blake3 hash of the file content.
    pub hash: String,
    /// Size of each individual file in the group.
    pub size: u64,
    /// The duplicate files.
    pub files: Vec<DuplicateFile>,
    /// True if all files in this group have the APFS clone flag set.
    pub is_clone_group: bool,
}

impl DuplicateGroup {
    /// Total wasted space (all copies minus one).
    pub fn wasted_space(&self) -> u64 {
        if self.files.len() <= 1 {
            0
        } else {
            self.size * (self.files.len() as u64 - 1)
        }
    }
}

/// Result of a duplicate scan.
#[derive(Debug)]
pub struct DupeResult {
    pub groups: Vec<DuplicateGroup>,
    pub total_wasted: u64,
    pub files_scanned: u64,
    pub empty_files_skipped: u64,
}

/// Progress counters for the duplicate scan, polled by the CLI.
pub struct DupeProgress {
    pub files_walked: AtomicU64,
    pub size_groups: AtomicU64,
    pub files_to_hash: AtomicU64,
    pub files_hashed: AtomicU64,
    pub bytes_hashed: AtomicU64,
}

impl Default for DupeProgress {
    fn default() -> Self {
        Self {
            files_walked: AtomicU64::new(0),
            size_groups: AtomicU64::new(0),
            files_to_hash: AtomicU64::new(0),
            files_hashed: AtomicU64::new(0),
            bytes_hashed: AtomicU64::new(0),
        }
    }
}

impl DupeProgress {
    pub fn new() -> Self {
        Self::default()
    }
}

/// Paths always excluded from duplicate scanning.
const DEFAULT_EXCLUSIONS: &[&str] = &[
    "/System",
    "/usr/bin",
    "/usr/lib",
    "/usr/libexec",
    "/usr/sbin",
    "/usr/share",
    "/bin",
    "/sbin",
    "/private/var/db",
];

/// Path components that signal an excluded tree.
const EXCLUDED_COMPONENTS: &[&str] = &[".git", ".app"];

/// Check if a path should be excluded.
fn should_exclude(path: &Path) -> bool {
    let path_str = path.to_string_lossy();
    for excl in DEFAULT_EXCLUSIONS {
        if path_str.starts_with(excl) {
            return true;
        }
    }
    for component in path.components() {
        let s = component.as_os_str().to_string_lossy();
        for excl in EXCLUDED_COMPONENTS {
            if s.ends_with(excl) {
                return true;
            }
        }
    }
    false
}

/// Size of the partial hash read (first 4KB).
const PARTIAL_HASH_SIZE: usize = 4096;

/// Files larger than this use mmap for full hashing.
const MMAP_THRESHOLD: u64 = 1024 * 1024; // 1 MB

/// Entry point: find duplicate files under the given paths.
pub fn find_duplicates(
    options: &DupeOptions,
    progress: Option<&DupeProgress>,
) -> Result<DupeResult> {
    let min_size = options.min_size.unwrap_or(1);
    let min_group = options.min_group_size.unwrap_or(2);

    // Stage 0: Walk filesystem and collect files, grouping by size.
    let (size_groups, empty_count) = walk_and_group_by_size(options, min_size, progress)?;

    if let Some(p) = progress {
        p.size_groups
            .store(size_groups.len() as u64, Ordering::Relaxed);
    }

    // Filter to groups with enough files to form duplicates.
    let size_groups: Vec<Vec<DuplicateFile>> = size_groups
        .into_values()
        .filter(|group| group.len() >= min_group)
        .collect();

    // Count files that need hashing for progress.
    let files_to_hash: u64 = size_groups.iter().map(|g| g.len() as u64).sum();
    if let Some(p) = progress {
        p.files_to_hash.store(files_to_hash, Ordering::Relaxed);
    }

    // Stage 1: Inode deduplication — remove hard links within each group.
    let deduped_groups: Vec<Vec<DuplicateFile>> = size_groups
        .into_iter()
        .map(dedup_by_inode)
        .filter(|g| g.len() >= min_group)
        .collect();

    // Stage 2: Partial hash (first 4KB) — group by partial blake3 hash.
    let partial_groups: Vec<Vec<DuplicateFile>> = deduped_groups
        .into_par_iter()
        .flat_map(|group| partial_hash_group(group, progress))
        .filter(|g| g.len() >= min_group)
        .collect();

    // Stage 3: Full hash — confirm duplicates with complete file hash.
    let final_groups: Vec<DuplicateGroup> = partial_groups
        .into_par_iter()
        .flat_map(|group| full_hash_group(group, progress))
        .filter(|g| g.files.len() >= min_group)
        .collect();

    // Sort by wasted space descending.
    let mut groups = final_groups;
    groups.sort_by_key(|g| std::cmp::Reverse(g.wasted_space()));

    let total_wasted: u64 = groups.iter().map(|g| g.wasted_space()).sum();
    let files_scanned = progress
        .map(|p| p.files_walked.load(Ordering::Relaxed))
        .unwrap_or(0);

    Ok(DupeResult {
        groups,
        total_wasted,
        files_scanned,
        empty_files_skipped: empty_count,
    })
}

/// Stage 0: Walk the filesystem and group files by size.
/// Returns (size_groups, empty_files_skipped).
fn walk_and_group_by_size(
    options: &DupeOptions,
    min_size: u64,
    progress: Option<&DupeProgress>,
) -> Result<(HashMap<u64, Vec<DuplicateFile>>, u64)> {
    let collector: Arc<Mutex<HashMap<u64, Vec<DuplicateFile>>>> =
        Arc::new(Mutex::new(HashMap::new()));
    let empty_count = Arc::new(AtomicU64::new(0));

    for root in &options.root_paths {
        if !root.exists() {
            continue;
        }

        let mut builder = WalkBuilder::new(root);
        builder
            .hidden(false)
            .git_ignore(true)
            .git_global(false)
            .git_exclude(false)
            .follow_links(false) // T2.4: Don't follow symlinks
            .threads(num_cpus());

        let walker = builder.build_parallel();
        let collector_ref = Arc::clone(&collector);
        let empty_ref = Arc::clone(&empty_count);

        walker.run(|| {
            let collector = Arc::clone(&collector_ref);
            let empty_count = Arc::clone(&empty_ref);

            Box::new(move |entry| {
                let Ok(entry) = entry else {
                    return ignore::WalkState::Continue;
                };

                let path = entry.path();

                if should_exclude(path) {
                    return ignore::WalkState::Skip;
                }

                let Some(file_type) = entry.file_type() else {
                    return ignore::WalkState::Continue;
                };
                if !file_type.is_file() {
                    return ignore::WalkState::Continue;
                }

                if let Some(p) = progress {
                    p.files_walked.fetch_add(1, Ordering::Relaxed);
                }

                let Ok(meta) = entry.metadata() else {
                    // T2.3: Skip unreadable files gracefully.
                    return ignore::WalkState::Continue;
                };

                let size = meta.len();

                // T2.1: Skip empty files.
                if size == 0 {
                    empty_count.fetch_add(1, Ordering::Relaxed);
                    return ignore::WalkState::Continue;
                }

                if size < min_size {
                    return ignore::WalkState::Continue;
                }

                let file = DuplicateFile {
                    path: path.to_path_buf(),
                    size,
                    mtime: meta.modified().ok(),
                    inode: meta.ino(),
                    device: meta.dev(),
                };

                if let Ok(mut map) = collector.lock() {
                    map.entry(size).or_default().push(file);
                }

                ignore::WalkState::Continue
            })
        });
    }

    let map = Arc::try_unwrap(collector)
        .map(|mutex| mutex.into_inner().unwrap_or_else(|e| e.into_inner()))
        .unwrap_or_else(|arc| arc.lock().unwrap().clone());

    let empty = empty_count.load(Ordering::Relaxed);

    Ok((map, empty))
}

/// Stage 1: Remove hard links (same device + inode) from a size group.
fn dedup_by_inode(group: Vec<DuplicateFile>) -> Vec<DuplicateFile> {
    let mut seen: HashMap<(u64, u64), usize> = HashMap::new();
    let mut result = Vec::new();

    for file in group {
        let key = (file.device, file.inode);
        if seen.contains_key(&key) {
            continue; // Hard link to an already-seen file
        }
        seen.insert(key, result.len());
        result.push(file);
    }

    result
}

/// Stage 2: Group files by partial hash (first 4KB).
/// Returns sub-groups that share the same partial hash.
fn partial_hash_group(
    group: Vec<DuplicateFile>,
    progress: Option<&DupeProgress>,
) -> Vec<Vec<DuplicateFile>> {
    let mut hash_groups: HashMap<[u8; 32], Vec<DuplicateFile>> = HashMap::new();

    for file in group {
        match partial_hash(&file.path) {
            Ok(hash) => {
                if let Some(p) = progress {
                    p.files_hashed.fetch_add(1, Ordering::Relaxed);
                    let hashed = PARTIAL_HASH_SIZE.min(file.size as usize) as u64;
                    p.bytes_hashed.fetch_add(hashed, Ordering::Relaxed);
                }
                hash_groups.entry(hash).or_default().push(file);
            }
            Err(_) => {
                // T2.3: Skip files we can't read.
            }
        }
    }

    hash_groups.into_values().collect()
}

/// Compute blake3 hash of the first 4KB of a file.
fn partial_hash(path: &Path) -> Result<[u8; 32]> {
    let mut file = File::open(path)?;
    let mut buf = vec![0u8; PARTIAL_HASH_SIZE];
    let n = file.read(&mut buf)?;
    buf.truncate(n);
    Ok(*blake3::hash(&buf).as_bytes())
}

/// Stage 3: Group files by full hash and build DuplicateGroup results.
fn full_hash_group(
    group: Vec<DuplicateFile>,
    progress: Option<&DupeProgress>,
) -> Vec<DuplicateGroup> {
    let mut hash_groups: HashMap<String, Vec<DuplicateFile>> = HashMap::new();

    for file in &group {
        match full_hash(&file.path, file.size) {
            Ok(hash) => {
                if let Some(p) = progress {
                    p.files_hashed.fetch_add(1, Ordering::Relaxed);
                    p.bytes_hashed.fetch_add(file.size, Ordering::Relaxed);
                }
                hash_groups.entry(hash).or_default().push(file.clone());
            }
            Err(_) => {
                // T2.3: Skip files we can't read.
            }
        }
    }

    let size = group.first().map(|f| f.size).unwrap_or(0);

    hash_groups
        .into_iter()
        .map(|(hash, files)| {
            let is_clone_group = check_apfs_clones(&files);
            DuplicateGroup {
                hash,
                size,
                files,
                is_clone_group,
            }
        })
        .collect()
}

/// Compute the full blake3 hash of a file.
/// Uses mmap for files > 1MB, buffered read otherwise.
fn full_hash(path: &Path, size: u64) -> Result<String> {
    if size > MMAP_THRESHOLD {
        full_hash_mmap(path)
    } else {
        full_hash_buffered(path)
    }
}

/// Full hash using buffered read (for files <= 1MB).
fn full_hash_buffered(path: &Path) -> Result<String> {
    let mut file = File::open(path)?;
    let mut hasher = blake3::Hasher::new();
    let mut buf = vec![0u8; 65536];
    loop {
        let n = file.read(&mut buf)?;
        if n == 0 {
            break;
        }
        hasher.update(&buf[..n]);
    }
    Ok(hasher.finalize().to_hex().to_string())
}

/// Full hash using memory-mapped I/O (for files > 1MB).
fn full_hash_mmap(path: &Path) -> Result<String> {
    let file = File::open(path)?;
    // SAFETY: We only read the file. The file could be modified by another process
    // during hashing, but that would just produce a different hash — not UB.
    let mmap = unsafe { memmap2::MmapOptions::new().map(&file)? };
    let hash = blake3::hash(&mmap);
    Ok(hash.to_hex().to_string())
}

/// T2.2: Check if all files in a group have the APFS clone flag (EF_MAY_SHARE_BLOCKS).
#[cfg(target_os = "macos")]
fn check_apfs_clones(files: &[DuplicateFile]) -> bool {
    use std::os::macos::fs::MetadataExt;

    /// EF_MAY_SHARE_BLOCKS flag value.
    const EF_MAY_SHARE_BLOCKS: u32 = 0x0000_0008;

    if files.is_empty() {
        return false;
    }

    files.iter().all(|f| {
        std::fs::metadata(&f.path)
            .map(|m| (m.st_flags() & EF_MAY_SHARE_BLOCKS) != 0)
            .unwrap_or(false)
    })
}

#[cfg(not(target_os = "macos"))]
fn check_apfs_clones(_files: &[DuplicateFile]) -> bool {
    false
}

/// Get a reasonable thread count.
fn num_cpus() -> usize {
    std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(4)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;

    fn setup_test_dir(name: &str) -> PathBuf {
        let tmp = std::env::temp_dir().join(format!("diskforge_test_dupe_{name}"));
        let _ = fs::remove_dir_all(&tmp);
        fs::create_dir_all(&tmp).unwrap();
        tmp
    }

    fn create_file(dir: &Path, name: &str, content: &[u8]) -> PathBuf {
        let path = dir.join(name);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).unwrap();
        }
        fs::write(&path, content).unwrap();
        path
    }

    #[test]
    fn unique_sizes_no_duplicates() {
        let tmp = setup_test_dir("unique_sizes");
        create_file(&tmp, "a.txt", &[1u8; 100]);
        create_file(&tmp, "b.txt", &[2u8; 200]);
        create_file(&tmp, "c.txt", &vec![3u8; 300]);

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        assert!(
            result.groups.is_empty(),
            "Files with unique sizes should produce no groups"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn same_size_different_content_no_duplicates() {
        let tmp = setup_test_dir("same_size_diff_content");
        create_file(&tmp, "a.txt", b"hello world!");
        create_file(&tmp, "b.txt", b"goodbye now!");

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        assert!(
            result.groups.is_empty(),
            "Same size but different content should not be grouped"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn identical_files_grouped() {
        let tmp = setup_test_dir("identical_files");
        let content = b"this is duplicate content for testing";
        create_file(&tmp, "a.txt", content);
        create_file(&tmp, "b.txt", content);
        create_file(&tmp, "c.txt", content);

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        assert_eq!(result.groups.len(), 1, "Should find 1 duplicate group");
        assert_eq!(
            result.groups[0].files.len(),
            3,
            "Group should contain 3 files"
        );
        assert_eq!(
            result.groups[0].wasted_space(),
            content.len() as u64 * 2,
            "Wasted = size * (count - 1)"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn hard_links_deduplicated() {
        let tmp = setup_test_dir("hard_links");
        let content = b"content for hard link test that is long enough";
        let original = create_file(&tmp, "original.txt", content);
        let link_path = tmp.join("hardlink.txt");
        fs::hard_link(&original, &link_path).unwrap();

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        assert!(
            result.groups.is_empty(),
            "Hard links to the same inode should not be reported as duplicates"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn empty_files_excluded() {
        let tmp = setup_test_dir("empty_files");
        create_file(&tmp, "a.txt", b"");
        create_file(&tmp, "b.txt", b"");
        create_file(&tmp, "c.txt", b"");

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: None, // Use default (1 byte)
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        assert!(result.groups.is_empty(), "Empty files should be excluded");
        assert_eq!(result.empty_files_skipped, 3);

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn min_size_filter() {
        let tmp = setup_test_dir("min_size");
        let small = vec![1u8; 50];
        let big = vec![2u8; 500];
        create_file(&tmp, "small1.txt", &small);
        create_file(&tmp, "small2.txt", &small);
        create_file(&tmp, "big1.txt", &big);
        create_file(&tmp, "big2.txt", &big);

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(100),
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        assert_eq!(
            result.groups.len(),
            1,
            "Only the big files should form a group"
        );
        assert_eq!(result.groups[0].size, 500);

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn partial_hash_eliminates_different_content() {
        let tmp = setup_test_dir("partial_hash");
        // Same size, different first 4KB
        let mut content_a = vec![0xAAu8; 5000];
        let mut content_b = vec![0xBBu8; 5000];
        // Make the tails the same to ensure partial hash catches the difference
        content_a[4096..].fill(0xFF);
        content_b[4096..].fill(0xFF);

        create_file(&tmp, "a.bin", &content_a);
        create_file(&tmp, "b.bin", &content_b);

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        assert!(
            result.groups.is_empty(),
            "Different first 4KB should be eliminated by partial hash"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn progress_counters_work() {
        let tmp = setup_test_dir("progress");
        let content = b"duplicate content for progress test!!";
        create_file(&tmp, "a.txt", content);
        create_file(&tmp, "b.txt", content);
        create_file(&tmp, "unique.txt", b"something unique here");

        let progress = DupeProgress::new();
        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            ..Default::default()
        };

        find_duplicates(&opts, Some(&progress)).unwrap();
        assert!(progress.files_walked.load(Ordering::Relaxed) >= 3);
        assert!(progress.files_hashed.load(Ordering::Relaxed) > 0);

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn symlinks_not_followed() {
        let tmp = setup_test_dir("symlinks");
        let content = b"real file content for symlink test";
        let real_file = create_file(&tmp, "real.txt", content);
        let link_path = tmp.join("link.txt");
        std::os::unix::fs::symlink(&real_file, &link_path).unwrap();

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        // Only 1 real file, symlink should be ignored, so no duplicate group
        assert!(
            result.groups.is_empty(),
            "Symlinks should not be treated as duplicate candidates"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn multiple_duplicate_groups() {
        let tmp = setup_test_dir("multi_groups");
        let content_a = b"group A duplicate content here!";
        let content_b = vec![0xFFu8; 200];

        create_file(&tmp, "a1.txt", content_a);
        create_file(&tmp, "a2.txt", content_a);
        create_file(&tmp, "b1.bin", &content_b);
        create_file(&tmp, "b2.bin", &content_b);

        let opts = DupeOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            ..Default::default()
        };

        let result = find_duplicates(&opts, None).unwrap();
        assert_eq!(result.groups.len(), 2, "Should find 2 duplicate groups");

        fs::remove_dir_all(&tmp).ok();
    }

    #[cfg(target_os = "macos")]
    #[test]
    fn apfs_clone_detection() {
        use std::process::Command;

        let tmp = setup_test_dir("apfs_clones");
        let content = vec![0u8; 8192]; // Needs to be at least a block for cloning
        let original = create_file(&tmp, "original.bin", &content);
        let clone_path = tmp.join("clone.bin");

        // Use `cp -c` to create an APFS clone
        let status = Command::new("cp")
            .args([
                "-c",
                &original.to_string_lossy(),
                &clone_path.to_string_lossy(),
            ])
            .status();

        match status {
            Ok(s) if s.success() => {
                let opts = DupeOptions {
                    root_paths: vec![tmp.clone()],
                    min_size: Some(1),
                    ..Default::default()
                };

                let result = find_duplicates(&opts, None).unwrap();
                // Clones should still be detected as duplicates
                if !result.groups.is_empty() {
                    // On APFS, the clone group flag should be set
                    // (May not be set on non-APFS volumes)
                    let _is_clone = result.groups[0].is_clone_group;
                    // We just verify it doesn't crash; actual flag depends on filesystem
                }
            }
            _ => {
                // cp -c not supported (e.g., non-APFS volume) — skip
            }
        }

        fs::remove_dir_all(&tmp).ok();
    }
}