1use crate::{
2 example::{Example, ExampleBuffer, ExampleState},
3 git,
4 headless::EpAppState,
5 progress::{InfoStyle, Progress, Step, StepProgress},
6};
7use anyhow::{Context as _, Result};
8use edit_prediction::udiff::{OpenedBuffers, refresh_worktree_entries};
9use edit_prediction::{
10 EditPredictionStore, cursor_excerpt::editable_and_context_ranges_for_cursor_position, zeta2,
11};
12use futures::AsyncWriteExt as _;
13use gpui::{AsyncApp, Entity};
14use language::{Anchor, Buffer, LanguageNotFound, OffsetRangeExt as _, ToOffset, ToPoint};
15use project::Project;
16use project::buffer_store::BufferStoreEvent;
17use std::{fs, path::PathBuf, sync::Arc};
18
19pub async fn run_load_project(
20 example: &mut Example,
21 app_state: Arc<EpAppState>,
22 mut cx: AsyncApp,
23) -> Result<()> {
24 if example.state.is_some() {
25 return Ok(());
26 }
27
28 let progress = Progress::global().start(Step::LoadProject, &example.spec.name);
29
30 let project = setup_project(example, &app_state, &progress, &mut cx).await?;
31
32 progress.set_substatus("applying edit history");
33 let open_buffers = apply_edit_history(example, &project, &mut cx).await?;
34
35 progress.set_substatus("resolving cursor");
36 let (buffer, cursor_position) =
37 cursor_position(example, &project, &open_buffers, &mut cx).await?;
38 buffer
39 .read_with(&cx, |buffer, _| buffer.parsing_idle())?
40 .await;
41 let (example_buffer, language_name) = buffer.read_with(&cx, |buffer, _cx| {
42 let cursor_point = cursor_position.to_point(&buffer);
43 let snapshot = buffer.snapshot();
44 let (editable_range, context_range) = editable_and_context_ranges_for_cursor_position(
45 cursor_point,
46 &snapshot,
47 zeta2::MAX_REWRITE_TOKENS,
48 zeta2::MAX_CONTEXT_TOKENS,
49 );
50 let editable_range = editable_range.to_offset(&snapshot);
51 let context_range = context_range.to_offset(&snapshot);
52 let language_name = buffer
53 .language()
54 .map(|l| l.name().to_string())
55 .unwrap_or_else(|| "Unknown".to_string());
56 (
57 ExampleBuffer {
58 content: buffer.text(),
59 cursor_row: cursor_point.row,
60 cursor_column: cursor_point.column,
61 cursor_offset: cursor_position.to_offset(&buffer),
62 context_range,
63 editable_range,
64 },
65 language_name,
66 )
67 })?;
68
69 progress.set_info(language_name, InfoStyle::Normal);
70
71 example.buffer = Some(example_buffer);
72 example.state = Some(ExampleState {
73 buffer,
74 project,
75 cursor_position,
76 _open_buffers: open_buffers,
77 });
78 Ok(())
79}
80
81async fn cursor_position(
82 example: &Example,
83 project: &Entity<Project>,
84 open_buffers: &OpenedBuffers,
85 cx: &mut AsyncApp,
86) -> Result<(Entity<Buffer>, Anchor)> {
87 let language_registry = project.read_with(cx, |project, _| project.languages().clone())?;
88 let result = language_registry
89 .load_language_for_file_path(&example.spec.cursor_path)
90 .await;
91
92 if let Err(error) = result
93 && !error.is::<LanguageNotFound>()
94 {
95 return Err(error);
96 }
97
98 let cursor_path_str = example.spec.cursor_path.to_string_lossy();
99 // We try open_buffers first because the file might be new and not saved to disk
100 let cursor_buffer = if let Some(buffer) = open_buffers.get(&cursor_path_str) {
101 buffer.clone()
102 } else {
103 // Since the worktree scanner is disabled, manually refresh entries for the cursor path.
104 if let Some(worktree) = project.read_with(cx, |project, cx| project.worktrees(cx).next())? {
105 refresh_worktree_entries(&worktree, [&*example.spec.cursor_path], cx).await?;
106 }
107
108 let cursor_path = project
109 .read_with(cx, |project, cx| {
110 project.find_project_path(&example.spec.cursor_path, cx)
111 })?
112 .with_context(|| {
113 format!(
114 "failed to find cursor path {}",
115 example.spec.cursor_path.display()
116 )
117 })?;
118
119 project
120 .update(cx, |project, cx| project.open_buffer(cursor_path, cx))?
121 .await?
122 };
123
124 let (cursor_excerpt, cursor_offset_within_excerpt) = example.spec.cursor_excerpt()?;
125
126 let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
127 let text = buffer.text();
128
129 let mut matches = text.match_indices(&cursor_excerpt);
130 let (excerpt_offset, _) = matches.next().with_context(|| {
131 format!(
132 "\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n.Example: {}\nCursor excerpt did not exist in buffer.",
133 example.spec.name
134 )
135 })?;
136 anyhow::ensure!(
137 matches.next().is_none(),
138 "More than one cursor position match found for {}",
139 &example.spec.name
140 );
141 Ok(excerpt_offset)
142 })??;
143
144 let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
145 let cursor_anchor =
146 cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
147
148 Ok((cursor_buffer, cursor_anchor))
149}
150
151async fn setup_project(
152 example: &mut Example,
153 app_state: &Arc<EpAppState>,
154 step_progress: &StepProgress,
155 cx: &mut AsyncApp,
156) -> Result<Entity<Project>> {
157 let ep_store = cx
158 .update(|cx| EditPredictionStore::try_global(cx))?
159 .context("Store should be initialized at init")?;
160
161 let worktree_path = setup_worktree(example, step_progress).await?;
162
163 if let Some(project) = app_state.project_cache.get(&example.spec.repository_url) {
164 ep_store.update(cx, |ep_store, _| {
165 ep_store.clear_history_for_project(&project);
166 })?;
167 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone())?;
168 let buffers = buffer_store.read_with(cx, |buffer_store, _| {
169 buffer_store.buffers().collect::<Vec<_>>()
170 })?;
171 for buffer in buffers {
172 buffer
173 .update(cx, |buffer, cx| buffer.reload(cx))?
174 .await
175 .ok();
176 }
177 return Ok(project);
178 }
179
180 let project = cx.update(|cx| {
181 Project::local(
182 app_state.client.clone(),
183 app_state.node_runtime.clone(),
184 app_state.user_store.clone(),
185 app_state.languages.clone(),
186 app_state.fs.clone(),
187 None,
188 false,
189 cx,
190 )
191 })?;
192
193 project
194 .update(cx, |project, cx| {
195 project.disable_worktree_scanner(cx);
196 project.create_worktree(&worktree_path, true, cx)
197 })?
198 .await?;
199
200 app_state
201 .project_cache
202 .insert(example.spec.repository_url.clone(), project.clone());
203
204 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone())?;
205 cx.subscribe(&buffer_store, {
206 let project = project.clone();
207 move |_, event, cx| match event {
208 BufferStoreEvent::BufferAdded(buffer) => {
209 ep_store.update(cx, |store, cx| store.register_buffer(&buffer, &project, cx));
210 }
211 _ => {}
212 }
213 })?
214 .detach();
215
216 Ok(project)
217}
218
219async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Result<PathBuf> {
220 let repo_name = example.repo_name().context("failed to get repo name")?;
221 let repo_dir = git::repo_path_for_url(&example.spec.repository_url)?;
222 let worktree_path = repo_name.worktree_path();
223 let repo_lock = git::lock_repo(&repo_dir).await;
224
225 // Clean up any stale git lock files from previous crashed runs.
226 // Safe-ish since we have our own lock.
227 // WARNING: Can corrupt worktrees if multiple processes of the CLI are running.
228 let worktree_git_dir = repo_dir
229 .join(".git/worktrees")
230 .join(repo_name.name.as_ref());
231 let index_lock = worktree_git_dir.join("index.lock");
232 if index_lock.exists() {
233 fs::remove_file(&index_lock).ok();
234 }
235
236 if !repo_dir.is_dir() {
237 step_progress.set_substatus(format!("cloning {}", repo_name.name));
238 fs::create_dir_all(&repo_dir)?;
239 git::run_git(&repo_dir, &["init"]).await?;
240 git::run_git(
241 &repo_dir,
242 &["remote", "add", "origin", &example.spec.repository_url],
243 )
244 .await?;
245 }
246
247 // Resolve the example to a revision, fetching it if needed.
248 step_progress.set_substatus("fetching");
249 let revision = git::fetch_if_needed(&repo_dir, &example.spec.revision).await?;
250
251 // Create the worktree for this example if needed.
252 step_progress.set_substatus("preparing worktree");
253 if worktree_path.is_dir() {
254 git::run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
255 git::run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
256 git::run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
257 } else {
258 let worktree_path_string = worktree_path.to_string_lossy();
259 let branch_name = example.spec.filename();
260 git::run_git(
261 &repo_dir,
262 &["branch", "-f", &branch_name, revision.as_str()],
263 )
264 .await?;
265 git::run_git(
266 &repo_dir,
267 &["worktree", "add", "-f", &worktree_path_string, &branch_name],
268 )
269 .await?;
270 }
271 drop(repo_lock);
272
273 // Apply the uncommitted diff for this example.
274 if !example.spec.uncommitted_diff.is_empty() {
275 step_progress.set_substatus("applying diff");
276 let mut apply_process = smol::process::Command::new("git")
277 .current_dir(&worktree_path)
278 .args(&["apply", "-"])
279 .stdin(std::process::Stdio::piped())
280 .spawn()?;
281
282 let mut stdin = apply_process.stdin.take().context("Failed to get stdin")?;
283 stdin
284 .write_all(example.spec.uncommitted_diff.as_bytes())
285 .await?;
286 stdin.close().await?;
287 drop(stdin);
288
289 let apply_result = apply_process.output().await?;
290 anyhow::ensure!(
291 apply_result.status.success(),
292 "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
293 apply_result.status,
294 String::from_utf8_lossy(&apply_result.stderr),
295 String::from_utf8_lossy(&apply_result.stdout),
296 );
297 }
298
299 step_progress.clear_substatus();
300 Ok(worktree_path)
301}
302
303async fn apply_edit_history(
304 example: &Example,
305 project: &Entity<Project>,
306 cx: &mut AsyncApp,
307) -> Result<OpenedBuffers> {
308 edit_prediction::udiff::apply_diff(&example.spec.edit_history, project, cx).await
309}