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