kernel/discovery/
duplicates.rs

1//! Detecting duplicate model weights by size then a cheap content fingerprint.
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::fs::File;
5use std::io::{Read, Seek, SeekFrom};
6use std::path::{Path, PathBuf};
7
8use sha2::{Digest, Sha256};
9
10/// The default minimum file size (256 MiB) considered for duplicate detection —
11/// small files are not worth reporting.
12pub const DEFAULT_THRESHOLD: i64 = 256 << 20;
13
14const SAMPLE_SIZE: u64 = 1 << 20;
15
16/// A set of files found to be duplicates of one another.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct DuplicateGroup {
19    /// The display names of the duplicate models, sorted.
20    pub names: Vec<String>,
21    /// The paths of the duplicate files, sorted.
22    pub paths: Vec<String>,
23    /// The bytes that could be reclaimed by keeping just one copy.
24    pub wasted_bytes: i64,
25}
26
27/// Find groups of duplicate weight files among `candidates` (name, path pairs).
28/// Files smaller than `threshold` are ignored. Candidates are bucketed by exact
29/// size, then confirmed by a fingerprint over the first and last megabyte, so
30/// only same-size, same-fingerprint files are reported. Groups are returned
31/// most-wasteful first.
32pub fn detect(candidates: &[(String, PathBuf)], threshold: i64) -> Vec<DuplicateGroup> {
33    let mut seen_paths: BTreeSet<PathBuf> = BTreeSet::new();
34    let mut by_size: BTreeMap<i64, Vec<(&str, &Path)>> = BTreeMap::new();
35    for (name, path) in candidates {
36        let identity = std::fs::canonicalize(path).unwrap_or_else(|_| path.clone());
37        if !seen_paths.insert(identity) {
38            continue;
39        }
40        let Ok(size) = file_size(path) else {
41            continue;
42        };
43        if size >= threshold && size > 0 {
44            by_size.entry(size).or_default().push((name, path));
45        }
46    }
47
48    let mut groups: Vec<DuplicateGroup> = Vec::new();
49    for (size, bucket) in by_size {
50        if bucket.len() < 2 {
51            continue;
52        }
53        let mut by_fingerprint: BTreeMap<[u8; 32], Vec<(&str, &Path)>> = BTreeMap::new();
54        for (name, path) in bucket {
55            if let Some(digest) = fingerprint(path, size as u64) {
56                by_fingerprint.entry(digest).or_default().push((name, path));
57            }
58        }
59        for matches in by_fingerprint.into_values() {
60            if matches.len() < 2 {
61                continue;
62            }
63            let mut names: Vec<String> =
64                matches.iter().map(|(name, _)| (*name).to_owned()).collect();
65            let mut paths: Vec<String> = matches
66                .iter()
67                .map(|(_, path)| path.to_string_lossy().into_owned())
68                .collect();
69            names.sort();
70            paths.sort();
71            groups.push(DuplicateGroup {
72                names,
73                paths,
74                wasted_bytes: (matches.len() as i64 - 1) * size,
75            });
76        }
77    }
78
79    groups.sort_by_key(|group| std::cmp::Reverse(group.wasted_bytes));
80    groups
81}
82
83fn file_size(path: &Path) -> std::io::Result<i64> {
84    Ok(std::fs::metadata(path)?.len() as i64)
85}
86
87/// A content fingerprint for the file at `path`: the lowercase-hex SHA-256 of its
88/// first and last megabyte (or the whole file, if small). `None` if unreadable.
89/// Two files with the same fingerprint and size are treated as the same weights.
90pub fn content_fingerprint(path: &Path) -> Option<String> {
91    let size = std::fs::metadata(path).ok()?.len();
92    fingerprint(path, size).map(hex::encode)
93}
94
95fn fingerprint(path: &Path, size: u64) -> Option<[u8; 32]> {
96    let mut file = File::open(path).ok()?;
97    let mut hasher = Sha256::new();
98    if size <= SAMPLE_SIZE * 2 {
99        let mut whole = Vec::new();
100        file.take(size).read_to_end(&mut whole).ok()?;
101        hasher.update(&whole);
102    } else {
103        let mut head = vec![0u8; SAMPLE_SIZE as usize];
104        file.read_exact(&mut head).ok()?;
105        file.seek(SeekFrom::Start(size - SAMPLE_SIZE)).ok()?;
106        let mut tail = vec![0u8; SAMPLE_SIZE as usize];
107        file.read_exact(&mut tail).ok()?;
108        hasher.update(&head);
109        hasher.update(&tail);
110    }
111    Some(hasher.finalize().into())
112}
kernel/discovery/duplicates.rs

kernel/discovery/
duplicates.rs
