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 Some(workspace.clone()),
191 Some(assistant_panel),
192 cx,
193 ))
194 })
195 }
196}
197
198impl WorkflowStepEdit {
199 pub fn new(
200 path: Option<String>,
201 operation: Option<String>,
202 search: Option<String>,
203 description: Option<String>,
204 ) -> Result<Self> {
205 let path = path.ok_or_else(|| anyhow!("missing path"))?;
206 let operation = operation.ok_or_else(|| anyhow!("missing operation"))?;
207
208 let kind = match operation.as_str() {
209 "update" => WorkflowStepEditKind::Update {
210 search: search.ok_or_else(|| anyhow!("missing search"))?,
211 description: description.ok_or_else(|| anyhow!("missing description"))?,
212 },
213 "insert_before" => WorkflowStepEditKind::InsertBefore {
214 search: search.ok_or_else(|| anyhow!("missing search"))?,
215 description: description.ok_or_else(|| anyhow!("missing description"))?,
216 },
217 "insert_after" => WorkflowStepEditKind::InsertAfter {
218 search: search.ok_or_else(|| anyhow!("missing search"))?,
219 description: description.ok_or_else(|| anyhow!("missing description"))?,
220 },
221 "delete" => WorkflowStepEditKind::Delete {
222 search: search.ok_or_else(|| anyhow!("missing search"))?,
223 },
224 "create" => WorkflowStepEditKind::Create {
225 description: description.ok_or_else(|| anyhow!("missing description"))?,
226 },
227 _ => Err(anyhow!("unknown operation {operation:?}"))?,
228 };
229
230 Ok(Self { path, kind })
231 }
232
233 pub async fn resolve(
234 &self,
235 project: Model<Project>,
236 mut cx: AsyncAppContext,
237 ) -> Result<(Model<Buffer>, super::WorkflowSuggestion)> {
238 let path = self.path.clone();
239 let kind = self.kind.clone();
240 let buffer = project
241 .update(&mut cx, |project, cx| {
242 let project_path = project
243 .find_project_path(Path::new(&path), cx)
244 .or_else(|| {
245 // If we couldn't find a project path for it, put it in the active worktree
246 // so that when we create the buffer, it can be saved.
247 let worktree = project
248 .active_entry()
249 .and_then(|entry_id| project.worktree_for_entry(entry_id, cx))
250 .or_else(|| project.worktrees(cx).next())?;
251 let worktree = worktree.read(cx);
252
253 Some(ProjectPath {
254 worktree_id: worktree.id(),
255 path: Arc::from(Path::new(&path)),
256 })
257 })
258 .with_context(|| format!("worktree not found for {:?}", path))?;
259 anyhow::Ok(project.open_buffer(project_path, cx))
260 })??
261 .await?;
262
263 let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
264 let suggestion = cx
265 .background_executor()
266 .spawn(async move {
267 match kind {
268 WorkflowStepEditKind::Update {
269 search,
270 description,
271 } => {
272 let range = Self::resolve_location(&snapshot, &search);
273 WorkflowSuggestion::Update { range, description }
274 }
275 WorkflowStepEditKind::Create { description } => {
276 WorkflowSuggestion::CreateFile { description }
277 }
278 WorkflowStepEditKind::InsertBefore {
279 search,
280 description,
281 } => {
282 let range = Self::resolve_location(&snapshot, &search);
283 WorkflowSuggestion::InsertBefore {
284 position: range.start,
285 description,
286 }
287 }
288 WorkflowStepEditKind::InsertAfter {
289 search,
290 description,
291 } => {
292 let range = Self::resolve_location(&snapshot, &search);
293 WorkflowSuggestion::InsertAfter {
294 position: range.end,
295 description,
296 }
297 }
298 WorkflowStepEditKind::Delete { search } => {
299 let range = Self::resolve_location(&snapshot, &search);
300 WorkflowSuggestion::Delete { range }
301 }
302 }
303 })
304 .await;
305
306 Ok((buffer, suggestion))
307 }
308
309 fn resolve_location(buffer: &text::BufferSnapshot, search_query: &str) -> Range<text::Anchor> {
310 const INSERTION_SCORE: f64 = -1.0;
311 const DELETION_SCORE: f64 = -1.0;
312 const REPLACEMENT_SCORE: f64 = -1.0;
313 const EQUALITY_SCORE: f64 = 5.0;
314
315 struct Matrix {
316 cols: usize,
317 data: Vec<f64>,
318 }
319
320 impl Matrix {
321 fn new(rows: usize, cols: usize) -> Self {
322 Matrix {
323 cols,
324 data: vec![0.0; rows * cols],
325 }
326 }
327
328 fn get(&self, row: usize, col: usize) -> f64 {
329 self.data[row * self.cols + col]
330 }
331
332 fn set(&mut self, row: usize, col: usize, value: f64) {
333 self.data[row * self.cols + col] = value;
334 }
335 }
336
337 let buffer_len = buffer.len();
338 let query_len = search_query.len();
339 let mut matrix = Matrix::new(query_len + 1, buffer_len + 1);
340
341 for (i, query_byte) in search_query.bytes().enumerate() {
342 for (j, buffer_byte) in buffer.bytes_in_range(0..buffer.len()).flatten().enumerate() {
343 let match_score = if query_byte == *buffer_byte {
344 EQUALITY_SCORE
345 } else {
346 REPLACEMENT_SCORE
347 };
348 let up = matrix.get(i + 1, j) + DELETION_SCORE;
349 let left = matrix.get(i, j + 1) + INSERTION_SCORE;
350 let diagonal = matrix.get(i, j) + match_score;
351 let score = up.max(left.max(diagonal)).max(0.);
352 matrix.set(i + 1, j + 1, score);
353 }
354 }
355
356 // Traceback to find the best match
357 let mut best_buffer_end = buffer_len;
358 let mut best_score = 0.0;
359 for col in 1..=buffer_len {
360 let score = matrix.get(query_len, col);
361 if score > best_score {
362 best_score = score;
363 best_buffer_end = col;
364 }
365 }
366
367 let mut query_ix = query_len;
368 let mut buffer_ix = best_buffer_end;
369 while query_ix > 0 && buffer_ix > 0 {
370 let current = matrix.get(query_ix, buffer_ix);
371 let up = matrix.get(query_ix - 1, buffer_ix);
372 let left = matrix.get(query_ix, buffer_ix - 1);
373 if current == left + INSERTION_SCORE {
374 buffer_ix -= 1;
375 } else if current == up + DELETION_SCORE {
376 query_ix -= 1;
377 } else {
378 query_ix -= 1;
379 buffer_ix -= 1;
380 }
381 }
382
383 let mut start = buffer.offset_to_point(buffer.clip_offset(buffer_ix, Bias::Left));
384 start.column = 0;
385 let mut end = buffer.offset_to_point(buffer.clip_offset(best_buffer_end, Bias::Right));
386 end.column = buffer.line_len(end.row);
387
388 buffer.anchor_after(start)..buffer.anchor_before(end)
389 }
390}
391
392#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
393#[serde(tag = "operation")]
394pub enum WorkflowStepEditKind {
395 /// Rewrites the specified text entirely based on the given description.
396 /// This operation completely replaces the given text.
397 Update {
398 /// A string in the source text to apply the update to.
399 search: String,
400 /// A brief description of the transformation to apply to the symbol.
401 description: String,
402 },
403 /// Creates a new file with the given path based on the provided description.
404 /// This operation adds a new file to the codebase.
405 Create {
406 /// A brief description of the file to be created.
407 description: String,
408 },
409 /// Inserts text before the specified text in the source file.
410 InsertBefore {
411 /// A string in the source text to insert text before.
412 search: String,
413 /// A brief description of how the new text should be generated.
414 description: String,
415 },
416 /// Inserts text after the specified text in the source file.
417 InsertAfter {
418 /// A string in the source text to insert text after.
419 search: String,
420 /// A brief description of how the new text should be generated.
421 description: String,
422 },
423 /// Deletes the specified symbol from the containing file.
424 Delete {
425 /// A string in the source text to delete.
426 search: String,
427 },
428}
429
430#[cfg(test)]
431mod tests {
432 use super::*;
433 use gpui::{AppContext, Context};
434 use text::{OffsetRangeExt, Point};
435
436 #[gpui::test]
437 fn test_resolve_location(cx: &mut AppContext) {
438 {
439 let buffer = cx.new_model(|cx| {
440 Buffer::local(
441 concat!(
442 " Lorem\n",
443 " ipsum\n",
444 " dolor sit amet\n",
445 " consecteur",
446 ),
447 cx,
448 )
449 });
450 let snapshot = buffer.read(cx).snapshot();
451 assert_eq!(
452 WorkflowStepEdit::resolve_location(&snapshot, "ipsum\ndolor").to_point(&snapshot),
453 Point::new(1, 0)..Point::new(2, 18)
454 );
455 }
456
457 {
458 let buffer = cx.new_model(|cx| {
459 Buffer::local(
460 concat!(
461 "fn foo1(a: usize) -> usize {\n",
462 " 42\n",
463 "}\n",
464 "\n",
465 "fn foo2(b: usize) -> usize {\n",
466 " 42\n",
467 "}\n",
468 ),
469 cx,
470 )
471 });
472 let snapshot = buffer.read(cx).snapshot();
473 assert_eq!(
474 WorkflowStepEdit::resolve_location(&snapshot, "fn foo1(b: usize) {\n42\n}")
475 .to_point(&snapshot),
476 Point::new(0, 0)..Point::new(2, 1)
477 );
478 }
479
480 {
481 let buffer = cx.new_model(|cx| {
482 Buffer::local(
483 concat!(
484 "fn main() {\n",
485 " Foo\n",
486 " .bar()\n",
487 " .baz()\n",
488 " .qux()\n",
489 "}\n",
490 "\n",
491 "fn foo2(b: usize) -> usize {\n",
492 " 42\n",
493 "}\n",
494 ),
495 cx,
496 )
497 });
498 let snapshot = buffer.read(cx).snapshot();
499 assert_eq!(
500 WorkflowStepEdit::resolve_location(&snapshot, "Foo.bar.baz.qux()")
501 .to_point(&snapshot),
502 Point::new(1, 0)..Point::new(4, 14)
503 );
504 }
505 }
506}