diff --git a/crates/edit_prediction/src/udiff.rs b/crates/edit_prediction/src/udiff.rs index da1b53875d..c4a9505de5 100644 --- a/crates/edit_prediction/src/udiff.rs +++ b/crates/edit_prediction/src/udiff.rs @@ -8,7 +8,7 @@ use std::{ }; use anyhow::{Context as _, Result, anyhow}; -use collections::HashMap; +use collections::{HashMap, hash_map::Entry}; use gpui::{AsyncApp, Entity}; use language::{Anchor, Buffer, OffsetRangeExt as _, TextBufferSnapshot, text_diff}; use postage::stream::Stream as _; @@ -17,7 +17,13 @@ use util::{paths::PathStyle, rel_path::RelPath}; use worktree::Worktree; #[derive(Clone, Debug)] -pub struct OpenedBuffers(#[allow(unused)] HashMap>); +pub struct OpenedBuffers(HashMap>); + +impl OpenedBuffers { + pub fn get(&self, path: &str) -> Option<&Entity> { + self.0.get(path) + } +} #[must_use] pub async fn apply_diff( @@ -42,12 +48,12 @@ pub async fn apply_diff( .collect(); refresh_worktree_entries(&worktree, paths.iter().map(|p| p.as_path()), cx).await?; - let mut included_files = HashMap::default(); + let mut included_files: HashMap> = HashMap::default(); let ranges = [Anchor::MIN..Anchor::MAX]; let mut diff = DiffParser::new(diff_str); let mut current_file = None; - let mut edits = vec![]; + let mut edits: Vec<(std::ops::Range, Arc)> = vec![]; while let Some(event) = diff.next()? { match event { @@ -58,21 +64,29 @@ pub async fn apply_diff( } => { let buffer = match current_file { None => { - let buffer = if is_new_file { - project - .update(cx, |project, cx| project.create_buffer(true, cx))? - .await? - } else { - let project_path = project - .update(cx, |project, cx| { - project.find_project_path(path.as_ref(), cx) - })? - .with_context(|| format!("no such path: {}", path))?; - project - .update(cx, |project, cx| project.open_buffer(project_path, cx))? - .await? + let buffer = match included_files.entry(path.to_string()) { + Entry::Occupied(entry) => entry.get().clone(), + Entry::Vacant(entry) => { + let buffer = if is_new_file { + project + .update(cx, |project, cx| project.create_buffer(true, cx))? + .await? + } else { + let project_path = project + .update(cx, |project, cx| { + project.find_project_path(path.as_ref(), cx) + })? + .with_context(|| format!("no such path: {}", path))?; + project + .update(cx, |project, cx| { + project.open_buffer(project_path, cx) + })? + .await? + }; + entry.insert(buffer.clone()); + buffer + } }; - included_files.insert(path.to_string(), buffer.clone()); current_file = Some(buffer); current_file.as_ref().unwrap() } diff --git a/crates/edit_prediction_cli/src/load_project.rs b/crates/edit_prediction_cli/src/load_project.rs index 83b300dc68..a3dc44a128 100644 --- a/crates/edit_prediction_cli/src/load_project.rs +++ b/crates/edit_prediction_cli/src/load_project.rs @@ -27,9 +27,10 @@ pub async fn run_load_project( let project = setup_project(example, &app_state, &progress, &mut cx).await?; - let _open_buffers = apply_edit_history(example, &project, &mut cx).await?; + let open_buffers = apply_edit_history(example, &project, &mut cx).await?; - let (buffer, cursor_position) = cursor_position(example, &project, &mut cx).await?; + let (buffer, cursor_position) = + cursor_position(example, &project, &open_buffers, &mut cx).await?; let (example_buffer, language_name) = buffer.read_with(&cx, |buffer, _cx| { let cursor_point = cursor_position.to_point(&buffer); let language_name = buffer @@ -54,7 +55,7 @@ pub async fn run_load_project( buffer, project, cursor_position, - _open_buffers, + _open_buffers: open_buffers, }); Ok(()) } @@ -62,6 +63,7 @@ pub async fn run_load_project( async fn cursor_position( example: &Example, project: &Entity, + open_buffers: &OpenedBuffers, cx: &mut AsyncApp, ) -> Result<(Entity, Anchor)> { let language_registry = project.read_with(cx, |project, _| project.languages().clone())?; @@ -75,25 +77,31 @@ async fn cursor_position( return Err(error); } - // Since the worktree scanner is disabled, manually refresh entries for the cursor path. - if let Some(worktree) = project.read_with(cx, |project, cx| project.worktrees(cx).next())? { - refresh_worktree_entries(&worktree, [&*example.spec.cursor_path], cx).await?; - } + let cursor_path_str = example.spec.cursor_path.to_string_lossy(); + // We try open_buffers first because the file might be new and not saved to disk + let cursor_buffer = if let Some(buffer) = open_buffers.get(&cursor_path_str) { + buffer.clone() + } else { + // Since the worktree scanner is disabled, manually refresh entries for the cursor path. + if let Some(worktree) = project.read_with(cx, |project, cx| project.worktrees(cx).next())? { + refresh_worktree_entries(&worktree, [&*example.spec.cursor_path], cx).await?; + } - let cursor_path = project - .read_with(cx, |project, cx| { - project.find_project_path(&example.spec.cursor_path, cx) - })? - .with_context(|| { - format!( - "failed to find cursor path {}", - example.spec.cursor_path.display() - ) - })?; + let cursor_path = project + .read_with(cx, |project, cx| { + project.find_project_path(&example.spec.cursor_path, cx) + })? + .with_context(|| { + format!( + "failed to find cursor path {}", + example.spec.cursor_path.display() + ) + })?; - let cursor_buffer = project - .update(cx, |project, cx| project.open_buffer(cursor_path, cx))? - .await?; + project + .update(cx, |project, cx| project.open_buffer(cursor_path, cx))? + .await? + }; let (cursor_excerpt, cursor_offset_within_excerpt) = example.spec.cursor_excerpt()?;