1use anyhow::{Context as _, Result};
2use serde::{Deserialize, Serialize};
3use std::{borrow::Cow, fmt::Write as _, mem, path::Path, sync::Arc};
4
5pub const CURSOR_POSITION_MARKER: &str = "[CURSOR_POSITION]";
6pub const INLINE_CURSOR_MARKER: &str = "<|user_cursor|>";
7
8#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
9pub struct ExampleSpec {
10 #[serde(default)]
11 pub name: String,
12 pub repository_url: String,
13 pub revision: String,
14 #[serde(default, skip_serializing_if = "Vec::is_empty")]
15 pub tags: Vec<String>,
16 #[serde(default, skip_serializing_if = "Option::is_none")]
17 pub reasoning: Option<String>,
18 #[serde(default)]
19 pub uncommitted_diff: String,
20 pub cursor_path: Arc<Path>,
21 pub cursor_position: String,
22 pub edit_history: String,
23 pub expected_patches: Vec<String>,
24 #[serde(default, skip_serializing_if = "Option::is_none")]
25 pub rejected_patch: Option<String>,
26}
27
28const REASONING_HEADING: &str = "Reasoning";
29const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
30const EDIT_HISTORY_HEADING: &str = "Edit History";
31const CURSOR_POSITION_HEADING: &str = "Cursor Position";
32const EXPECTED_PATCH_HEADING: &str = "Expected Patch";
33const REJECTED_PATCH_HEADING: &str = "Rejected Patch";
34
35#[derive(Serialize, Deserialize)]
36struct FrontMatter<'a> {
37 repository_url: Cow<'a, str>,
38 revision: Cow<'a, str>,
39 #[serde(default, skip_serializing_if = "Vec::is_empty")]
40 tags: Vec<String>,
41}
42
43impl ExampleSpec {
44 /// Generate a sanitized filename for this example.
45 pub fn filename(&self) -> String {
46 self.name
47 .chars()
48 .map(|c| match c {
49 ' ' | ':' | '~' | '^' | '?' | '*' | '[' | '\\' | '@' | '{' | '/' | '<' | '>'
50 | '|' | '"' => '-',
51 c => c,
52 })
53 .collect()
54 }
55
56 /// Format this example spec as markdown.
57 pub fn to_markdown(&self) -> String {
58 use std::fmt::Write as _;
59
60 let front_matter = FrontMatter {
61 repository_url: Cow::Borrowed(&self.repository_url),
62 revision: Cow::Borrowed(&self.revision),
63 tags: self.tags.clone(),
64 };
65 let front_matter_toml =
66 toml::to_string_pretty(&front_matter).unwrap_or_else(|_| String::new());
67
68 let mut markdown = String::new();
69
70 _ = writeln!(markdown, "+++");
71 markdown.push_str(&front_matter_toml);
72 if !markdown.ends_with('\n') {
73 markdown.push('\n');
74 }
75 _ = writeln!(markdown, "+++");
76 markdown.push('\n');
77
78 _ = writeln!(markdown, "# {}", self.name);
79 markdown.push('\n');
80
81 if let Some(reasoning) = &self.reasoning {
82 _ = writeln!(markdown, "## {}", REASONING_HEADING);
83 markdown.push('\n');
84 markdown.push_str(reasoning);
85 if !markdown.ends_with('\n') {
86 markdown.push('\n');
87 }
88 markdown.push('\n');
89 }
90
91 if !self.uncommitted_diff.is_empty() {
92 _ = writeln!(markdown, "## {}", UNCOMMITTED_DIFF_HEADING);
93 _ = writeln!(markdown);
94 _ = writeln!(markdown, "```diff");
95 markdown.push_str(&self.uncommitted_diff);
96 if !markdown.ends_with('\n') {
97 markdown.push('\n');
98 }
99 _ = writeln!(markdown, "```");
100 markdown.push('\n');
101 }
102
103 _ = writeln!(markdown, "## {}", EDIT_HISTORY_HEADING);
104 _ = writeln!(markdown);
105
106 if self.edit_history.is_empty() {
107 _ = writeln!(markdown, "(No edit history)");
108 _ = writeln!(markdown);
109 } else {
110 _ = writeln!(markdown, "```diff");
111 markdown.push_str(&self.edit_history);
112 if !markdown.ends_with('\n') {
113 markdown.push('\n');
114 }
115 _ = writeln!(markdown, "```");
116 markdown.push('\n');
117 }
118
119 _ = writeln!(markdown, "## {}", CURSOR_POSITION_HEADING);
120 _ = writeln!(markdown);
121 _ = writeln!(markdown, "```{}", self.cursor_path.to_string_lossy());
122 markdown.push_str(&self.cursor_position);
123 if !markdown.ends_with('\n') {
124 markdown.push('\n');
125 }
126 _ = writeln!(markdown, "```");
127 markdown.push('\n');
128
129 _ = writeln!(markdown, "## {}", EXPECTED_PATCH_HEADING);
130 markdown.push('\n');
131 for patch in &self.expected_patches {
132 _ = writeln!(markdown, "```diff");
133 markdown.push_str(patch);
134 if !markdown.ends_with('\n') {
135 markdown.push('\n');
136 }
137 _ = writeln!(markdown, "```");
138 markdown.push('\n');
139 }
140
141 if let Some(rejected_patch) = &self.rejected_patch {
142 _ = writeln!(markdown, "## {}", REJECTED_PATCH_HEADING);
143 markdown.push('\n');
144 _ = writeln!(markdown, "```diff");
145 markdown.push_str(rejected_patch);
146 if !markdown.ends_with('\n') {
147 markdown.push('\n');
148 }
149 _ = writeln!(markdown, "```");
150 markdown.push('\n');
151 }
152
153 markdown
154 }
155
156 /// Parse an example spec from markdown.
157 pub fn from_markdown(mut input: &str) -> anyhow::Result<Self> {
158 use pulldown_cmark::{CodeBlockKind, CowStr, Event, HeadingLevel, Parser, Tag, TagEnd};
159
160 let mut spec = ExampleSpec {
161 name: String::new(),
162 repository_url: String::new(),
163 revision: String::new(),
164 tags: Vec::new(),
165 reasoning: None,
166 uncommitted_diff: String::new(),
167 cursor_path: Path::new("").into(),
168 cursor_position: String::new(),
169 edit_history: String::new(),
170 expected_patches: Vec::new(),
171 rejected_patch: None,
172 };
173
174 if let Some(rest) = input.strip_prefix("+++\n")
175 && let Some((front_matter, rest)) = rest.split_once("+++\n")
176 {
177 if let Ok(data) = toml::from_str::<FrontMatter<'_>>(front_matter) {
178 spec.repository_url = data.repository_url.into_owned();
179 spec.revision = data.revision.into_owned();
180 spec.tags = data.tags;
181 }
182 input = rest.trim_start();
183 }
184
185 let parser = Parser::new(input);
186 let mut text = String::new();
187 let mut block_info: CowStr = "".into();
188
189 #[derive(PartialEq)]
190 enum Section {
191 Start,
192 UncommittedDiff,
193 EditHistory,
194 CursorPosition,
195 ExpectedPatch,
196 RejectedPatch,
197 Other,
198 }
199
200 let mut current_section = Section::Start;
201
202 for event in parser {
203 match event {
204 Event::Text(line) => {
205 text.push_str(&line);
206 }
207 Event::End(TagEnd::Heading(HeadingLevel::H1)) => {
208 spec.name = mem::take(&mut text);
209 }
210 Event::End(TagEnd::Heading(HeadingLevel::H2)) => {
211 let title = mem::take(&mut text);
212 current_section = if title.eq_ignore_ascii_case(UNCOMMITTED_DIFF_HEADING) {
213 Section::UncommittedDiff
214 } else if title.eq_ignore_ascii_case(EDIT_HISTORY_HEADING) {
215 Section::EditHistory
216 } else if title.eq_ignore_ascii_case(CURSOR_POSITION_HEADING) {
217 Section::CursorPosition
218 } else if title.eq_ignore_ascii_case(EXPECTED_PATCH_HEADING) {
219 Section::ExpectedPatch
220 } else if title.eq_ignore_ascii_case(REJECTED_PATCH_HEADING) {
221 Section::RejectedPatch
222 } else {
223 Section::Other
224 };
225 }
226 Event::End(TagEnd::Heading(HeadingLevel::H3)) => {
227 mem::take(&mut text);
228 }
229 Event::End(TagEnd::Heading(HeadingLevel::H4)) => {
230 mem::take(&mut text);
231 }
232 Event::End(TagEnd::Heading(level)) => {
233 anyhow::bail!("Unexpected heading level: {level}");
234 }
235 Event::Start(Tag::CodeBlock(kind)) => {
236 match kind {
237 CodeBlockKind::Fenced(info) => {
238 block_info = info;
239 }
240 CodeBlockKind::Indented => {
241 anyhow::bail!("Unexpected indented codeblock");
242 }
243 };
244 }
245 Event::Start(_) => {
246 text.clear();
247 block_info = "".into();
248 }
249 Event::End(TagEnd::CodeBlock) => {
250 let block_info = block_info.trim();
251 match current_section {
252 Section::UncommittedDiff => {
253 spec.uncommitted_diff = mem::take(&mut text);
254 }
255 Section::EditHistory => {
256 spec.edit_history.push_str(&mem::take(&mut text));
257 }
258 Section::CursorPosition => {
259 spec.cursor_path = Path::new(block_info).into();
260 spec.cursor_position = mem::take(&mut text);
261 }
262 Section::ExpectedPatch => {
263 spec.expected_patches.push(mem::take(&mut text));
264 }
265 Section::RejectedPatch => {
266 spec.rejected_patch = Some(mem::take(&mut text));
267 }
268 Section::Start | Section::Other => {}
269 }
270 }
271 _ => {}
272 }
273 }
274
275 if spec.cursor_path.as_ref() == Path::new("") || spec.cursor_position.is_empty() {
276 anyhow::bail!("Missing cursor position codeblock");
277 }
278
279 Ok(spec)
280 }
281
282 /// Returns the excerpt of text around the cursor, and the offset of the cursor within that
283 /// excerpt.
284 ///
285 /// The cursor's position is marked with a special comment that appears
286 /// below the cursor line, which contains the string `[CURSOR_POSITION]`,
287 /// preceded by an arrow marking the cursor's column. The arrow can be
288 /// either:
289 /// - `^` - The cursor column is at the position of the `^` character (pointing up to the cursor)
290 /// - `<` - The cursor column is at the first non-whitespace character on that line.
291 pub fn cursor_excerpt(&self) -> Result<(String, usize)> {
292 let input = &self.cursor_position;
293
294 // Check for inline cursor marker first
295 if let Some(inline_offset) = input.find(INLINE_CURSOR_MARKER) {
296 let excerpt = input[..inline_offset].to_string()
297 + &input[inline_offset + INLINE_CURSOR_MARKER.len()..];
298 return Ok((excerpt, inline_offset));
299 }
300
301 let marker_offset = input
302 .find(CURSOR_POSITION_MARKER)
303 .context("missing [CURSOR_POSITION] marker")?;
304 let marker_line_start = input[..marker_offset]
305 .rfind('\n')
306 .map(|pos| pos + 1)
307 .unwrap_or(0);
308 let marker_line_end = input[marker_line_start..]
309 .find('\n')
310 .map(|pos| marker_line_start + pos + 1)
311 .unwrap_or(input.len());
312 let marker_line = &input[marker_line_start..marker_line_end].trim_end_matches('\n');
313
314 let cursor_column = if let Some(cursor_offset) = marker_line.find('^') {
315 cursor_offset
316 } else if let Some(less_than_pos) = marker_line.find('<') {
317 marker_line
318 .find(|c: char| !c.is_whitespace())
319 .unwrap_or(less_than_pos)
320 } else {
321 anyhow::bail!(
322 "cursor position marker line must contain '^' or '<' before [CURSOR_POSITION]"
323 );
324 };
325
326 let mut excerpt = input[..marker_line_start].to_string() + &input[marker_line_end..];
327 excerpt.truncate(excerpt.trim_end_matches('\n').len());
328
329 // The cursor is on the line above the marker line.
330 let cursor_line_end = marker_line_start.saturating_sub(1);
331 let cursor_line_start = excerpt[..cursor_line_end]
332 .rfind('\n')
333 .map(|pos| pos + 1)
334 .unwrap_or(0);
335 let cursor_offset = cursor_line_start + cursor_column;
336
337 Ok((excerpt, cursor_offset))
338 }
339
340 /// Sets the cursor position excerpt from a plain excerpt and cursor byte offset.
341 ///
342 /// The `line_comment_prefix` is used to format the marker line as a comment.
343 /// If the cursor column is less than the comment prefix length, the `<` format is used.
344 /// Otherwise, the `^` format is used.
345 pub fn set_cursor_excerpt(
346 &mut self,
347 excerpt: &str,
348 cursor_offset: usize,
349 line_comment_prefix: &str,
350 ) {
351 // Find which line the cursor is on and its column
352 let cursor_line_start = excerpt[..cursor_offset]
353 .rfind('\n')
354 .map(|pos| pos + 1)
355 .unwrap_or(0);
356 let cursor_line_end = excerpt[cursor_line_start..]
357 .find('\n')
358 .map(|pos| cursor_line_start + pos + 1)
359 .unwrap_or(excerpt.len());
360 let cursor_line = &excerpt[cursor_line_start..cursor_line_end];
361 let cursor_line_indent = &cursor_line[..cursor_line.len() - cursor_line.trim_start().len()];
362 let cursor_column = cursor_offset - cursor_line_start;
363
364 // Build the marker line
365 let mut marker_line = String::new();
366 if cursor_column < line_comment_prefix.len() {
367 for _ in 0..cursor_column {
368 marker_line.push(' ');
369 }
370 marker_line.push_str(line_comment_prefix);
371 write!(marker_line, " <{}", CURSOR_POSITION_MARKER).unwrap();
372 } else {
373 if cursor_column >= cursor_line_indent.len() + line_comment_prefix.len() {
374 marker_line.push_str(cursor_line_indent);
375 }
376 marker_line.push_str(line_comment_prefix);
377 while marker_line.len() < cursor_column {
378 marker_line.push(' ');
379 }
380 write!(marker_line, "^{}", CURSOR_POSITION_MARKER).unwrap();
381 }
382
383 // Build the final cursor_position string
384 let mut result = String::with_capacity(excerpt.len() + marker_line.len() + 2);
385 result.push_str(&excerpt[..cursor_line_end]);
386 if !result.ends_with('\n') {
387 result.push('\n');
388 }
389 result.push_str(&marker_line);
390 if cursor_line_end < excerpt.len() {
391 result.push('\n');
392 result.push_str(&excerpt[cursor_line_end..]);
393 }
394
395 self.cursor_position = result;
396 }
397}
398
399#[cfg(test)]
400mod tests {
401 use super::*;
402 use indoc::indoc;
403
404 #[test]
405 fn test_cursor_excerpt_with_caret() {
406 let mut spec = ExampleSpec {
407 name: String::new(),
408 repository_url: String::new(),
409 revision: String::new(),
410 tags: Vec::new(),
411 reasoning: None,
412 uncommitted_diff: String::new(),
413 cursor_path: Path::new("test.rs").into(),
414 cursor_position: String::new(),
415 edit_history: String::new(),
416 expected_patches: Vec::new(),
417 rejected_patch: None,
418 };
419
420 // Cursor before `42`
421 let excerpt = indoc! {"
422 fn main() {
423 let x = 42;
424 println!(\"{}\", x);
425 }"
426 };
427 let offset = excerpt.find("42").unwrap();
428 let position_string = indoc! {"
429 fn main() {
430 let x = 42;
431 // ^[CURSOR_POSITION]
432 println!(\"{}\", x);
433 }"
434 }
435 .to_string();
436
437 spec.set_cursor_excerpt(excerpt, offset, "//");
438 assert_eq!(spec.cursor_position, position_string);
439 assert_eq!(
440 spec.cursor_excerpt().unwrap(),
441 (excerpt.to_string(), offset)
442 );
443
444 // Cursor after `l` in `let`
445 let offset = excerpt.find("et x").unwrap();
446 let position_string = indoc! {"
447 fn main() {
448 let x = 42;
449 // ^[CURSOR_POSITION]
450 println!(\"{}\", x);
451 }"
452 }
453 .to_string();
454
455 spec.set_cursor_excerpt(excerpt, offset, "//");
456 assert_eq!(spec.cursor_position, position_string);
457 assert_eq!(
458 spec.cursor_excerpt().unwrap(),
459 (excerpt.to_string(), offset)
460 );
461
462 // Cursor before `let`
463 let offset = excerpt.find("let").unwrap();
464 let position_string = indoc! {"
465 fn main() {
466 let x = 42;
467 // ^[CURSOR_POSITION]
468 println!(\"{}\", x);
469 }"
470 }
471 .to_string();
472
473 spec.set_cursor_excerpt(excerpt, offset, "//");
474 assert_eq!(spec.cursor_position, position_string);
475 assert_eq!(
476 spec.cursor_excerpt().unwrap(),
477 (excerpt.to_string(), offset)
478 );
479
480 // Cursor at beginning of the line with `let`
481 let offset = excerpt.find(" let").unwrap();
482 let position_string = indoc! {"
483 fn main() {
484 let x = 42;
485 // <[CURSOR_POSITION]
486 println!(\"{}\", x);
487 }"
488 }
489 .to_string();
490
491 spec.set_cursor_excerpt(excerpt, offset, "//");
492 assert_eq!(spec.cursor_position, position_string);
493 assert_eq!(
494 spec.cursor_excerpt().unwrap(),
495 (excerpt.to_string(), offset)
496 );
497
498 // Cursor at end of line, after the semicolon
499 let offset = excerpt.find(';').unwrap() + 1;
500 let position_string = indoc! {"
501 fn main() {
502 let x = 42;
503 // ^[CURSOR_POSITION]
504 println!(\"{}\", x);
505 }"
506 }
507 .to_string();
508
509 spec.set_cursor_excerpt(excerpt, offset, "//");
510 assert_eq!(spec.cursor_position, position_string);
511 assert_eq!(
512 spec.cursor_excerpt().unwrap(),
513 (excerpt.to_string(), offset)
514 );
515
516 // Caret at end of file (no trailing newline)
517 let excerpt = indoc! {"
518 fn main() {
519 let x = 42;"
520 };
521 let offset = excerpt.find(';').unwrap() + 1;
522 let position_string = indoc! {"
523 fn main() {
524 let x = 42;
525 // ^[CURSOR_POSITION]"
526 }
527 .to_string();
528
529 spec.set_cursor_excerpt(excerpt, offset, "//");
530 assert_eq!(spec.cursor_position, position_string);
531 assert_eq!(
532 spec.cursor_excerpt().unwrap(),
533 (excerpt.to_string(), offset)
534 );
535 }
536
537 #[test]
538 fn test_cursor_excerpt_with_inline_marker() {
539 let mut spec = ExampleSpec {
540 name: String::new(),
541 repository_url: String::new(),
542 revision: String::new(),
543 tags: Vec::new(),
544 reasoning: None,
545 uncommitted_diff: String::new(),
546 cursor_path: Path::new("test.rs").into(),
547 cursor_position: String::new(),
548 edit_history: String::new(),
549 expected_patches: Vec::new(),
550 rejected_patch: None,
551 };
552
553 // Cursor before `42` using inline marker
554 spec.cursor_position = indoc! {"
555 fn main() {
556 let x = <|user_cursor|>42;
557 println!(\"{}\", x);
558 }"
559 }
560 .to_string();
561
562 let expected_excerpt = indoc! {"
563 fn main() {
564 let x = 42;
565 println!(\"{}\", x);
566 }"
567 };
568 let expected_offset = expected_excerpt.find("42").unwrap();
569
570 assert_eq!(
571 spec.cursor_excerpt().unwrap(),
572 (expected_excerpt.to_string(), expected_offset)
573 );
574
575 // Cursor at beginning of line
576 spec.cursor_position = indoc! {"
577 fn main() {
578 <|user_cursor|> let x = 42;
579 }"
580 }
581 .to_string();
582
583 let expected_excerpt = indoc! {"
584 fn main() {
585 let x = 42;
586 }"
587 };
588 let expected_offset = expected_excerpt.find(" let").unwrap();
589
590 assert_eq!(
591 spec.cursor_excerpt().unwrap(),
592 (expected_excerpt.to_string(), expected_offset)
593 );
594
595 // Cursor at end of file
596 spec.cursor_position = "fn main() {}<|user_cursor|>".to_string();
597 let expected_excerpt = "fn main() {}";
598 let expected_offset = expected_excerpt.len();
599
600 assert_eq!(
601 spec.cursor_excerpt().unwrap(),
602 (expected_excerpt.to_string(), expected_offset)
603 );
604 }
605}