DarkRoom/ui/dr-ui/src/develop/denoise_cache.rs

//! TRACES: FR-DEV-3g
//! The learned denoise's results, kept on disk (docs/dev/denoise.md §7.1).
//!
//! The network takes seconds per photograph and is on by default, so a
//! photograph reopened, or exported after it was developed, must not pay
//! again. A result is the network's output as it is — linear camera RGB at
//! the crop's size — written as half floats: about 120 MB for 20 MP, and no
//! compressor to link on Android. The strength slider is applied afterwards
//! and is not part of the key, so moving it never invalidates anything.
//!
//! **Keyed on the file's bytes and the model**: a SHA-256 of what was
//! decoded, and the model file's name and size. Anything that changes the
//! input or the network changes the key; the edit does not.
//!
//! **Bounded by a budget**, oldest first: a hit refreshes an entry's time, a
//! write evicts what no longer fits. Disposable — a peer of the inference
//! engine's cache, never synced — so an entry that cannot be read is simply
//! recomputed.

use std::io::{Read, Write};
use std::path::{Path, PathBuf};

use sha2::{Digest, Sha256};

/// Bytes the cache may hold before the oldest entries go: about forty 20 MP
/// photographs.
pub const BUDGET: u64 = 5 * 1024 * 1024 * 1024;

const MAGIC: &[u8; 8] = b"DRDN1\0\0\0";
const HEADER: usize = 8 + 4 + 4 + 1;

/// A cached result: the RGB samples, their size, and where the noise
/// figures came from.
pub struct Entry {
    pub rgb: Vec<f32>,
    pub width: u32,
    pub height: u32,
    pub source: dr_denoise::Source,
}

/// The cache's directory: beside the inference engine's, under the data
/// root, which is writable on every platform.
pub fn dir() -> PathBuf {
    crate::library::inference_cache_dir()
        .parent()
        .map(|p| p.join("denoise-cache"))
        .unwrap_or_else(|| PathBuf::from("denoise-cache"))
}

/// The key for a file's bytes under a model.
pub fn key(bytes: &[u8], model: &Path) -> String {
    let mut h = Sha256::new();
    h.update(bytes);
    if let Some(name) = model.file_name() {
        h.update(name.to_string_lossy().as_bytes());
    }
    let size = std::fs::metadata(model).map(|m| m.len()).unwrap_or(0);
    h.update(size.to_le_bytes());
    let digest = h.finalize();
    digest.iter().map(|b| format!("{b:02x}")).collect()
}

fn path_in(dir: &Path, key: &str) -> PathBuf {
    dir.join(format!("{key}.drdn"))
}

fn source_code(s: dr_denoise::Source) -> u8 {
    match s {
        dr_denoise::Source::Table => 0,
        dr_denoise::Source::DngProfile => 1,
        dr_denoise::Source::Measured => 2,
    }
}

fn source_from(code: u8) -> Option<dr_denoise::Source> {
    Some(match code {
        0 => dr_denoise::Source::Table,
        1 => dr_denoise::Source::DngProfile,
        2 => dr_denoise::Source::Measured,
        _ => return None,
    })
}

/// The entry for `key`, if one is held and reads back whole. A hit
/// refreshes its time, so what is in use outlives what is not.
pub fn load(dir: &Path, key: &str) -> Option<Entry> {
    let path = path_in(dir, key);
    let mut file = std::fs::File::open(&path).ok()?;
    let mut head = [0u8; HEADER];
    file.read_exact(&mut head).ok()?;
    if &head[..8] != MAGIC {
        return None;
    }
    let width = u32::from_le_bytes(head[8..12].try_into().ok()?);
    let height = u32::from_le_bytes(head[12..16].try_into().ok()?);
    let source = source_from(head[16])?;
    let samples = (width as usize)
        .checked_mul(height as usize)?
        .checked_mul(3)?;
    let mut raw = vec![0u8; samples.checked_mul(2)?];
    file.read_exact(&mut raw).ok()?;
    let rgb = raw
        .chunks_exact(2)
        .map(|b| half::f16::from_le_bytes([b[0], b[1]]).to_f32())
        .collect();
    let _ = file.set_modified(std::time::SystemTime::now());
    Some(Entry {
        rgb,
        width,
        height,
        source,
    })
}

