modelshelf 0.1.0

A shared local LLM model registry: discover, deduplicate, download, and update models across desktop apps.
Documentation
//! Content-based model identity: full hashing, quick fingerprints, and a
//! stat-keyed cache that avoids re-hashing unchanged multi-gigabyte files.

use std::collections::BTreeMap;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::Path;

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::registry::ModelId;
use crate::{Error, Result};

/// Bytes hashed at each end of the file for the quick fingerprint.
const FINGERPRINT_WINDOW: u64 = 1024 * 1024;

/// A cheap pre-filter for duplicate detection: two files can only be
/// identical if their fingerprints match. Only fingerprint-colliding files
/// are worth a full hash.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct QuickFingerprint {
    /// File size in bytes.
    pub size_bytes: u64,
    /// sha256 of the first `min(size, 1 MiB)` bytes, hex.
    pub head_sha256: String,
    /// sha256 of the last `min(size, 1 MiB)` bytes, hex.
    pub tail_sha256: String,
}

/// Compute the full sha256 of a file, streaming.
pub fn full_sha256(path: &Path) -> Result<String> {
    let mut f = File::open(path).map_err(|e| Error::io(path, e))?;
    let mut hasher = Sha256::new();
    let mut buf = vec![0u8; 1024 * 1024];
    loop {
        let n = f.read(&mut buf).map_err(|e| Error::io(path, e))?;
        if n == 0 {
            break;
        }
        hasher.update(&buf[..n]);
    }
    Ok(hex(&hasher.finalize()))
}

/// Compute the [`QuickFingerprint`] of a file (reads at most 2 MiB).
pub fn quick_fingerprint(path: &Path) -> Result<QuickFingerprint> {
    let mut f = File::open(path).map_err(|e| Error::io(path, e))?;
    let size = f.metadata().map_err(|e| Error::io(path, e))?.len();
    let window = FINGERPRINT_WINDOW.min(size);

    let mut head = vec![0u8; window as usize];
    f.read_exact(&mut head).map_err(|e| Error::io(path, e))?;

    f.seek(SeekFrom::Start(size - window))
        .map_err(|e| Error::io(path, e))?;
    let mut tail = vec![0u8; window as usize];
    f.read_exact(&mut tail).map_err(|e| Error::io(path, e))?;

    Ok(QuickFingerprint {
        size_bytes: size,
        head_sha256: hex(&Sha256::digest(&head)),
        tail_sha256: hex(&Sha256::digest(&tail)),
    })
}

fn hex(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        use std::fmt::Write;
        let _ = write!(s, "{b:02x}");
    }
    s
}

// ---- hash cache ---------------------------------------------------------

/// One cached hash, valid only while the file's stat signature is unchanged.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
struct CachedHash {
    size_bytes: u64,
    mtime_unix_ms: i64,
    sha256: String,
}

/// A persistent path→hash cache keyed on `(size, mtime)`.
///
/// Stored as `hash_cache.json` in the shelf home. Purely an optimization:
/// deleting the file is always safe.
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct HashCache {
    #[serde(default)]
    entries: BTreeMap<String, CachedHash>,
    #[serde(skip)]
    dirty: bool,
}

/// File name of the hash cache inside the shelf home.
pub const HASH_CACHE_FILE: &str = "hash_cache.json";

impl HashCache {
    /// Load the cache, or start empty if missing/unreadable (it is only an
    /// optimization, so an unreadable cache is discarded, not an error).
    pub fn load(shelf_home: &Path) -> Self {
        let path = shelf_home.join(HASH_CACHE_FILE);
        match std::fs::read(&path) {
            Ok(bytes) => serde_json::from_slice(&bytes).unwrap_or_default(),
            Err(_) => HashCache::default(),
        }
    }

    /// Persist the cache if it changed since loading.
    pub fn save(&mut self, shelf_home: &Path) -> Result<()> {
        if !self.dirty {
            return Ok(());
        }
        let bytes = serde_json::to_vec(&self)
            .map_err(|e| Error::io(shelf_home, std::io::Error::other(e)))?;
        crate::registry::lock::write_atomic(&shelf_home.join(HASH_CACHE_FILE), &bytes)?;
        self.dirty = false;
        Ok(())
    }

    /// Full sha256 of `path`, reusing the cached value when the file's size
    /// and mtime are unchanged. A cache hit does not open the file at all.
    pub fn sha256_of(&mut self, path: &Path) -> Result<String> {
        let meta = std::fs::metadata(path).map_err(|e| Error::io(path, e))?;
        let stamp = mtime_unix_ms(&meta);
        let key = path.to_string_lossy().into_owned();
        if let Some(hit) = self.entries.get(&key) {
            if hit.size_bytes == meta.len() && hit.mtime_unix_ms == stamp {
                return Ok(hit.sha256.clone());
            }
        }
        let sha = full_sha256(path)?;
        self.entries.insert(
            key,
            CachedHash {
                size_bytes: meta.len(),
                mtime_unix_ms: stamp,
                sha256: sha.clone(),
            },
        );
        self.dirty = true;
        Ok(sha)
    }

