use crate::{
diff_review::{ReviewLineKind, ReviewSnapshot},
output::sanitize_display_text,
rendering::{
DisplayLine,
highlight::{SourceHighlightBudget, SourceHighlightJob},
},
};
use std::{collections::BTreeMap, sync::Arc};
pub(super) struct FileProjection {
pub original: Arc<Vec<DisplayLine>>,
pub current: Arc<Vec<DisplayLine>>,
pub removed: Vec<bool>,
pending: bool,
}
pub(super) struct TreeEntry {
pub identity: String,
pub label: String,
pub file: Option<usize>,
pub comment: Option<usize>,
}
pub(crate) struct PreparedSnapshot {
source: ReviewSnapshot,
pub(super) files_display: Vec<Arc<FileProjection>>,
pub(super) tree: Vec<TreeEntry>,
pub(super) highlighting_pending: bool,
}
impl std::ops::Deref for PreparedSnapshot {
type Target = ReviewSnapshot;
fn deref(&self) -> &Self::Target {
&self.source
}
}
#[derive(Default)]
pub(super) struct ProjectionCache {
root: std::path::PathBuf,
files: BTreeMap<String, [SourceHighlightJob; 2]>,
current_first: bool,
}
impl PreparedSnapshot {
#[cfg(test)]
pub(super) fn new(source: ReviewSnapshot, previous: Option<&Self>) -> Self {
Self::prepare(
source,
previous,
None,
&crate::cancellation::AgentCancellation::default(),
&mut ProjectionCache::default(),
)
.unwrap()
}
pub(super) fn prepare(
source: ReviewSnapshot,
previous: Option<&Self>,
selected_path: Option<&str>,
cancellation: &crate::cancellation::AgentCancellation,
cache: &mut ProjectionCache,
) -> anyhow::Result<Self> {
Self::with_budget(
source,
previous,
selected_path,
&SourceHighlightBudget::new(cancellation),
cache,
usize::MAX,
)
}
pub(super) fn with_budget(
source: ReviewSnapshot,
previous: Option<&Self>,
selected_path: Option<&str>,
budget: &SourceHighlightBudget<'_>,
cache: &mut ProjectionCache,
max_lines: usize,
) -> anyhow::Result<Self> {
budget.cancellation.check()?;
if cache.root != source.root {
cache.files.clear();
cache.root = source.root.clone();
}
let paths: std::collections::BTreeSet<_> =
source.files.iter().map(|file| file.path.as_str()).collect();
cache.files.retain(|path, _| paths.contains(path.as_str()));
let selected = selected_path
.and_then(|path| source.files.iter().position(|file| file.path == path))
.unwrap_or(0);
let previous_files: BTreeMap<_, _> = previous
.filter(|old| old.root == source.root)
.into_iter()
.flat_map(|old| old.files.iter().zip(&old.files_display))
.map(|(file, display)| (file.path.as_str(), (file, display)))
.collect();
let mut files_display = vec![None; source.files.len()];
let order = (selected..source.files.len().min(selected + 1))
.chain((0..source.files.len()).filter(|index| *index != selected));
for index in order {
budget.cancellation.check()?;
let file = &source.files[index];
if let Some((old, display)) = previous_files.get(file.path.as_str())
&& !display.pending
&& old.original == file.original
&& old.current == file.current
&& old.rows.len() == file.rows.len()
&& old.rows.iter().zip(&file.rows).all(|(a, b)| {
a.old_line == b.old_line
&& a.new_line == b.new_line
&& a.kind == b.kind
&& a.text == b.text
})
{
files_display[index] = Some(Arc::clone(display));
continue;
}
let language = std::path::Path::new(&file.path)
.extension()
.and_then(|s| s.to_str());
let sources = [
sanitize_display_text(&file.original),
sanitize_display_text(&file.current),
];
let jobs = cache.files.entry(file.path.clone()).or_insert_with(|| {
[
SourceHighlightJob::new(sources[0].clone(), language),
SourceHighlightJob::new(sources[1].clone(), language),
]
});
for side in 0..2 {
if !jobs[side].matches(&sources[side]) {
jobs[side] = SourceHighlightJob::new(sources[side].clone(), language);
}
}
let first = usize::from(cache.current_first);
for side in [first, 1 - first] {
let slice = SourceHighlightBudget {
deadline: budget
.deadline
.min(std::time::Instant::now() + std::time::Duration::from_millis(150)),
cancellation: budget.cancellation,
};
jobs[side].advance(&slice, max_lines)?;
}
let original = jobs[0].display();
let current = jobs[1].display();
let mut removed = vec![false; original.len()];
for row in &file.rows {
if row.kind == ReviewLineKind::Removed
&& let Some(index) = row.old_line.and_then(|line| line.checked_sub(1))
&& let Some(removed) = removed.get_mut(index)
{
*removed = true;
}
}
files_display[index] = Some(Arc::new(FileProjection {
original,
current,
removed,
pending: jobs.iter().any(SourceHighlightJob::pending),
}));
}
cache.current_first = !cache.current_first;
budget.cancellation.check()?;
let tree = build_tree(&source);
let files_display: Vec<_> = files_display
.into_iter()
.map(|file| file.expect("every file projected"))
.collect();
let highlighting_pending = files_display.iter().any(|file| file.pending);
Ok(Self {
source,
files_display,
tree,
highlighting_pending,
})
}
}
#[derive(Default)]
struct Folder<'a> {
folders: BTreeMap<&'a str, Folder<'a>>,
files: BTreeMap<&'a str, usize>,
}
fn build_tree(snapshot: &ReviewSnapshot) -> Vec<TreeEntry> {
let mut root = Folder::default();
for (index, file) in snapshot.files.iter().enumerate() {
let mut folder = &mut root;
let mut parts = file.path.split('/').peekable();
while let Some(part) = parts.next() {
if parts.peek().is_some() {
folder = folder.folders.entry(part).or_default();
} else {
folder.files.insert(part, index);
}
}
}
let mut entries = Vec::new();
append_folder(&root, "", 0, snapshot, &mut entries);
entries
}
fn append_folder(
folder: &Folder<'_>,
path: &str,
depth: usize,
snapshot: &ReviewSnapshot,
entries: &mut Vec<TreeEntry>,
) {
for (name, child) in &folder.folders {
let path = format!("{path}{name}/");
entries.push(TreeEntry {
identity: format!("folder:{path}"),
label: sanitize_display_text(&format!("{}{name}/", " ".repeat(depth))),
file: None,
comment: None,
});
append_folder(child, &path, depth + 1, snapshot, entries);
}
for (name, index) in &folder.files {
let file = &snapshot.files[*index];
entries.push(TreeEntry {
identity: format!("file:{}", file.path),
label: sanitize_display_text(&format!("{}{name}", " ".repeat(depth))),
file: Some(*index),
comment: None,
});
for (comment_index, comment) in snapshot
.comments
.iter()
.enumerate()
.filter(|(_, c)| c.path == file.path)
{
entries.push(TreeEntry {
identity: format!("comment:{}", comment.id),
label: sanitize_display_text(&format!(
"{}{} {}:{} {}",
" ".repeat(depth + 1),
super::render::comment_marker(comment.stale, comment.resolved),
if comment.side == crate::diff_review::ReviewSide::Original {
"O"
} else {
"D"
},
comment.line,
comment.text.lines().next().unwrap_or("")
)),
file: Some(*index),
comment: Some(comment_index),
});
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{cancellation::AgentCancellation, diff_review::ReviewFile, rendering::DisplayRole};
use std::time::Instant;
fn snapshot() -> ReviewSnapshot {
ReviewSnapshot {
root: "/review".into(),
files: vec![ReviewFile {
path: "file.rs".into(),
original: "/* open\nstill comment\n*/\nfn original() {}".into(),
current: "/* open\nstill comment\n*/\nfn current() {}".into(),
rows: vec![],
notice: None,
}],
comments: vec![],
}
}
#[test]
fn both_sides_resume_on_unchanged_selection_and_keep_multiline_state() {
crate::rendering::highlight::prewarm();
let token = AgentCancellation::default();
let mut cache = ProjectionCache::default();
let mut previous = None;
for step in 0..4 {
let start = Instant::now();
let prepared = PreparedSnapshot::with_budget(
snapshot(),
previous.as_ref(),
Some("file.rs"),
&SourceHighlightBudget::new(&token),
&mut cache,
1,
)
.unwrap();
eprintln!("highlight continuation {step}: {:?}", start.elapsed());
for side in [
&prepared.files_display[0].original,
&prepared.files_display[0].current,
] {
assert_eq!(side.len(), 4);
assert!(
side[step]
.spans
.iter()
.any(|span| span.role != DisplayRole::FallbackCode)
);
if step == 1 {
assert!(
side[1]
.spans
.iter()
.all(|span| span.role == DisplayRole::Comment)
);
}
if step < 3 {
assert_eq!(side[step + 1].spans[0].role, DisplayRole::FallbackCode);
}
}
assert_eq!(prepared.highlighting_pending, step < 3);
previous = Some(prepared);
}
let first = previous.unwrap();
let mut source = snapshot();
source.files[0].current = "fn changed() {}".into();
let changed =
PreparedSnapshot::prepare(source, Some(&first), None, &token, &mut cache).unwrap();
assert!(Arc::ptr_eq(
&first.files_display[0].original,
&changed.files_display[0].original
));
assert!(
changed.files_display[0].current[0]
.spans
.iter()
.any(|s| s.text == "changed" && s.role == DisplayRole::Function)
);
}
#[test]
fn expired_budget_resumes_without_selection_change_and_cancellation_is_prompt() {
let (token, cancel) = AgentCancellation::default().child_token();
let mut cache = ProjectionCache::default();
let expired = SourceHighlightBudget {
deadline: Instant::now(),
cancellation: &token,
};
let first =
PreparedSnapshot::with_budget(snapshot(), None, None, &expired, &mut cache, 1).unwrap();
assert!(first.highlighting_pending);
let next =
PreparedSnapshot::prepare(snapshot(), Some(&first), None, &token, &mut cache).unwrap();
assert!(!next.highlighting_pending);
cancel.cancel();
let start = Instant::now();
assert!(
PreparedSnapshot::prepare(snapshot(), Some(&next), None, &token, &mut cache).is_err()
);
eprintln!("canceled projection: {:?}", start.elapsed());
}
}