/// Keep `entry` under `key`, then evict to `budget`. Written beside its name
/// and renamed, so a reader never sees half a file. Failure only costs a
/// recompute next time, so it is logged and swallowed.
pub fn store(dir: &Path, key: &str, entry: &Entry, budget: u64) {
    let result = (|| -> std::io::Result<()> {
        std::fs::create_dir_all(dir)?;
        let path = path_in(dir, key);
        let partial = path.with_extension("part");
        let mut out = std::io::BufWriter::new(std::fs::File::create(&partial)?);
        out.write_all(MAGIC)?;
        out.write_all(&entry.width.to_le_bytes())?;
        out.write_all(&entry.height.to_le_bytes())?;
        out.write_all(&[source_code(entry.source)])?;
        for v in &entry.rgb {
            out.write_all(&half::f16::from_f32(*v).to_le_bytes())?;
        }
        out.into_inner().map_err(|e| e.into_error())?.sync_all()?;
        std::fs::rename(&partial, &path)
    })();
    if let Err(e) = result {
        log::warn!("denoise cache: not kept ({e})");
        return;
    }
    evict(dir, budget);
}

/// Remove the oldest entries until what is left fits `budget`.
pub fn evict(dir: &Path, budget: u64) {
    let Ok(read) = std::fs::read_dir(dir) else {
        return;
    };
    let mut entries: Vec<(std::time::SystemTime, u64, PathBuf)> = read
        .flatten()
        .filter(|e| e.path().extension().is_some_and(|x| x == "drdn"))
        .filter_map(|e| {
            let m = e.metadata().ok()?;
            Some((m.modified().ok()?, m.len(), e.path()))
        })
        .collect();
    let mut total: u64 = entries.iter().map(|(_, len, _)| len).sum();
    entries.sort_by_key(|(t, _, _)| *t);
    for (_, len, path) in entries {
        if total <= budget {
            break;
        }
        if std::fs::remove_file(&path).is_ok() {
            total -= len;
        }
    }
}

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

    fn scratch(name: &str) -> PathBuf {
        let d =
            std::env::temp_dir().join(format!("dr-denoise-cache-{name}-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&d);
        d
    }

    fn entry(w: u32, h: u32) -> Entry {
        Entry {
            rgb: (0..w * h * 3).map(|i| i as f32 / 100.0).collect(),
            width: w,
            height: h,
            source: dr_denoise::Source::DngProfile,
        }
    }

    #[test]
    fn a_result_comes_back_as_it_went_in_to_half_precision() {
        let dir = scratch("roundtrip");
        let e = entry(4, 3);
        store(&dir, "k", &e, BUDGET);
        let back = load(&dir, "k").expect("a hit");
        assert_eq!((back.width, back.height), (4, 3));
        assert_eq!(back.source, dr_denoise::Source::DngProfile);
        for (a, b) in e.rgb.iter().zip(&back.rgb) {
            assert!((a - b).abs() <= a.abs() * 1e-3 + 1e-4, "{a} {b}");
        }
        assert!(load(&dir, "other").is_none());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn the_key_follows_the_bytes_and_the_model() {
        let dir = scratch("key");
        std::fs::create_dir_all(&dir).unwrap();
        let model = dir.join("m.onnx");
        std::fs::write(&model, b"weights").unwrap();
        let a = key(b"photo", &model);
        assert_eq!(a, key(b"photo", &model));
        assert_ne!(a, key(b"photo2", &model));
        std::fs::write(&model, b"other weights").unwrap();
        assert_ne!(a, key(b"photo", &model));
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn the_oldest_entries_go_first_past_the_budget() {
        let dir = scratch("evict");
        let e = entry(10, 10);
        store(&dir, "old", &e, BUDGET);
        let one = std::fs::metadata(path_in(&dir, "old")).unwrap().len();
        std::thread::sleep(std::time::Duration::from_millis(20));
        store(&dir, "mid", &e, BUDGET);
        std::thread::sleep(std::time::Duration::from_millis(20));
        // A hit makes the oldest the most recent.
        load(&dir, "old").unwrap();
        std::thread::sleep(std::time::Duration::from_millis(20));
        store(&dir, "new", &e, 2 * one);
        assert!(load(&dir, "mid").is_none(), "the least recently used went");
        assert!(load(&dir, "old").is_some() && load(&dir, "new").is_some());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn a_damaged_entry_is_a_miss() {
        let dir = scratch("damaged");
        store(&dir, "k", &entry(4, 4), BUDGET);
        let p = path_in(&dir, "k");
        let bytes = std::fs::read(&p).unwrap();
        std::fs::write(&p, &bytes[..bytes.len() / 2]).unwrap();
        assert!(load(&dir, "k").is_none());
        let _ = std::fs::remove_dir_all(&dir);
    }
}