1use anyhow::{anyhow, Result};
2use gpui::{AppContext, Model, Task, WeakModel};
3use project::{Fs, Project, ProjectPath, Worktree};
4use std::{cmp::Ordering, fmt::Write as _, ops::Range, sync::Arc};
5use sum_tree::TreeMap;
6
7pub struct ProjectContext {
8 files: TreeMap<ProjectPath, PathState>,
9 project: WeakModel<Project>,
10 fs: Arc<dyn Fs>,
11}
12
13#[derive(Debug, Clone)]
14enum PathState {
15 PathOnly,
16 EntireFile,
17 Excerpts { ranges: Vec<Range<usize>> },
18}
19
20impl ProjectContext {
21 pub fn new(project: WeakModel<Project>, fs: Arc<dyn Fs>) -> Self {
22 Self {
23 files: TreeMap::default(),
24 fs,
25 project,
26 }
27 }
28
29 pub fn add_path(&mut self, project_path: ProjectPath) {
30 if self.files.get(&project_path).is_none() {
31 self.files.insert(project_path, PathState::PathOnly);
32 }
33 }
34
35 pub fn add_excerpts(&mut self, project_path: ProjectPath, new_ranges: &[Range<usize>]) {
36 let previous_state = self
37 .files
38 .get(&project_path)
39 .unwrap_or(&PathState::PathOnly);
40
41 let mut ranges = match previous_state {
42 PathState::EntireFile => return,
43 PathState::PathOnly => Vec::new(),
44 PathState::Excerpts { ranges } => ranges.to_vec(),
45 };
46
47 for new_range in new_ranges {
48 let ix = ranges.binary_search_by(|probe| {
49 if probe.end < new_range.start {
50 Ordering::Less
51 } else if probe.start > new_range.end {
52 Ordering::Greater
53 } else {
54 Ordering::Equal
55 }
56 });
57
58 match ix {
59 Ok(mut ix) => {
60 let existing = &mut ranges[ix];
61 existing.start = existing.start.min(new_range.start);
62 existing.end = existing.end.max(new_range.end);
63 while ix + 1 < ranges.len() && ranges[ix + 1].start <= ranges[ix].end {
64 ranges[ix].end = ranges[ix].end.max(ranges[ix + 1].end);
65 ranges.remove(ix + 1);
66 }
67 while ix > 0 && ranges[ix - 1].end >= ranges[ix].start {
68 ranges[ix].start = ranges[ix].start.min(ranges[ix - 1].start);
69 ranges.remove(ix - 1);
70 ix -= 1;
71 }
72 }
73 Err(ix) => {
74 ranges.insert(ix, new_range.clone());
75 }
76 }
77 }
78
79 self.files
80 .insert(project_path, PathState::Excerpts { ranges });
81 }
82
83 pub fn add_file(&mut self, project_path: ProjectPath) {
84 self.files.insert(project_path, PathState::EntireFile);
85 }
86
87 pub fn generate_system_message(&self, cx: &mut AppContext) -> Task<Result<String>> {
88 let project = self
89 .project
90 .upgrade()
91 .ok_or_else(|| anyhow!("project dropped"));
92 let files = self.files.clone();
93 let fs = self.fs.clone();
94 cx.spawn(|cx| async move {
95 let project = project?;
96 let mut result = "project structure:\n".to_string();
97
98 let mut last_worktree: Option<Model<Worktree>> = None;
99 for (project_path, path_state) in files.iter() {
100 if let Some(worktree) = &last_worktree {
101 if worktree.read_with(&cx, |tree, _| tree.id())? != project_path.worktree_id {
102 last_worktree = None;
103 }
104 }
105
106 let worktree;
107 if let Some(last_worktree) = &last_worktree {
108 worktree = last_worktree.clone();
109 } else if let Some(tree) = project.read_with(&cx, |project, cx| {
110 project.worktree_for_id(project_path.worktree_id, cx)
111 })? {
112 worktree = tree;
113 last_worktree = Some(worktree.clone());
114 let worktree_name =
115 worktree.read_with(&cx, |tree, _cx| tree.root_name().to_string())?;
116 writeln!(&mut result, "# {}", worktree_name).unwrap();
117 } else {
118 continue;
119 }
120
121 let worktree_abs_path = worktree.read_with(&cx, |tree, _cx| tree.abs_path())?;
122 let path = &project_path.path;
123 writeln!(&mut result, "## {}", path.display()).unwrap();
124
125 match path_state {
126 PathState::PathOnly => {}
127 PathState::EntireFile => {
128 let text = fs.load(&worktree_abs_path.join(&path)).await?;
129 writeln!(&mut result, "~~~\n{text}\n~~~").unwrap();
130 }
131 PathState::Excerpts { ranges } => {
132 let text = fs.load(&worktree_abs_path.join(&path)).await?;
133
134 writeln!(&mut result, "~~~").unwrap();
135
136 // Assumption: ranges are in order, not overlapping
137 let mut prev_range_end = 0;
138 for range in ranges {
139 if range.start > prev_range_end {
140 writeln!(&mut result, "...").unwrap();
141 prev_range_end = range.end;
142 }
143
144 let mut start = range.start;
145 let mut end = range.end.min(text.len());
146 while !text.is_char_boundary(start) {
147 start += 1;
148 }
149 while !text.is_char_boundary(end) {
150 end -= 1;
151 }
152 result.push_str(&text[start..end]);
153 if !result.ends_with('\n') {
154 result.push('\n');
155 }
156 }
157
158 if prev_range_end < text.len() {
159 writeln!(&mut result, "...").unwrap();
160 }
161
162 writeln!(&mut result, "~~~").unwrap();
163 }
164 }
165 }
166 Ok(result)
167 })
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use std::path::Path;
174
175 use super::*;
176 use gpui::TestAppContext;
177 use project::FakeFs;
178 use serde_json::json;
179 use settings::SettingsStore;
180
181 use unindent::Unindent as _;
182
183 #[gpui::test]
184 async fn test_system_message_generation(cx: &mut TestAppContext) {
185 init_test(cx);
186
187 let file_3_contents = r#"
188 fn test1() {}
189 fn test2() {}
190 fn test3() {}
191 "#
192 .unindent();
193
194 let fs = FakeFs::new(cx.executor());
195 fs.insert_tree(
196 "/code",
197 json!({
198 "root1": {
199 "lib": {
200 "file1.rs": "mod example;",
201 "file2.rs": "",
202 },
203 "test": {
204 "file3.rs": file_3_contents,
205 }
206 },
207 "root2": {
208 "src": {
209 "main.rs": ""
210 }
211 }
212 }),
213 )
214 .await;
215
216 let project = Project::test(
217 fs.clone(),
218 ["/code/root1".as_ref(), "/code/root2".as_ref()],
219 cx,
220 )
221 .await;
222
223 let worktree_ids = project.read_with(cx, |project, cx| {
224 project
225 .worktrees()
226 .map(|worktree| worktree.read(cx).id())
227 .collect::<Vec<_>>()
228 });
229
230 let mut ax = ProjectContext::new(project.downgrade(), fs);
231
232 ax.add_file(ProjectPath {
233 worktree_id: worktree_ids[0],
234 path: Path::new("lib/file1.rs").into(),
235 });
236
237 let message = cx
238 .update(|cx| ax.generate_system_message(cx))
239 .await
240 .unwrap();
241 assert_eq!(
242 r#"
243 project structure:
244 # root1
245 ## lib/file1.rs
246 ~~~
247 mod example;
248 ~~~
249 "#
250 .unindent(),
251 message
252 );
253
254 ax.add_excerpts(
255 ProjectPath {
256 worktree_id: worktree_ids[0],
257 path: Path::new("test/file3.rs").into(),
258 },
259 &[
260 file_3_contents.find("fn test2").unwrap()
261 ..file_3_contents.find("fn test3").unwrap(),
262 ],
263 );
264
265 let message = cx
266 .update(|cx| ax.generate_system_message(cx))
267 .await
268 .unwrap();
269 assert_eq!(
270 r#"
271 project structure:
272 # root1
273 ## lib/file1.rs
274 ~~~
275 mod example;
276 ~~~
277 ## test/file3.rs
278 ~~~
279 ...
280 fn test2() {}
281 ...
282 ~~~
283 "#
284 .unindent(),
285 message
286 );
287 }
288
289 fn init_test(cx: &mut TestAppContext) {
290 cx.update(|cx| {
291 let settings_store = SettingsStore::test(cx);
292 cx.set_global(settings_store);
293 Project::init_settings(cx);
294 });
295 }
296}