use crate::diff::model::Changeset;
use crate::ui::highlight::Highlighter;
use crate::ui::render_rows::sanitize_line;
use crate::ui::theme::theme;
use ratatui::style::Color;
use std::collections::HashMap;
use std::sync::mpsc::{self, Sender};
use std::sync::{Arc, Mutex, MutexGuard};
pub type LineRuns = Arc<Vec<(Color, String)>>;
type FileCache = HashMap<String, LineRuns>;
type SharedCache = Arc<Mutex<HashMap<usize, FileCache>>>;
const PREFETCH_RADIUS: usize = 2;
const KEEP_FILES: usize = 2 * PREFETCH_RADIUS + 1;
fn prefetch_order(focus: usize, n: usize) -> Vec<usize> {
if focus >= n {
return Vec::new();
}
let mut order = Vec::with_capacity(KEEP_FILES);
order.push(focus);
for d in 1..=PREFETCH_RADIUS {
if focus.checked_add(d).is_some_and(|up| up < n) {
order.push(focus + d);
}
if focus >= d {
order.push(focus - d);
}
}
order
}
fn lock_cache(
cache: &Mutex<HashMap<usize, FileCache>>,
) -> MutexGuard<'_, HashMap<usize, FileCache>> {
cache.lock().unwrap_or_else(|e| e.into_inner())
}
fn spawn_warm_worker(cache: SharedCache, changeset: Arc<Changeset>) -> Sender<usize> {
let (tx, rx) = mpsc::channel::<usize>();
let n = changeset.files.len();
std::thread::spawn(move || {
let hl = Highlighter::new();
let mut pending: Option<usize> = None;
'jobs: loop {
let mut focus = match pending.take() {
Some(i) => i,
None => match rx.recv() {
Ok(i) => i,
Err(_) => return, },
};
while let Ok(newer) = rx.try_recv() {
focus = newer;
}
for file_idx in prefetch_order(focus, n) {
let Some(file) = changeset.files.get(file_idx) else {
continue;
};
let syntax = hl.syntax_for(file.display_path());
for line in file.hunks.iter().flat_map(|h| h.lines.iter()) {
match rx.try_recv() {
Ok(newer) => {
pending = Some(newer);
continue 'jobs;
}
Err(mpsc::TryRecvError::Disconnected) => return,
Err(mpsc::TryRecvError::Empty) => {}
}
let text = sanitize_line(&line.text);
if lock_cache(&cache)
.get(&file_idx)
.is_some_and(|m| m.contains_key(&text))
{
continue;
}
let runs = Arc::new(hl.line(syntax, &text));
lock_cache(&cache)
.entry(file_idx)
.or_default()
.entry(text)
.or_insert(runs);
}
}
}
});
tx
}
pub struct HighlightCache {
changeset: Arc<Changeset>,
highlighter: Highlighter,
cache: SharedCache,
tx: Sender<usize>,
warmed_file: Option<usize>,
}
impl HighlightCache {
pub fn new(changeset: Arc<Changeset>) -> Self {
let cache: SharedCache = Arc::new(Mutex::new(HashMap::new()));
let tx = spawn_warm_worker(cache.clone(), changeset.clone());
HighlightCache {
changeset,
highlighter: Highlighter::new(),
cache,
tx,
warmed_file: None,
}
}
pub fn runs(&self, file_idx: usize, text: &str) -> LineRuns {
if let Some(v) = lock_cache(&self.cache)
.get(&file_idx)
.and_then(|m| m.get(text))
{
return v.clone();
}
let spans = match self.changeset.files.get(file_idx) {
Some(f) => {
let syntax = self.highlighter.syntax_for(f.display_path());
self.highlighter.line(syntax, text)
}
None => vec![(theme().text, text.to_string())],
};
let rc = Arc::new(spans);
lock_cache(&self.cache)
.entry(file_idx)
.or_default()
.entry(text.to_string())
.or_insert(rc)
.clone()
}
pub fn warm(&mut self, file_idx: usize) {
if self.warmed_file == Some(file_idx) {
return;
}
self.warmed_file = Some(file_idx);
self.evict_outside_prefetch_window(file_idx);
let _ = self.tx.send(file_idx);
}
fn evict_outside_prefetch_window(&mut self, file_idx: usize) {
let keep = prefetch_order(file_idx, self.changeset.files.len());
lock_cache(&self.cache).retain(|fi, _| keep.contains(fi));
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::diff::model::{Changeset, DiffFile, DiffLine, Hunk, LineKind};
fn file_with(lines: &[&str]) -> DiffFile {
DiffFile {
old_path: "a.rs".into(),
new_path: "a.rs".into(),
is_binary: false,
hunks: vec![Hunk {
old_start: 1,
old_count: lines.len() as u32,
new_start: 1,
new_count: lines.len() as u32,
section: None,
lines: lines
.iter()
.map(|t| DiffLine {
kind: LineKind::Context,
old_line: Some(1),
new_line: Some(1),
text: (*t).to_string(),
})
.collect(),
}],
}
}
fn cs(n_files: usize) -> Arc<Changeset> {
Arc::new(Changeset {
files: (0..n_files).map(|_| file_with(&["let x = 1;"])).collect(),
})
}
#[test]
fn runs_are_cached_and_stable() {
let hc = HighlightCache::new(cs(1));
let a = hc.runs(0, "let x = 1;");
let b = hc.runs(0, "let x = 1;");
assert!(Arc::ptr_eq(&a, &b));
assert!(!a.is_empty());
}
#[test]
fn missing_file_falls_back_to_plain() {
let hc = HighlightCache::new(cs(1));
let runs = hc.runs(99, "anything");
assert_eq!(runs.len(), 1);
assert_eq!(runs[0].1, "anything");
}
#[test]
fn prefetch_order_is_focus_then_outward_neighbors() {
assert_eq!(prefetch_order(5, 10), vec![5, 6, 4, 7, 3]);
assert_eq!(prefetch_order(0, 10), vec![0, 1, 2]);
assert_eq!(prefetch_order(9, 10), vec![9, 8, 7]);
assert!(prefetch_order(9, 0).is_empty());
assert!(prefetch_order(10, 10).is_empty());
assert!(prefetch_order(15, 10).is_empty());
assert!(prefetch_order(usize::MAX, 10).is_empty());
}
#[test]
fn cache_is_bounded_by_prefetch_window() {
let mut hc = HighlightCache::new(cs(10));
for fi in 0..10 {
hc.runs(fi, "let x = 1;");
hc.warm(fi);
}
let len = lock_cache(&hc.cache).len();
assert!(len <= KEEP_FILES, "cache should be bounded, got {len}");
}
}