    /// [`ModelId`] of `path`, via [`Self::sha256_of`].
    pub fn model_id_of(&mut self, path: &Path) -> Result<ModelId> {
        Ok(ModelId::from_sha256_hex(&self.sha256_of(path)?))
    }
}

fn mtime_unix_ms(meta: &std::fs::Metadata) -> i64 {
    meta.modified()
        .ok()
        .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
        .map(|d| d.as_millis() as i64)
        .unwrap_or(0)
}

/// Whether two paths refer to the same physical file (same device + inode on
/// Unix, same volume + file index on Windows). Hardlinked duplicates are
/// detected this way without reading any content.
pub fn same_inode(a: &Path, b: &Path) -> bool {
    #[cfg(unix)]
    {
        use std::os::unix::fs::MetadataExt;
        match (std::fs::metadata(a), std::fs::metadata(b)) {
            (Ok(ma), Ok(mb)) => ma.dev() == mb.dev() && ma.ino() == mb.ino(),
            _ => false,
        }
    }
    #[cfg(windows)]
    {
        // Stable Rust has no direct file-index accessor; equal content hash +
        // nlink>1 heuristics are handled elsewhere. Conservatively report
        // false so Windows never skips a verification step.
        let _ = (a, b);
        false
    }
    #[cfg(not(any(unix, windows)))]
    {
        let _ = (a, b);
        false
    }
}

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

    /// sha256("modelshelf") — precomputed known answer.
    const KNOWN: &str = "03e357c72c3f71b1d631a3623c6ac3232e36f2719d1c2f6a0fc568e13e6cb9c6";

    #[test]
    fn full_hash_known_answer() {
        let tmp = tempfile::tempdir().unwrap();
        let p = tmp.path().join("f");
        std::fs::write(&p, b"modelshelf").unwrap();
        assert_eq!(full_sha256(&p).unwrap(), KNOWN);
    }

    #[test]
    fn fingerprint_groups_only_identical_content() {
        let tmp = tempfile::tempdir().unwrap();
        let a = tmp.path().join("a");
        let b = tmp.path().join("b");
        let c = tmp.path().join("c");
        std::fs::write(&a, b"same-bytes-here").unwrap();
        std::fs::write(&b, b"same-bytes-here").unwrap();
        std::fs::write(&c, b"diff-bytes-here").unwrap(); // same length!
        let fa = quick_fingerprint(&a).unwrap();
        let fb = quick_fingerprint(&b).unwrap();
        let fc = quick_fingerprint(&c).unwrap();
        assert_eq!(fa, fb);
        assert_ne!(fa, fc);
        assert_eq!(fa.size_bytes, fc.size_bytes);
    }

    #[test]
    fn cache_hit_does_not_read_the_file() {
        let tmp = tempfile::tempdir().unwrap();
        let home = tmp.path();
        let p = home.join("model.bin");
        std::fs::write(&p, b"modelshelf").unwrap();

        let mut cache = HashCache::load(home);
        assert_eq!(cache.sha256_of(&p).unwrap(), KNOWN);
        cache.save(home).unwrap();

        // Poison the cache entry for the same stat signature: if sha256_of
        // returns the poisoned value, it proves the file was not re-read.
        let mut reloaded = HashCache::load(home);
        let key = p.to_string_lossy().into_owned();
        let entry = reloaded.entries.get_mut(&key).unwrap();
        entry.sha256 = "poisoned".into();
        assert_eq!(reloaded.sha256_of(&p).unwrap(), "poisoned");
    }

    #[test]
    fn cache_misses_when_size_changes() {
        let tmp = tempfile::tempdir().unwrap();
        let home = tmp.path();
        let p = home.join("model.bin");
        std::fs::write(&p, b"modelshelf").unwrap();
        let mut cache = HashCache::load(home);
        cache.sha256_of(&p).unwrap();
        std::fs::write(&p, b"modelshelf-changed").unwrap();
        assert_ne!(cache.sha256_of(&p).unwrap(), KNOWN);
    }

    #[test]
    fn unreadable_cache_is_discarded() {
        let tmp = tempfile::tempdir().unwrap();
        std::fs::write(tmp.path().join(HASH_CACHE_FILE), b"garbage!").unwrap();
        let cache = HashCache::load(tmp.path());
        assert!(cache.entries.is_empty());
    }

    #[cfg(unix)]
    #[test]
    fn same_inode_detects_hardlinks() {
        let tmp = tempfile::tempdir().unwrap();
        let a = tmp.path().join("a");
        let b = tmp.path().join("b");
        let c = tmp.path().join("c");
        std::fs::write(&a, b"x").unwrap();
        std::fs::hard_link(&a, &b).unwrap();
        std::fs::write(&c, b"x").unwrap();
        assert!(same_inode(&a, &b));
        assert!(!same_inode(&a, &c));
    }
}