1use std::ops::Range;
2use std::sync::Arc;
3use std::time::Duration;
4
5use anyhow::Context as _;
6use clock::Global;
7use collections::HashMap;
8use futures::FutureExt as _;
9use futures::future::{Shared, join_all};
10use gpui::{AppContext as _, Context, Entity, SharedString, Task};
11use itertools::Itertools;
12use language::Buffer;
13use lsp::LanguageServerId;
14use settings::Settings as _;
15use text::Anchor;
16
17use crate::lsp_command::{GetFoldingRanges, LspCommand as _};
18use crate::lsp_store::LspStore;
19use crate::project_settings::ProjectSettings;
20
21#[derive(Clone, Debug)]
22pub struct LspFoldingRange {
23 pub range: Range<Anchor>,
24 pub collapsed_text: Option<SharedString>,
25}
26
27pub(super) type FoldingRangeTask =
28 Shared<Task<std::result::Result<Vec<LspFoldingRange>, Arc<anyhow::Error>>>>;
29
30#[derive(Debug, Default)]
31pub(super) struct FoldingRangeData {
32 pub(super) ranges: HashMap<LanguageServerId, Vec<LspFoldingRange>>,
33 ranges_update: Option<(Global, FoldingRangeTask)>,
34}
35
36impl LspStore {
37 /// Returns a task that resolves to the folding ranges for the given buffer.
38 ///
39 /// Caches results per buffer version so repeated calls for the same version
40 /// return immediately. Deduplicates concurrent in-flight requests.
41 pub fn fetch_folding_ranges(
42 &mut self,
43 buffer: &Entity<Buffer>,
44 cx: &mut Context<Self>,
45 ) -> Task<Vec<LspFoldingRange>> {
46 let version_queried_for = buffer.read(cx).version();
47 let buffer_id = buffer.read(cx).remote_id();
48
49 let current_language_servers = self.as_local().map(|local| {
50 local
51 .buffers_opened_in_servers
52 .get(&buffer_id)
53 .cloned()
54 .unwrap_or_default()
55 });
56
57 if let Some(lsp_data) = self.current_lsp_data(buffer_id) {
58 if let Some(cached) = &lsp_data.folding_ranges {
59 if !version_queried_for.changed_since(&lsp_data.buffer_version) {
60 let has_different_servers =
61 current_language_servers.is_some_and(|current_language_servers| {
62 current_language_servers != cached.ranges.keys().copied().collect()
63 });
64 if !has_different_servers {
65 let snapshot = buffer.read(cx).snapshot();
66 return Task::ready(
67 cached
68 .ranges
69 .values()
70 .flatten()
71 .cloned()
72 .sorted_by(|a, b| a.range.start.cmp(&b.range.start, &snapshot))
73 .collect(),
74 );
75 }
76 }
77 }
78 }
79
80 let folding_lsp_data = self
81 .latest_lsp_data(buffer, cx)
82 .folding_ranges
83 .get_or_insert_default();
84 if let Some((updating_for, running_update)) = &folding_lsp_data.ranges_update {
85 if !version_queried_for.changed_since(updating_for) {
86 let running = running_update.clone();
87 return cx.background_spawn(async move { running.await.unwrap_or_default() });
88 }
89 }
90
91 let buffer = buffer.clone();
92 let query_version = version_queried_for.clone();
93 let new_task = cx
94 .spawn(async move |lsp_store, cx| {
95 cx.background_executor()
96 .timer(Duration::from_millis(30))
97 .await;
98
99 let fetched = lsp_store
100 .update(cx, |lsp_store, cx| {
101 lsp_store.fetch_folding_ranges_for_buffer(&buffer, cx)
102 })
103 .map_err(Arc::new)?
104 .await
105 .context("fetching folding ranges")
106 .map_err(Arc::new);
107
108 let fetched = match fetched {
109 Ok(fetched) => fetched,
110 Err(e) => {
111 lsp_store
112 .update(cx, |lsp_store, _| {
113 if let Some(lsp_data) = lsp_store.lsp_data.get_mut(&buffer_id) {
114 if let Some(folding_ranges) = &mut lsp_data.folding_ranges {
115 folding_ranges.ranges_update = None;
116 }
117 }
118 })
119 .ok();
120 return Err(e);
121 }
122 };
123
124 lsp_store
125 .update(cx, |lsp_store, cx| {
126 let lsp_data = lsp_store.latest_lsp_data(&buffer, cx);
127 let folding = lsp_data.folding_ranges.get_or_insert_default();
128
129 if let Some(fetched_ranges) = fetched {
130 if lsp_data.buffer_version == query_version {
131 folding.ranges.extend(fetched_ranges);
132 } else if !lsp_data.buffer_version.changed_since(&query_version) {
133 lsp_data.buffer_version = query_version;
134 folding.ranges = fetched_ranges;
135 }
136 }
137 folding.ranges_update = None;
138 let snapshot = buffer.read(cx).snapshot();
139 folding
140 .ranges
141 .values()
142 .flatten()
143 .cloned()
144 .sorted_by(|a, b| a.range.start.cmp(&b.range.start, &snapshot))
145 .collect()
146 })
147 .map_err(Arc::new)
148 })
149 .shared();
150
151 folding_lsp_data.ranges_update = Some((version_queried_for, new_task.clone()));
152
153 cx.background_spawn(async move { new_task.await.unwrap_or_default() })
154 }
155
156 fn fetch_folding_ranges_for_buffer(
157 &mut self,
158 buffer: &Entity<Buffer>,
159 cx: &mut Context<Self>,
160 ) -> Task<anyhow::Result<Option<HashMap<LanguageServerId, Vec<LspFoldingRange>>>>> {
161 if let Some((client, project_id)) = self.upstream_client() {
162 let request = GetFoldingRanges;
163 if !self.is_capable_for_proto_request(buffer, &request, cx) {
164 return Task::ready(Ok(None));
165 }
166
167 let request_timeout = ProjectSettings::get_global(cx)
168 .global_lsp_settings
169 .get_request_timeout();
170 let request_task = client.request_lsp(
171 project_id,
172 None,
173 request_timeout,
174 cx.background_executor().clone(),
175 request.to_proto(project_id, buffer.read(cx)),
176 );
177 let buffer = buffer.clone();
178 cx.spawn(async move |weak_lsp_store, cx| {
179 let Some(lsp_store) = weak_lsp_store.upgrade() else {
180 return Ok(None);
181 };
182 let Some(responses) = request_task.await? else {
183 return Ok(None);
184 };
185
186 let folding_ranges = join_all(responses.payload.into_iter().map(|response| {
187 let lsp_store = lsp_store.clone();
188 let buffer = buffer.clone();
189 let cx = cx.clone();
190 async move {
191 (
192 LanguageServerId::from_proto(response.server_id),
193 GetFoldingRanges
194 .response_from_proto(response.response, lsp_store, buffer, cx)
195 .await,
196 )
197 }
198 }))
199 .await;
200
201 let mut has_errors = false;
202 let result = folding_ranges
203 .into_iter()
204 .filter_map(|(server_id, ranges)| match ranges {
205 Ok(ranges) => Some((server_id, ranges)),
206 Err(e) => {
207 has_errors = true;
208 log::error!("Failed to fetch folding ranges: {e:#}");
209 None
210 }
211 })
212 .collect::<HashMap<_, _>>();
213 anyhow::ensure!(
214 !has_errors || !result.is_empty(),
215 "Failed to fetch folding ranges"
216 );
217 Ok(Some(result))
218 })
219 } else {
220 let folding_task =
221 self.request_multiple_lsp_locally(buffer, None::<usize>, GetFoldingRanges, cx);
222 cx.background_spawn(async move { Ok(Some(folding_task.await.into_iter().collect())) })
223 }
224 }
225}