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