1use crate::{
2 example::{Example, ExampleBuffer, ExampleState},
3 headless::EpAppState,
4};
5use anyhow::{Result, anyhow};
6use collections::HashMap;
7use edit_prediction::EditPredictionStore;
8use edit_prediction::udiff::OpenedBuffers;
9use futures::{
10 AsyncWriteExt as _,
11 lock::{Mutex, OwnedMutexGuard},
12};
13use gpui::{AsyncApp, Entity};
14use language::{Anchor, Buffer, LanguageNotFound, ToOffset, ToPoint};
15use project::buffer_store::BufferStoreEvent;
16use project::{Project, ProjectPath};
17use std::{
18 cell::RefCell,
19 fs,
20 path::{Path, PathBuf},
21 sync::Arc,
22};
23use util::{paths::PathStyle, rel_path::RelPath};
24use zeta_prompt::CURSOR_MARKER;
25
26pub async fn run_load_project(example: &mut Example, app_state: Arc<EpAppState>, mut cx: AsyncApp) {
27 if example.state.is_some() {
28 return;
29 }
30
31 let project = setup_project(example, &app_state, &mut cx).await;
32 let buffer_store = project
33 .read_with(&cx, |project, _| project.buffer_store().clone())
34 .unwrap();
35
36 let ep_store = cx
37 .update(|cx| EditPredictionStore::try_global(cx).unwrap())
38 .unwrap();
39
40 cx.subscribe(&buffer_store, {
41 let project = project.clone();
42 move |_, event, cx| match event {
43 BufferStoreEvent::BufferAdded(buffer) => {
44 ep_store.update(cx, |store, cx| store.register_buffer(&buffer, &project, cx));
45 }
46 _ => {}
47 }
48 })
49 .unwrap()
50 .detach();
51
52 let _open_buffers = apply_edit_history(example, &project, &mut cx)
53 .await
54 .unwrap();
55 let (buffer, cursor_position) = cursor_position(example, &project, &mut cx).await;
56 example.buffer = buffer
57 .read_with(&cx, |buffer, _cx| {
58 let cursor_point = cursor_position.to_point(&buffer);
59 Some(ExampleBuffer {
60 content: buffer.text(),
61 cursor_row: cursor_point.row,
62 cursor_column: cursor_point.column,
63 cursor_offset: cursor_position.to_offset(&buffer),
64 })
65 })
66 .unwrap();
67 example.state = Some(ExampleState {
68 buffer,
69 project,
70 cursor_position,
71 _open_buffers,
72 });
73}
74
75async fn cursor_position(
76 example: &Example,
77 project: &Entity<Project>,
78 cx: &mut AsyncApp,
79) -> (Entity<Buffer>, Anchor) {
80 let language_registry = project
81 .read_with(cx, |project, _| project.languages().clone())
82 .unwrap();
83 let result = language_registry
84 .load_language_for_file_path(&example.cursor_path)
85 .await;
86
87 if let Err(error) = result
88 && !error.is::<LanguageNotFound>()
89 {
90 panic!("Failed to load language for file path: {}", error);
91 }
92
93 let worktree = project
94 .read_with(cx, |project, cx| {
95 project.visible_worktrees(cx).next().unwrap()
96 })
97 .unwrap();
98
99 let cursor_path = RelPath::new(&example.cursor_path, PathStyle::Posix)
100 .unwrap()
101 .into_arc();
102 let cursor_buffer = project
103 .update(cx, |project, cx| {
104 project.open_buffer(
105 ProjectPath {
106 worktree_id: worktree.read(cx).id(),
107 path: cursor_path,
108 },
109 cx,
110 )
111 })
112 .unwrap()
113 .await
114 .unwrap();
115 let cursor_offset_within_excerpt = example
116 .cursor_position
117 .find(CURSOR_MARKER)
118 .ok_or_else(|| anyhow!("missing cursor marker"))
119 .unwrap();
120 let mut cursor_excerpt = example.cursor_position.clone();
121 cursor_excerpt.replace_range(
122 cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
123 "",
124 );
125 let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
126 let text = buffer.text();
127
128 let mut matches = text.match_indices(&cursor_excerpt);
129 let (excerpt_offset, _) = matches.next().unwrap_or_else(|| {
130 panic!(
131 "\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n.Example: {}\nCursor excerpt did not exist in buffer.",
132 example.name
133 );
134 });
135 assert!(matches.next().is_none(), "More than one cursor position match found for {}", &example.name);
136 excerpt_offset
137 }).unwrap();
138
139 let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
140 let cursor_anchor = cursor_buffer
141 .read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))
142 .unwrap();
143
144 (cursor_buffer, cursor_anchor)
145}
146
147async fn setup_project(
148 example: &mut Example,
149 app_state: &Arc<EpAppState>,
150 cx: &mut AsyncApp,
151) -> Entity<Project> {
152 setup_worktree(example).await;
153
154 let project = cx
155 .update(|cx| {
156 Project::local(
157 app_state.client.clone(),
158 app_state.node_runtime.clone(),
159 app_state.user_store.clone(),
160 app_state.languages.clone(),
161 app_state.fs.clone(),
162 None,
163 cx,
164 )
165 })
166 .unwrap();
167
168 project
169 .update(cx, |project, cx| {
170 project.disable_worktree_scanner(cx);
171 })
172 .unwrap();
173
174 let worktree = project
175 .update(cx, |project, cx| {
176 project.create_worktree(&example.worktree_path(), true, cx)
177 })
178 .unwrap()
179 .await
180 .unwrap();
181 worktree
182 .read_with(cx, |worktree, _cx| {
183 worktree.as_local().unwrap().scan_complete()
184 })
185 .unwrap()
186 .await;
187 project
188}
189
190pub async fn setup_worktree(example: &Example) {
191 let repo_dir = example.repo_path();
192 let repo_lock = lock_repo(&repo_dir).await;
193
194 if !repo_dir.is_dir() {
195 fs::create_dir_all(&repo_dir).unwrap();
196 run_git(&repo_dir, &["init"]).await.unwrap();
197 run_git(
198 &repo_dir,
199 &["remote", "add", "origin", &example.repository_url],
200 )
201 .await
202 .unwrap();
203 }
204
205 // Resolve the example to a revision, fetching it if needed.
206 let revision = run_git(
207 &repo_dir,
208 &["rev-parse", &format!("{}^{{commit}}", example.revision)],
209 )
210 .await;
211 let revision = if let Ok(revision) = revision {
212 revision
213 } else {
214 if run_git(
215 &repo_dir,
216 &["fetch", "--depth", "1", "origin", &example.revision],
217 )
218 .await
219 .is_err()
220 {
221 run_git(&repo_dir, &["fetch", "origin"]).await.unwrap();
222 }
223 let revision = run_git(&repo_dir, &["rev-parse", "FETCH_HEAD"])
224 .await
225 .unwrap();
226 revision
227 };
228
229 // Create the worktree for this example if needed.
230 let worktree_path = example.worktree_path();
231 if worktree_path.is_dir() {
232 run_git(&worktree_path, &["clean", "--force", "-d"])
233 .await
234 .unwrap();
235 run_git(&worktree_path, &["reset", "--hard", "HEAD"])
236 .await
237 .unwrap();
238 run_git(&worktree_path, &["checkout", revision.as_str()])
239 .await
240 .unwrap();
241 } else {
242 let worktree_path_string = worktree_path.to_string_lossy();
243 run_git(
244 &repo_dir,
245 &["branch", "-f", &example.name, revision.as_str()],
246 )
247 .await
248 .unwrap();
249 run_git(
250 &repo_dir,
251 &[
252 "worktree",
253 "add",
254 "-f",
255 &worktree_path_string,
256 &example.name,
257 ],
258 )
259 .await
260 .unwrap();
261 }
262 drop(repo_lock);
263
264 // Apply the uncommitted diff for this example.
265 if !example.uncommitted_diff.is_empty() {
266 let mut apply_process = smol::process::Command::new("git")
267 .current_dir(&worktree_path)
268 .args(&["apply", "-"])
269 .stdin(std::process::Stdio::piped())
270 .spawn()
271 .unwrap();
272
273 let mut stdin = apply_process.stdin.take().unwrap();
274 stdin
275 .write_all(example.uncommitted_diff.as_bytes())
276 .await
277 .unwrap();
278 stdin.close().await.unwrap();
279 drop(stdin);
280
281 let apply_result = apply_process.output().await.unwrap();
282 if !apply_result.status.success() {
283 panic!(
284 "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
285 apply_result.status,
286 String::from_utf8_lossy(&apply_result.stderr),
287 String::from_utf8_lossy(&apply_result.stdout),
288 );
289 }
290 }
291}
292
293async fn apply_edit_history(
294 example: &Example,
295 project: &Entity<Project>,
296 cx: &mut AsyncApp,
297) -> Result<OpenedBuffers> {
298 edit_prediction::udiff::apply_diff(&example.edit_history, project, cx).await
299}
300
301thread_local! {
302 static REPO_LOCKS: RefCell<HashMap<PathBuf, Arc<Mutex<()>>>> = RefCell::new(HashMap::default());
303}
304
305#[must_use]
306pub async fn lock_repo(path: impl AsRef<Path>) -> OwnedMutexGuard<()> {
307 REPO_LOCKS
308 .with(|cell| {
309 cell.borrow_mut()
310 .entry(path.as_ref().to_path_buf())
311 .or_default()
312 .clone()
313 })
314 .lock_owned()
315 .await
316}
317
318async fn run_git(repo_path: &Path, args: &[&str]) -> Result<String> {
319 let output = smol::process::Command::new("git")
320 .current_dir(repo_path)
321 .args(args)
322 .output()
323 .await?;
324
325 anyhow::ensure!(
326 output.status.success(),
327 "`git {}` within `{}` failed with status: {}\nstderr:\n{}\nstdout:\n{}",
328 args.join(" "),
329 repo_path.display(),
330 output.status,
331 String::from_utf8_lossy(&output.stderr),
332 String::from_utf8_lossy(&output.stdout),
333 );
334 Ok(String::from_utf8(output.stdout)?.trim().to_string())
335}