1use crate::{AssistantPanel, InlineAssistId, InlineAssistant};
2use anyhow::{anyhow, Context as _, Result};
3use collections::HashMap;
4use editor::Editor;
5use gpui::AsyncAppContext;
6use gpui::{Model, Task, UpdateGlobal as _, View, WeakView, WindowContext};
7use language::{Buffer, BufferSnapshot};
8use project::{Project, ProjectPath};
9use schemars::JsonSchema;
10use serde::{Deserialize, Serialize};
11use std::{ops::Range, path::Path, sync::Arc};
12use text::Bias;
13use workspace::Workspace;
14
15#[derive(Debug)]
16pub(crate) struct WorkflowStep {
17 pub range: Range<language::Anchor>,
18 pub leading_tags_end: text::Anchor,
19 pub trailing_tag_start: Option<text::Anchor>,
20 pub edits: Arc<[Result<WorkflowStepEdit>]>,
21 pub resolution_task: Option<Task<()>>,
22 pub resolution: Option<Arc<Result<WorkflowStepResolution>>>,
23}
24
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub(crate) struct WorkflowStepEdit {
27 pub path: String,
28 pub kind: WorkflowStepEditKind,
29}
30
31#[derive(Clone, Debug, Eq, PartialEq)]
32pub(crate) struct WorkflowStepResolution {
33 pub title: String,
34 pub suggestion_groups: HashMap<Model<Buffer>, Vec<WorkflowSuggestionGroup>>,
35}
36
37#[derive(Clone, Debug, Eq, PartialEq)]
38pub struct WorkflowSuggestionGroup {
39 pub context_range: Range<language::Anchor>,
40 pub suggestions: Vec<WorkflowSuggestion>,
41}
42
43#[derive(Clone, Debug, Eq, PartialEq)]
44pub enum WorkflowSuggestion {
45 Update {
46 range: Range<language::Anchor>,
47 description: String,
48 },
49 CreateFile {
50 description: String,
51 },
52 InsertBefore {
53 position: language::Anchor,
54 description: String,
55 },
56 InsertAfter {
57 position: language::Anchor,
58 description: String,
59 },
60 Delete {
61 range: Range<language::Anchor>,
62 },
63}
64
65impl WorkflowSuggestion {
66 pub fn range(&self) -> Range<language::Anchor> {
67 match self {
68 Self::Update { range, .. } => range.clone(),
69 Self::CreateFile { .. } => language::Anchor::MIN..language::Anchor::MAX,
70 Self::InsertBefore { position, .. } | Self::InsertAfter { position, .. } => {
71 *position..*position
72 }
73 Self::Delete { range, .. } => range.clone(),
74 }
75 }
76
77 pub fn description(&self) -> Option<&str> {
78 match self {
79 Self::Update { description, .. }
80 | Self::CreateFile { description }
81 | Self::InsertBefore { description, .. }
82 | Self::InsertAfter { description, .. } => Some(description),
83 Self::Delete { .. } => None,
84 }
85 }
86
87 fn description_mut(&mut self) -> Option<&mut String> {
88 match self {
89 Self::Update { description, .. }
90 | Self::CreateFile { description }
91 | Self::InsertBefore { description, .. }
92 | Self::InsertAfter { description, .. } => Some(description),
93 Self::Delete { .. } => None,
94 }
95 }
96
97 pub fn try_merge(&mut self, other: &Self, buffer: &BufferSnapshot) -> bool {
98 let range = self.range();
99 let other_range = other.range();
100
101 // Don't merge if we don't contain the other suggestion.
102 if range.start.cmp(&other_range.start, buffer).is_gt()
103 || range.end.cmp(&other_range.end, buffer).is_lt()
104 {
105 return false;
106 }
107
108 if let Some(description) = self.description_mut() {
109 if let Some(other_description) = other.description() {
110 description.push('\n');
111 description.push_str(other_description);
112 }
113 }
114 true
115 }
116
117 pub fn show(
118 &self,
119 editor: &View<Editor>,
120 excerpt_id: editor::ExcerptId,
121 workspace: &WeakView<Workspace>,
122 assistant_panel: &View<AssistantPanel>,
123 cx: &mut WindowContext,
124 ) -> Option<InlineAssistId> {
125 let mut initial_transaction_id = None;
126 let initial_prompt;
127 let suggestion_range;
128 let buffer = editor.read(cx).buffer().clone();
129 let snapshot = buffer.read(cx).snapshot(cx);
130
131 match self {
132 Self::Update {
133 range, description, ..
134 } => {
135 initial_prompt = description.clone();
136 suggestion_range = snapshot.anchor_in_excerpt(excerpt_id, range.start)?
137 ..snapshot.anchor_in_excerpt(excerpt_id, range.end)?;
138 }
139 Self::CreateFile { description } => {
140 initial_prompt = description.clone();
141 suggestion_range = editor::Anchor::min()..editor::Anchor::min();
142 }
143 Self::InsertBefore {
144 position,
145 description,
146 ..
147 } => {
148 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
149 initial_prompt = description.clone();
150 suggestion_range = buffer.update(cx, |buffer, cx| {
151 buffer.start_transaction(cx);
152 let line_start = buffer.insert_empty_line(position, true, true, cx);
153 initial_transaction_id = buffer.end_transaction(cx);
154 buffer.refresh_preview(cx);
155
156 let line_start = buffer.read(cx).anchor_before(line_start);
157 line_start..line_start
158 });
159 }
160 Self::InsertAfter {
161 position,
162 description,
163 ..
164 } => {
165 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
166 initial_prompt = description.clone();
167 suggestion_range = buffer.update(cx, |buffer, cx| {
168 buffer.start_transaction(cx);
169 let line_start = buffer.insert_empty_line(position, true, true, cx);
170 initial_transaction_id = buffer.end_transaction(cx);
171 buffer.refresh_preview(cx);
172
173 let line_start = buffer.read(cx).anchor_before(line_start);
174 line_start..line_start
175 });
176 }
177 Self::Delete { range, .. } => {
178 initial_prompt = "Delete".to_string();
179 suggestion_range = snapshot.anchor_in_excerpt(excerpt_id, range.start)?
180 ..snapshot.anchor_in_excerpt(excerpt_id, range.end)?;
181 }
182 }
183
184 InlineAssistant::update_global(cx, |inline_assistant, cx| {
185 Some(inline_assistant.suggest_assist(
186 editor,
187 suggestion_range,
188 initial_prompt,
189 initial_transaction_id,
190 false,
191 Some(workspace.clone()),
192 Some(assistant_panel),
193 cx,
194 ))
195 })
196 }
197}
198
199impl WorkflowStepEdit {
200 pub fn new(
201 path: Option<String>,
202 operation: Option<String>,
203 search: Option<String>,
204 description: Option<String>,
205 ) -> Result<Self> {
206 let path = path.ok_or_else(|| anyhow!("missing path"))?;
207 let operation = operation.ok_or_else(|| anyhow!("missing operation"))?;
208
209 let kind = match operation.as_str() {
210 "update" => WorkflowStepEditKind::Update {
211 search: search.ok_or_else(|| anyhow!("missing search"))?,
212 description: description.ok_or_else(|| anyhow!("missing description"))?,
213 },
214 "insert_before" => WorkflowStepEditKind::InsertBefore {
215 search: search.ok_or_else(|| anyhow!("missing search"))?,
216 description: description.ok_or_else(|| anyhow!("missing description"))?,
217 },
218 "insert_after" => WorkflowStepEditKind::InsertAfter {
219 search: search.ok_or_else(|| anyhow!("missing search"))?,
220 description: description.ok_or_else(|| anyhow!("missing description"))?,
221 },
222 "delete" => WorkflowStepEditKind::Delete {
223 search: search.ok_or_else(|| anyhow!("missing search"))?,
224 },
225 "create" => WorkflowStepEditKind::Create {
226 description: description.ok_or_else(|| anyhow!("missing description"))?,
227 },
228 _ => Err(anyhow!("unknown operation {operation:?}"))?,
229 };
230
231 Ok(Self { path, kind })
232 }
233
234 pub async fn resolve(
235 &self,
236 project: Model<Project>,
237 mut cx: AsyncAppContext,
238 ) -> Result<(Model<Buffer>, super::WorkflowSuggestion)> {
239 let path = self.path.clone();
240 let kind = self.kind.clone();
241 let buffer = project
242 .update(&mut cx, |project, cx| {
243 let project_path = project
244 .find_project_path(Path::new(&path), cx)
245 .or_else(|| {
246 // If we couldn't find a project path for it, put it in the active worktree
247 // so that when we create the buffer, it can be saved.
248 let worktree = project
249 .active_entry()
250 .and_then(|entry_id| project.worktree_for_entry(entry_id, cx))
251 .or_else(|| project.worktrees(cx).next())?;
252 let worktree = worktree.read(cx);
253
254 Some(ProjectPath {
255 worktree_id: worktree.id(),
256 path: Arc::from(Path::new(&path)),
257 })
258 })
259 .with_context(|| format!("worktree not found for {:?}", path))?;
260 anyhow::Ok(project.open_buffer(project_path, cx))
261 })??
262 .await?;
263
264 let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
265 let suggestion = cx
266 .background_executor()
267 .spawn(async move {
268 match kind {
269 WorkflowStepEditKind::Update {
270 search,
271 description,
272 } => {
273 let range = Self::resolve_location(&snapshot, &search);
274 WorkflowSuggestion::Update { range, description }
275 }
276 WorkflowStepEditKind::Create { description } => {
277 WorkflowSuggestion::CreateFile { description }
278 }
279 WorkflowStepEditKind::InsertBefore {
280 search,
281 description,
282 } => {
283 let range = Self::resolve_location(&snapshot, &search);
284 WorkflowSuggestion::InsertBefore {
285 position: range.start,
286 description,
287 }
288 }
289 WorkflowStepEditKind::InsertAfter {
290 search,
291 description,
292 } => {
293 let range = Self::resolve_location(&snapshot, &search);
294 WorkflowSuggestion::InsertAfter {
295 position: range.end,
296 description,
297 }
298 }
299 WorkflowStepEditKind::Delete { search } => {
300 let range = Self::resolve_location(&snapshot, &search);
301 WorkflowSuggestion::Delete { range }
302 }
303 }
304 })
305 .await;
306
307 Ok((buffer, suggestion))
308 }
309
310 fn resolve_location(buffer: &text::BufferSnapshot, search_query: &str) -> Range<text::Anchor> {
311 const INSERTION_SCORE: f64 = -1.0;
312 const DELETION_SCORE: f64 = -1.0;
313 const REPLACEMENT_SCORE: f64 = -1.0;
314 const EQUALITY_SCORE: f64 = 5.0;
315
316 struct Matrix {
317 cols: usize,
318 data: Vec<f64>,
319 }
320
321 impl Matrix {
322 fn new(rows: usize, cols: usize) -> Self {
323 Matrix {
324 cols,
325 data: vec![0.0; rows * cols],
326 }
327 }
328
329 fn get(&self, row: usize, col: usize) -> f64 {
330 self.data[row * self.cols + col]
331 }
332
333 fn set(&mut self, row: usize, col: usize, value: f64) {
334 self.data[row * self.cols + col] = value;
335 }
336 }
337
338 let buffer_len = buffer.len();
339 let query_len = search_query.len();
340 let mut matrix = Matrix::new(query_len + 1, buffer_len + 1);
341
342 for (i, query_byte) in search_query.bytes().enumerate() {
343 for (j, buffer_byte) in buffer.bytes_in_range(0..buffer.len()).flatten().enumerate() {
344 let match_score = if query_byte == *buffer_byte {
345 EQUALITY_SCORE
346 } else {
347 REPLACEMENT_SCORE
348 };
349 let up = matrix.get(i + 1, j) + DELETION_SCORE;
350 let left = matrix.get(i, j + 1) + INSERTION_SCORE;
351 let diagonal = matrix.get(i, j) + match_score;
352 let score = up.max(left.max(diagonal)).max(0.);
353 matrix.set(i + 1, j + 1, score);
354 }
355 }
356
357 // Traceback to find the best match
358 let mut best_buffer_end = buffer_len;
359 let mut best_score = 0.0;
360 for col in 1..=buffer_len {
361 let score = matrix.get(query_len, col);
362 if score > best_score {
363 best_score = score;
364 best_buffer_end = col;
365 }
366 }
367
368 let mut query_ix = query_len;
369 let mut buffer_ix = best_buffer_end;
370 while query_ix > 0 && buffer_ix > 0 {
371 let current = matrix.get(query_ix, buffer_ix);
372 let up = matrix.get(query_ix - 1, buffer_ix);
373 let left = matrix.get(query_ix, buffer_ix - 1);
374 if current == left + INSERTION_SCORE {
375 buffer_ix -= 1;
376 } else if current == up + DELETION_SCORE {
377 query_ix -= 1;
378 } else {
379 query_ix -= 1;
380 buffer_ix -= 1;
381 }
382 }
383
384 let mut start = buffer.offset_to_point(buffer.clip_offset(buffer_ix, Bias::Left));
385 start.column = 0;
386 let mut end = buffer.offset_to_point(buffer.clip_offset(best_buffer_end, Bias::Right));
387 end.column = buffer.line_len(end.row);
388
389 buffer.anchor_after(start)..buffer.anchor_before(end)
390 }
391}
392
393#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
394#[serde(tag = "operation")]
395pub enum WorkflowStepEditKind {
396 /// Rewrites the specified text entirely based on the given description.
397 /// This operation completely replaces the given text.
398 Update {
399 /// A string in the source text to apply the update to.
400 search: String,
401 /// A brief description of the transformation to apply to the symbol.
402 description: String,
403 },
404 /// Creates a new file with the given path based on the provided description.
405 /// This operation adds a new file to the codebase.
406 Create {
407 /// A brief description of the file to be created.
408 description: String,
409 },
410 /// Inserts text before the specified text in the source file.
411 InsertBefore {
412 /// A string in the source text to insert text before.
413 search: String,
414 /// A brief description of how the new text should be generated.
415 description: String,
416 },
417 /// Inserts text after the specified text in the source file.
418 InsertAfter {
419 /// A string in the source text to insert text after.
420 search: String,
421 /// A brief description of how the new text should be generated.
422 description: String,
423 },
424 /// Deletes the specified symbol from the containing file.
425 Delete {
426 /// A string in the source text to delete.
427 search: String,
428 },
429}
430
431#[cfg(test)]
432mod tests {
433 use super::*;
434 use gpui::{AppContext, Context};
435 use text::{OffsetRangeExt, Point};
436
437 #[gpui::test]
438 fn test_resolve_location(cx: &mut AppContext) {
439 {
440 let buffer = cx.new_model(|cx| {
441 Buffer::local(
442 concat!(
443 " Lorem\n",
444 " ipsum\n",
445 " dolor sit amet\n",
446 " consecteur",
447 ),
448 cx,
449 )
450 });
451 let snapshot = buffer.read(cx).snapshot();
452 assert_eq!(
453 WorkflowStepEdit::resolve_location(&snapshot, "ipsum\ndolor").to_point(&snapshot),
454 Point::new(1, 0)..Point::new(2, 18)
455 );
456 }
457
458 {
459 let buffer = cx.new_model(|cx| {
460 Buffer::local(
461 concat!(
462 "fn foo1(a: usize) -> usize {\n",
463 " 42\n",
464 "}\n",
465 "\n",
466 "fn foo2(b: usize) -> usize {\n",
467 " 42\n",
468 "}\n",
469 ),
470 cx,
471 )
472 });
473 let snapshot = buffer.read(cx).snapshot();
474 assert_eq!(
475 WorkflowStepEdit::resolve_location(&snapshot, "fn foo1(b: usize) {\n42\n}")
476 .to_point(&snapshot),
477 Point::new(0, 0)..Point::new(2, 1)
478 );
479 }
480
481 {
482 let buffer = cx.new_model(|cx| {
483 Buffer::local(
484 concat!(
485 "fn main() {\n",
486 " Foo\n",
487 " .bar()\n",
488 " .baz()\n",
489 " .qux()\n",
490 "}\n",
491 "\n",
492 "fn foo2(b: usize) -> usize {\n",
493 " 42\n",
494 "}\n",
495 ),
496 cx,
497 )
498 });
499 let snapshot = buffer.read(cx).snapshot();
500 assert_eq!(
501 WorkflowStepEdit::resolve_location(&snapshot, "Foo.bar.baz.qux()")
502 .to_point(&snapshot),
503 Point::new(1, 0)..Point::new(4, 14)
504 );
505 }
506 }
507}