use std::{ borrow::Cow, cell::RefCell, fmt::{self, Display}, fs, io::Write, mem, path::{Path, PathBuf}, sync::{Arc, OnceLock}, }; use crate::headless::ZetaCliAppState; use anyhow::{Context as _, Result, anyhow}; use clap::ValueEnum; use cloud_zeta2_prompt::CURSOR_MARKER; use collections::HashMap; use edit_prediction_context::Line; use futures::{ AsyncWriteExt as _, lock::{Mutex, OwnedMutexGuard}, }; use futures::{FutureExt as _, future::Shared}; use gpui::{AsyncApp, Entity, Task, http_client::Url}; use language::{Anchor, Buffer}; use project::{Project, ProjectPath}; use pulldown_cmark::CowStr; use serde::{Deserialize, Serialize}; use util::{paths::PathStyle, rel_path::RelPath}; use zeta2::udiff::OpenedBuffers; use crate::paths::{REPOS_DIR, WORKTREES_DIR}; const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff"; const EDIT_HISTORY_HEADING: &str = "Edit History"; const CURSOR_POSITION_HEADING: &str = "Cursor Position"; const EXPECTED_PATCH_HEADING: &str = "Expected Patch"; const EXPECTED_CONTEXT_HEADING: &str = "Expected Context"; const REPOSITORY_URL_FIELD: &str = "repository_url"; const REVISION_FIELD: &str = "revision"; #[derive(Debug, Clone)] pub struct NamedExample { pub name: String, pub example: Example, } #[derive(Clone, Debug, Serialize, Deserialize)] pub struct Example { pub repository_url: String, pub revision: String, pub uncommitted_diff: String, pub cursor_path: PathBuf, pub cursor_position: String, pub edit_history: String, pub expected_patch: String, pub expected_context: Vec, } pub type ActualExcerpt = Excerpt; #[derive(Clone, Debug, Serialize, Deserialize)] pub struct Excerpt { pub path: PathBuf, pub text: String, } #[derive(Default, Clone, Debug, Serialize, Deserialize)] pub struct ExpectedContextEntry { pub heading: String, pub alternatives: Vec, } #[derive(Default, Clone, Debug, Serialize, Deserialize)] pub struct ExpectedExcerptSet { pub heading: String, pub excerpts: Vec, } #[derive(Clone, Debug, Serialize, Deserialize)] pub struct ExpectedExcerpt { pub path: PathBuf, pub text: String, pub required_lines: Vec, } #[derive(ValueEnum, Debug, Clone)] pub enum ExampleFormat { Json, Toml, Md, } impl NamedExample { pub fn load(path: impl AsRef) -> Result { let path = path.as_ref(); let content = std::fs::read_to_string(path)?; let ext = path.extension(); match ext.and_then(|s| s.to_str()) { Some("json") => Ok(Self { name: path.file_stem().unwrap_or_default().display().to_string(), example: serde_json::from_str(&content)?, }), Some("toml") => Ok(Self { name: path.file_stem().unwrap_or_default().display().to_string(), example: toml::from_str(&content)?, }), Some("md") => Self::parse_md(&content), Some(_) => { anyhow::bail!("Unrecognized example extension: {}", ext.unwrap().display()); } None => { anyhow::bail!( "Failed to determine example type since the file does not have an extension." ); } } } pub fn parse_md(input: &str) -> Result { use pulldown_cmark::{CodeBlockKind, Event, HeadingLevel, Parser, Tag, TagEnd}; let parser = Parser::new(input); let mut named = NamedExample { name: String::new(), example: Example { repository_url: String::new(), revision: String::new(), uncommitted_diff: String::new(), cursor_path: PathBuf::new(), cursor_position: String::new(), edit_history: String::new(), expected_patch: String::new(), expected_context: Vec::new(), }, }; let mut text = String::new(); let mut block_info: CowStr = "".into(); #[derive(PartialEq)] enum Section { UncommittedDiff, EditHistory, CursorPosition, ExpectedExcerpts, ExpectedPatch, Other, } let mut current_section = Section::Other; for event in parser { match event { Event::Text(line) => { text.push_str(&line); if !named.name.is_empty() && current_section == Section::Other // in h1 section && let Some((field, value)) = line.split_once('=') { match field.trim() { REPOSITORY_URL_FIELD => { named.example.repository_url = value.trim().to_string(); } REVISION_FIELD => { named.example.revision = value.trim().to_string(); } _ => {} } } } Event::End(TagEnd::Heading(HeadingLevel::H1)) => { if !named.name.is_empty() { anyhow::bail!( "Found multiple H1 headings. There should only be one with the name of the example." ); } named.name = mem::take(&mut text); } Event::End(TagEnd::Heading(HeadingLevel::H2)) => { let title = mem::take(&mut text); current_section = if title.eq_ignore_ascii_case(UNCOMMITTED_DIFF_HEADING) { Section::UncommittedDiff } else if title.eq_ignore_ascii_case(EDIT_HISTORY_HEADING) { Section::EditHistory } else if title.eq_ignore_ascii_case(CURSOR_POSITION_HEADING) { Section::CursorPosition } else if title.eq_ignore_ascii_case(EXPECTED_PATCH_HEADING) { Section::ExpectedPatch } else if title.eq_ignore_ascii_case(EXPECTED_CONTEXT_HEADING) { Section::ExpectedExcerpts } else { Section::Other }; } Event::End(TagEnd::Heading(HeadingLevel::H3)) => { let heading = mem::take(&mut text); match current_section { Section::ExpectedExcerpts => { named.example.expected_context.push(ExpectedContextEntry { heading, alternatives: Vec::new(), }); } _ => {} } } Event::End(TagEnd::Heading(HeadingLevel::H4)) => { let heading = mem::take(&mut text); match current_section { Section::ExpectedExcerpts => { let expected_context = &mut named.example.expected_context; let last_entry = expected_context.last_mut().unwrap(); last_entry.alternatives.push(ExpectedExcerptSet { heading, excerpts: Vec::new(), }) } _ => {} } } Event::End(TagEnd::Heading(level)) => { anyhow::bail!("Unexpected heading level: {level}"); } Event::Start(Tag::CodeBlock(kind)) => { match kind { CodeBlockKind::Fenced(info) => { block_info = info; } CodeBlockKind::Indented => { anyhow::bail!("Unexpected indented codeblock"); } }; } Event::Start(_) => { text.clear(); block_info = "".into(); } Event::End(TagEnd::CodeBlock) => { let block_info = block_info.trim(); match current_section { Section::UncommittedDiff => { named.example.uncommitted_diff = mem::take(&mut text); } Section::EditHistory => { named.example.edit_history.push_str(&mem::take(&mut text)); } Section::CursorPosition => { named.example.cursor_path = block_info.into(); named.example.cursor_position = mem::take(&mut text); } Section::ExpectedExcerpts => { let text = mem::take(&mut text); for excerpt in text.split("\n…\n") { let (mut text, required_lines) = extract_required_lines(&excerpt); if !text.ends_with('\n') { text.push('\n'); } if named.example.expected_context.is_empty() { named.example.expected_context.push(Default::default()); } let alternatives = &mut named .example .expected_context .last_mut() .unwrap() .alternatives; if alternatives.is_empty() { alternatives.push(ExpectedExcerptSet { heading: String::new(), excerpts: vec![], }); } alternatives .last_mut() .unwrap() .excerpts .push(ExpectedExcerpt { path: block_info.into(), text, required_lines, }); } } Section::ExpectedPatch => { named.example.expected_patch = mem::take(&mut text); } Section::Other => {} } } _ => {} } } if named.example.cursor_path.as_path() == Path::new("") || named.example.cursor_position.is_empty() { anyhow::bail!("Missing cursor position codeblock"); } Ok(named) } pub fn write(&self, format: ExampleFormat, mut out: impl Write) -> Result<()> { match format { ExampleFormat::Json => Ok(serde_json::to_writer(out, &self.example)?), ExampleFormat::Toml => { Ok(out.write_all(toml::to_string_pretty(&self.example)?.as_bytes())?) } ExampleFormat::Md => Ok(write!(out, "{}", self)?), } } pub async fn setup_project( &self, app_state: &Arc, cx: &mut AsyncApp, ) -> Result> { let worktree_path = self.setup_worktree().await?; static AUTHENTICATED: OnceLock>> = OnceLock::new(); AUTHENTICATED .get_or_init(|| { let client = app_state.client.clone(); cx.spawn(async move |cx| { client .sign_in_with_optional_connect(true, cx) .await .unwrap(); }) .shared() }) .clone() .await; let project = cx.update(|cx| { Project::local( app_state.client.clone(), app_state.node_runtime.clone(), app_state.user_store.clone(), app_state.languages.clone(), app_state.fs.clone(), None, cx, ) })?; let worktree = project .update(cx, |project, cx| { project.create_worktree(&worktree_path, true, cx) })? .await?; worktree .read_with(cx, |worktree, _cx| { worktree.as_local().unwrap().scan_complete() })? .await; anyhow::Ok(project) } pub async fn setup_worktree(&self) -> Result { let (repo_owner, repo_name) = self.repo_name()?; let file_name = self.file_name(); let repo_dir = REPOS_DIR.join(repo_owner.as_ref()).join(repo_name.as_ref()); let repo_lock = lock_repo(&repo_dir).await; if !repo_dir.is_dir() { fs::create_dir_all(&repo_dir)?; run_git(&repo_dir, &["init"]).await?; run_git( &repo_dir, &["remote", "add", "origin", &self.example.repository_url], ) .await?; } // Resolve the example to a revision, fetching it if needed. let revision = run_git( &repo_dir, &[ "rev-parse", &format!("{}^{{commit}}", self.example.revision), ], ) .await; let revision = if let Ok(revision) = revision { revision } else { run_git( &repo_dir, &["fetch", "--depth", "1", "origin", &self.example.revision], ) .await?; let revision = run_git(&repo_dir, &["rev-parse", "FETCH_HEAD"]).await?; if revision != self.example.revision { run_git(&repo_dir, &["tag", &self.example.revision, &revision]).await?; } revision }; // Create the worktree for this example if needed. let worktree_path = WORKTREES_DIR.join(&file_name).join(repo_name.as_ref()); if worktree_path.is_dir() { run_git(&worktree_path, &["clean", "--force", "-d"]).await?; run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?; run_git(&worktree_path, &["checkout", revision.as_str()]).await?; } else { let worktree_path_string = worktree_path.to_string_lossy(); run_git(&repo_dir, &["branch", "-f", &file_name, revision.as_str()]).await?; run_git( &repo_dir, &["worktree", "add", "-f", &worktree_path_string, &file_name], ) .await?; } drop(repo_lock); // Apply the uncommitted diff for this example. if !self.example.uncommitted_diff.is_empty() { let mut apply_process = smol::process::Command::new("git") .current_dir(&worktree_path) .args(&["apply", "-"]) .stdin(std::process::Stdio::piped()) .spawn()?; let mut stdin = apply_process.stdin.take().unwrap(); stdin .write_all(self.example.uncommitted_diff.as_bytes()) .await?; stdin.close().await?; drop(stdin); let apply_result = apply_process.output().await?; if !apply_result.status.success() { anyhow::bail!( "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}", apply_result.status, String::from_utf8_lossy(&apply_result.stderr), String::from_utf8_lossy(&apply_result.stdout), ); } } Ok(worktree_path) } pub fn file_name(&self) -> String { self.name .chars() .map(|c| { if c.is_whitespace() { '-' } else { c.to_ascii_lowercase() } }) .collect() } fn repo_name(&self) -> Result<(Cow<'_, str>, Cow<'_, str>)> { // git@github.com:owner/repo.git if self.example.repository_url.contains('@') { let (owner, repo) = self .example .repository_url .split_once(':') .context("expected : in git url")? .1 .split_once('/') .context("expected / in git url")?; Ok(( Cow::Borrowed(owner), Cow::Borrowed(repo.trim_end_matches(".git")), )) // http://github.com/owner/repo.git } else { let url = Url::parse(&self.example.repository_url)?; let mut segments = url.path_segments().context("empty http url")?; let owner = segments .next() .context("expected owner path segment")? .to_string(); let repo = segments .next() .context("expected repo path segment")? .trim_end_matches(".git") .to_string(); assert!(segments.next().is_none()); Ok((owner.into(), repo.into())) } } pub async fn cursor_position( &self, project: &Entity, cx: &mut AsyncApp, ) -> Result<(Entity, Anchor)> { let worktree = project.read_with(cx, |project, cx| { project.visible_worktrees(cx).next().unwrap() })?; let cursor_path = RelPath::new(&self.example.cursor_path, PathStyle::Posix)?.into_arc(); let cursor_buffer = project .update(cx, |project, cx| { project.open_buffer( ProjectPath { worktree_id: worktree.read(cx).id(), path: cursor_path, }, cx, ) })? .await?; let cursor_offset_within_excerpt = self .example .cursor_position .find(CURSOR_MARKER) .ok_or_else(|| anyhow!("missing cursor marker"))?; let mut cursor_excerpt = self.example.cursor_position.clone(); cursor_excerpt.replace_range( cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()), "", ); let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| { let text = buffer.text(); let mut matches = text.match_indices(&cursor_excerpt); let Some((excerpt_offset, _)) = matches.next() else { anyhow::bail!( "\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n.Cursor excerpt did not exist in buffer." ); }; assert!(matches.next().is_none()); Ok(excerpt_offset) })??; let cursor_offset = excerpt_offset + cursor_offset_within_excerpt; let cursor_anchor = cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?; Ok((cursor_buffer, cursor_anchor)) } #[must_use] pub async fn apply_edit_history( &self, project: &Entity, cx: &mut AsyncApp, ) -> Result> { zeta2::udiff::apply_diff(&self.example.edit_history, project, cx).await } } fn extract_required_lines(text: &str) -> (String, Vec) { const MARKER: &str = "[ZETA]"; let mut new_text = String::new(); let mut required_lines = Vec::new(); let mut skipped_lines = 0_u32; for (row, mut line) in text.split('\n').enumerate() { if let Some(marker_column) = line.find(MARKER) { let mut strip_column = marker_column; while strip_column > 0 { let prev_char = line[strip_column - 1..].chars().next().unwrap(); if prev_char.is_whitespace() || ['/', '#'].contains(&prev_char) { strip_column -= 1; } else { break; } } let metadata = &line[marker_column + MARKER.len()..]; if metadata.contains("required") { required_lines.push(Line(row as u32 - skipped_lines)); } if strip_column == 0 { skipped_lines += 1; continue; } line = &line[..strip_column]; } new_text.push_str(line); new_text.push('\n'); } new_text.pop(); (new_text, required_lines) } async fn run_git(repo_path: &Path, args: &[&str]) -> Result { let output = smol::process::Command::new("git") .current_dir(repo_path) .args(args) .output() .await?; anyhow::ensure!( output.status.success(), "`git {}` within `{}` failed with status: {}\nstderr:\n{}\nstdout:\n{}", args.join(" "), repo_path.display(), output.status, String::from_utf8_lossy(&output.stderr), String::from_utf8_lossy(&output.stdout), ); Ok(String::from_utf8(output.stdout)?.trim().to_string()) } impl Display for NamedExample { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "# {}\n\n", self.name)?; write!( f, "{REPOSITORY_URL_FIELD} = {}\n", self.example.repository_url )?; write!(f, "{REVISION_FIELD} = {}\n\n", self.example.revision)?; write!(f, "## {UNCOMMITTED_DIFF_HEADING}\n\n")?; write!(f, "`````diff\n")?; write!(f, "{}", self.example.uncommitted_diff)?; write!(f, "`````\n")?; if !self.example.edit_history.is_empty() { write!(f, "`````diff\n{}`````\n", self.example.edit_history)?; } write!( f, "## {CURSOR_POSITION_HEADING}\n\n`````{}\n{}`````\n", self.example.cursor_path.display(), self.example.cursor_position )?; write!(f, "## {EDIT_HISTORY_HEADING}\n\n")?; if !self.example.expected_patch.is_empty() { write!( f, "\n## {EXPECTED_PATCH_HEADING}\n\n`````diff\n{}`````\n", self.example.expected_patch )?; } if !self.example.expected_context.is_empty() { write!(f, "\n## {EXPECTED_CONTEXT_HEADING}\n\n")?; for entry in &self.example.expected_context { write!(f, "\n### {}\n\n", entry.heading)?; let skip_h4 = entry.alternatives.len() == 1 && entry.alternatives[0].heading.is_empty(); for excerpt_set in &entry.alternatives { if !skip_h4 { write!(f, "\n#### {}\n\n", excerpt_set.heading)?; } for excerpt in &excerpt_set.excerpts { write!( f, "`````{}{}\n{}`````\n\n", excerpt .path .extension() .map(|ext| format!("{} ", ext.to_string_lossy())) .unwrap_or_default(), excerpt.path.display(), excerpt.text )?; } } } } Ok(()) } } thread_local! { static REPO_LOCKS: RefCell>>> = RefCell::new(HashMap::default()); } #[must_use] pub async fn lock_repo(path: impl AsRef) -> OwnedMutexGuard<()> { REPO_LOCKS .with(|cell| { cell.borrow_mut() .entry(path.as_ref().to_path_buf()) .or_default() .clone() }) .lock_owned() .await } #[cfg(test)] mod tests { use super::*; use indoc::indoc; use pretty_assertions::assert_eq; #[test] fn test_extract_required_lines() { let input = indoc! {" zero one // [ZETA] required two // [ZETA] something three four # [ZETA] required five "}; let expected_updated_input = indoc! {" zero one two three four five "}; let expected_required_lines = vec![Line(1), Line(4)]; let (updated_input, required_lines) = extract_required_lines(input); assert_eq!(updated_input, expected_updated_input); assert_eq!(required_lines, expected_required_lines); } }