1use std::{
2 borrow::Cow,
3 cell::RefCell,
4 fmt::{self, Display},
5 fs,
6 io::Write,
7 mem,
8 path::{Path, PathBuf},
9 sync::Arc,
10};
11
12use anyhow::{Context as _, Result, anyhow};
13use clap::ValueEnum;
14use cloud_zeta2_prompt::CURSOR_MARKER;
15use collections::HashMap;
16use futures::{
17 AsyncWriteExt as _,
18 lock::{Mutex, OwnedMutexGuard},
19};
20use gpui::{AsyncApp, Entity, http_client::Url};
21use language::{Anchor, Buffer};
22use project::{Project, ProjectPath};
23use pulldown_cmark::CowStr;
24use serde::{Deserialize, Serialize};
25use util::{paths::PathStyle, rel_path::RelPath};
26use zeta2::udiff::OpenedBuffers;
27
28use crate::paths::{REPOS_DIR, WORKTREES_DIR};
29
30const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
31const EDIT_HISTORY_HEADING: &str = "Edit History";
32const CURSOR_POSITION_HEADING: &str = "Cursor Position";
33const EXPECTED_PATCH_HEADING: &str = "Expected Patch";
34const EXPECTED_EXCERPTS_HEADING: &str = "Expected Excerpts";
35const REPOSITORY_URL_FIELD: &str = "repository_url";
36const REVISION_FIELD: &str = "revision";
37
38#[derive(Debug, Clone)]
39pub struct NamedExample {
40 pub name: String,
41 pub example: Example,
42}
43
44#[derive(Clone, Debug, Serialize, Deserialize)]
45pub struct Example {
46 pub repository_url: String,
47 pub revision: String,
48 pub uncommitted_diff: String,
49 pub cursor_path: PathBuf,
50 pub cursor_position: String,
51 pub edit_history: String,
52 pub expected_patch: String,
53 pub expected_excerpts: Vec<ExpectedExcerpt>,
54}
55
56pub type ExpectedExcerpt = Excerpt;
57pub type ActualExcerpt = Excerpt;
58
59#[derive(Clone, Debug, Serialize, Deserialize)]
60pub struct Excerpt {
61 pub path: PathBuf,
62 pub text: String,
63}
64
65#[derive(ValueEnum, Debug, Clone)]
66pub enum ExampleFormat {
67 Json,
68 Toml,
69 Md,
70}
71
72impl NamedExample {
73 pub fn load(path: impl AsRef<Path>) -> Result<Self> {
74 let path = path.as_ref();
75 let content = std::fs::read_to_string(path)?;
76 let ext = path.extension();
77
78 match ext.and_then(|s| s.to_str()) {
79 Some("json") => Ok(Self {
80 name: path.file_stem().unwrap_or_default().display().to_string(),
81 example: serde_json::from_str(&content)?,
82 }),
83 Some("toml") => Ok(Self {
84 name: path.file_stem().unwrap_or_default().display().to_string(),
85 example: toml::from_str(&content)?,
86 }),
87 Some("md") => Self::parse_md(&content),
88 Some(_) => {
89 anyhow::bail!("Unrecognized example extension: {}", ext.unwrap().display());
90 }
91 None => {
92 anyhow::bail!(
93 "Failed to determine example type since the file does not have an extension."
94 );
95 }
96 }
97 }
98
99 pub fn parse_md(input: &str) -> Result<Self> {
100 use pulldown_cmark::{CodeBlockKind, Event, HeadingLevel, Parser, Tag, TagEnd};
101
102 let parser = Parser::new(input);
103
104 let mut named = NamedExample {
105 name: String::new(),
106 example: Example {
107 repository_url: String::new(),
108 revision: String::new(),
109 uncommitted_diff: String::new(),
110 cursor_path: PathBuf::new(),
111 cursor_position: String::new(),
112 edit_history: String::new(),
113 expected_patch: String::new(),
114 expected_excerpts: Vec::new(),
115 },
116 };
117
118 let mut text = String::new();
119 let mut current_section = String::new();
120 let mut block_info: CowStr = "".into();
121
122 for event in parser {
123 match event {
124 Event::Text(line) => {
125 text.push_str(&line);
126
127 if !named.name.is_empty()
128 && current_section.is_empty()
129 // in h1 section
130 && let Some((field, value)) = line.split_once('=')
131 {
132 match field.trim() {
133 REPOSITORY_URL_FIELD => {
134 named.example.repository_url = value.trim().to_string();
135 }
136 REVISION_FIELD => {
137 named.example.revision = value.trim().to_string();
138 }
139 _ => {
140 eprintln!("Warning: Unrecognized field `{field}`");
141 }
142 }
143 }
144 }
145 Event::End(TagEnd::Heading(HeadingLevel::H1)) => {
146 if !named.name.is_empty() {
147 anyhow::bail!(
148 "Found multiple H1 headings. There should only be one with the name of the example."
149 );
150 }
151 named.name = mem::take(&mut text);
152 }
153 Event::End(TagEnd::Heading(HeadingLevel::H2)) => {
154 current_section = mem::take(&mut text);
155 }
156 Event::End(TagEnd::Heading(level)) => {
157 anyhow::bail!("Unexpected heading level: {level}");
158 }
159 Event::Start(Tag::CodeBlock(kind)) => {
160 match kind {
161 CodeBlockKind::Fenced(info) => {
162 block_info = info;
163 }
164 CodeBlockKind::Indented => {
165 anyhow::bail!("Unexpected indented codeblock");
166 }
167 };
168 }
169 Event::Start(_) => {
170 text.clear();
171 block_info = "".into();
172 }
173 Event::End(TagEnd::CodeBlock) => {
174 let block_info = block_info.trim();
175 if current_section.eq_ignore_ascii_case(UNCOMMITTED_DIFF_HEADING) {
176 named.example.uncommitted_diff = mem::take(&mut text);
177 } else if current_section.eq_ignore_ascii_case(EDIT_HISTORY_HEADING) {
178 named.example.edit_history.push_str(&mem::take(&mut text));
179 } else if current_section.eq_ignore_ascii_case(CURSOR_POSITION_HEADING) {
180 named.example.cursor_path = block_info.into();
181 named.example.cursor_position = mem::take(&mut text);
182 } else if current_section.eq_ignore_ascii_case(EXPECTED_PATCH_HEADING) {
183 named.example.expected_patch = mem::take(&mut text);
184 } else if current_section.eq_ignore_ascii_case(EXPECTED_EXCERPTS_HEADING) {
185 // TODO: "…" should not be a part of the excerpt
186 named.example.expected_excerpts.push(ExpectedExcerpt {
187 path: block_info.into(),
188 text: mem::take(&mut text),
189 });
190 } else {
191 eprintln!("Warning: Unrecognized section `{current_section:?}`")
192 }
193 }
194 _ => {}
195 }
196 }
197
198 if named.example.cursor_path.as_path() == Path::new("")
199 || named.example.cursor_position.is_empty()
200 {
201 anyhow::bail!("Missing cursor position codeblock");
202 }
203
204 Ok(named)
205 }
206
207 pub fn write(&self, format: ExampleFormat, mut out: impl Write) -> Result<()> {
208 match format {
209 ExampleFormat::Json => Ok(serde_json::to_writer(out, &self.example)?),
210 ExampleFormat::Toml => {
211 Ok(out.write_all(toml::to_string_pretty(&self.example)?.as_bytes())?)
212 }
213 ExampleFormat::Md => Ok(write!(out, "{}", self)?),
214 }
215 }
216
217 pub async fn setup_worktree(&self) -> Result<PathBuf> {
218 let (repo_owner, repo_name) = self.repo_name()?;
219 let file_name = self.file_name();
220
221 fs::create_dir_all(&*REPOS_DIR)?;
222 fs::create_dir_all(&*WORKTREES_DIR)?;
223
224 let repo_dir = REPOS_DIR.join(repo_owner.as_ref()).join(repo_name.as_ref());
225 let repo_lock = lock_repo(&repo_dir).await;
226
227 if !repo_dir.is_dir() {
228 fs::create_dir_all(&repo_dir)?;
229 run_git(&repo_dir, &["init"]).await?;
230 run_git(
231 &repo_dir,
232 &["remote", "add", "origin", &self.example.repository_url],
233 )
234 .await?;
235 }
236
237 // Resolve the example to a revision, fetching it if needed.
238 let revision = run_git(&repo_dir, &["rev-parse", &self.example.revision]).await;
239 let revision = if let Ok(revision) = revision {
240 revision
241 } else {
242 run_git(
243 &repo_dir,
244 &["fetch", "--depth", "1", "origin", &self.example.revision],
245 )
246 .await?;
247 let revision = run_git(&repo_dir, &["rev-parse", "FETCH_HEAD"]).await?;
248 if revision != self.example.revision {
249 run_git(&repo_dir, &["tag", &self.example.revision, &revision]).await?;
250 }
251 revision
252 };
253
254 // Create the worktree for this example if needed.
255 let worktree_path = WORKTREES_DIR.join(&file_name);
256 if worktree_path.is_dir() {
257 run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
258 run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
259 run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
260 } else {
261 let worktree_path_string = worktree_path.to_string_lossy();
262 run_git(&repo_dir, &["branch", "-f", &file_name, revision.as_str()]).await?;
263 run_git(
264 &repo_dir,
265 &["worktree", "add", "-f", &worktree_path_string, &file_name],
266 )
267 .await?;
268 }
269 drop(repo_lock);
270
271 // Apply the uncommitted diff for this example.
272 if !self.example.uncommitted_diff.is_empty() {
273 let mut apply_process = smol::process::Command::new("git")
274 .current_dir(&worktree_path)
275 .args(&["apply", "-"])
276 .stdin(std::process::Stdio::piped())
277 .spawn()?;
278
279 let mut stdin = apply_process.stdin.take().unwrap();
280 stdin
281 .write_all(self.example.uncommitted_diff.as_bytes())
282 .await?;
283 stdin.close().await?;
284 drop(stdin);
285
286 let apply_result = apply_process.output().await?;
287 if !apply_result.status.success() {
288 anyhow::bail!(
289 "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
290 apply_result.status,
291 String::from_utf8_lossy(&apply_result.stderr),
292 String::from_utf8_lossy(&apply_result.stdout),
293 );
294 }
295 }
296
297 Ok(worktree_path)
298 }
299
300 fn file_name(&self) -> String {
301 self.name
302 .chars()
303 .map(|c| {
304 if c.is_whitespace() {
305 '-'
306 } else {
307 c.to_ascii_lowercase()
308 }
309 })
310 .collect()
311 }
312
313 fn repo_name(&self) -> Result<(Cow<'_, str>, Cow<'_, str>)> {
314 // git@github.com:owner/repo.git
315 if self.example.repository_url.contains('@') {
316 let (owner, repo) = self
317 .example
318 .repository_url
319 .split_once(':')
320 .context("expected : in git url")?
321 .1
322 .split_once('/')
323 .context("expected / in git url")?;
324 Ok((
325 Cow::Borrowed(owner),
326 Cow::Borrowed(repo.trim_end_matches(".git")),
327 ))
328 // http://github.com/owner/repo.git
329 } else {
330 let url = Url::parse(&self.example.repository_url)?;
331 let mut segments = url.path_segments().context("empty http url")?;
332 let owner = segments
333 .next()
334 .context("expected owner path segment")?
335 .to_string();
336 let repo = segments
337 .next()
338 .context("expected repo path segment")?
339 .trim_end_matches(".git")
340 .to_string();
341 assert!(segments.next().is_none());
342
343 Ok((owner.into(), repo.into()))
344 }
345 }
346
347 pub async fn cursor_position(
348 &self,
349 project: &Entity<Project>,
350 cx: &mut AsyncApp,
351 ) -> Result<(Entity<Buffer>, Anchor)> {
352 let worktree = project.read_with(cx, |project, cx| {
353 project.visible_worktrees(cx).next().unwrap()
354 })?;
355 let cursor_path = RelPath::new(&self.example.cursor_path, PathStyle::Posix)?.into_arc();
356 let cursor_buffer = project
357 .update(cx, |project, cx| {
358 project.open_buffer(
359 ProjectPath {
360 worktree_id: worktree.read(cx).id(),
361 path: cursor_path,
362 },
363 cx,
364 )
365 })?
366 .await?;
367 let cursor_offset_within_excerpt = self
368 .example
369 .cursor_position
370 .find(CURSOR_MARKER)
371 .ok_or_else(|| anyhow!("missing cursor marker"))?;
372 let mut cursor_excerpt = self.example.cursor_position.clone();
373 cursor_excerpt.replace_range(
374 cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
375 "",
376 );
377 let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
378 let text = buffer.text();
379
380 let mut matches = text.match_indices(&cursor_excerpt);
381 let Some((excerpt_offset, _)) = matches.next() else {
382 anyhow::bail!(
383 "Cursor excerpt did not exist in buffer.\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n"
384 );
385 };
386 assert!(matches.next().is_none());
387
388 Ok(excerpt_offset)
389 })??;
390
391 let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
392 let cursor_anchor =
393 cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
394 Ok((cursor_buffer, cursor_anchor))
395 }
396
397 #[must_use]
398 pub async fn apply_edit_history(
399 &self,
400 project: &Entity<Project>,
401 cx: &mut AsyncApp,
402 ) -> Result<OpenedBuffers<'_>> {
403 zeta2::udiff::apply_diff(&self.example.edit_history, project, cx).await
404 }
405}
406
407async fn run_git(repo_path: &Path, args: &[&str]) -> Result<String> {
408 let output = smol::process::Command::new("git")
409 .current_dir(repo_path)
410 .args(args)
411 .output()
412 .await?;
413
414 anyhow::ensure!(
415 output.status.success(),
416 "`git {}` within `{}` failed with status: {}\nstderr:\n{}\nstdout:\n{}",
417 args.join(" "),
418 repo_path.display(),
419 output.status,
420 String::from_utf8_lossy(&output.stderr),
421 String::from_utf8_lossy(&output.stdout),
422 );
423 Ok(String::from_utf8(output.stdout)?.trim().to_string())
424}
425
426impl Display for NamedExample {
427 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
428 write!(f, "# {}\n\n", self.name)?;
429 write!(
430 f,
431 "{REPOSITORY_URL_FIELD} = {}\n",
432 self.example.repository_url
433 )?;
434 write!(f, "{REVISION_FIELD} = {}\n\n", self.example.revision)?;
435
436 write!(f, "## {UNCOMMITTED_DIFF_HEADING}\n\n")?;
437 write!(f, "`````diff\n")?;
438 write!(f, "{}", self.example.uncommitted_diff)?;
439 write!(f, "`````\n")?;
440
441 if !self.example.edit_history.is_empty() {
442 write!(f, "`````diff\n{}`````\n", self.example.edit_history)?;
443 }
444
445 write!(
446 f,
447 "## {CURSOR_POSITION_HEADING}\n\n`````{}\n{}`````\n",
448 self.example.cursor_path.display(),
449 self.example.cursor_position
450 )?;
451 write!(f, "## {EDIT_HISTORY_HEADING}\n\n")?;
452
453 if !self.example.expected_patch.is_empty() {
454 write!(
455 f,
456 "\n## {EXPECTED_PATCH_HEADING}\n\n`````diff\n{}`````\n",
457 self.example.expected_patch
458 )?;
459 }
460
461 if !self.example.expected_excerpts.is_empty() {
462 write!(f, "\n## {EXPECTED_EXCERPTS_HEADING}\n\n")?;
463
464 for excerpt in &self.example.expected_excerpts {
465 write!(
466 f,
467 "`````{}{}\n{}`````\n\n",
468 excerpt
469 .path
470 .extension()
471 .map(|ext| format!("{} ", ext.to_string_lossy()))
472 .unwrap_or_default(),
473 excerpt.path.display(),
474 excerpt.text
475 )?;
476 }
477 }
478
479 Ok(())
480 }
481}
482
483thread_local! {
484 static REPO_LOCKS: RefCell<HashMap<PathBuf, Arc<Mutex<()>>>> = RefCell::new(HashMap::default());
485}
486
487#[must_use]
488pub async fn lock_repo(path: impl AsRef<Path>) -> OwnedMutexGuard<()> {
489 REPO_LOCKS
490 .with(|cell| {
491 cell.borrow_mut()
492 .entry(path.as_ref().to_path_buf())
493 .or_default()
494 .clone()
495 })
496 .lock_owned()
497 .await
498}