1use anyhow::{anyhow, Context as _, Result};
2use collections::HashMap;
3use gpui::{AppContext, AsyncAppContext, EntityId, EventEmitter, Model, ModelContext, WeakModel};
4use rpc::{
5 proto::{self, AnyProtoClient},
6 TypedEnvelope,
7};
8use text::ReplicaId;
9use worktree::{ProjectEntryId, Worktree, WorktreeId};
10
11pub struct WorktreeStore {
12 is_shared: bool,
13 worktrees: Vec<WorktreeHandle>,
14 worktrees_reordered: bool,
15}
16
17pub enum WorktreeStoreEvent {
18 WorktreeAdded(Model<Worktree>),
19 WorktreeRemoved(EntityId, WorktreeId),
20 WorktreeOrderChanged,
21}
22
23impl EventEmitter<WorktreeStoreEvent> for WorktreeStore {}
24
25impl WorktreeStore {
26 pub fn new(retain_worktrees: bool) -> Self {
27 Self {
28 is_shared: retain_worktrees,
29 worktrees: Vec::new(),
30 worktrees_reordered: false,
31 }
32 }
33
34 /// Iterates through all worktrees, including ones that don't appear in the project panel
35 pub fn worktrees(&self) -> impl '_ + DoubleEndedIterator<Item = Model<Worktree>> {
36 self.worktrees
37 .iter()
38 .filter_map(move |worktree| worktree.upgrade())
39 }
40
41 /// Iterates through all user-visible worktrees, the ones that appear in the project panel.
42 pub fn visible_worktrees<'a>(
43 &'a self,
44 cx: &'a AppContext,
45 ) -> impl 'a + DoubleEndedIterator<Item = Model<Worktree>> {
46 self.worktrees()
47 .filter(|worktree| worktree.read(cx).is_visible())
48 }
49
50 pub fn worktree_for_id(&self, id: WorktreeId, cx: &AppContext) -> Option<Model<Worktree>> {
51 self.worktrees()
52 .find(|worktree| worktree.read(cx).id() == id)
53 }
54
55 pub fn worktree_for_entry(
56 &self,
57 entry_id: ProjectEntryId,
58 cx: &AppContext,
59 ) -> Option<Model<Worktree>> {
60 self.worktrees()
61 .find(|worktree| worktree.read(cx).contains_entry(entry_id))
62 }
63
64 pub fn add(&mut self, worktree: &Model<Worktree>, cx: &mut ModelContext<Self>) {
65 let push_strong_handle = self.is_shared || worktree.read(cx).is_visible();
66 let handle = if push_strong_handle {
67 WorktreeHandle::Strong(worktree.clone())
68 } else {
69 WorktreeHandle::Weak(worktree.downgrade())
70 };
71 if self.worktrees_reordered {
72 self.worktrees.push(handle);
73 } else {
74 let i = match self
75 .worktrees
76 .binary_search_by_key(&Some(worktree.read(cx).abs_path()), |other| {
77 other.upgrade().map(|worktree| worktree.read(cx).abs_path())
78 }) {
79 Ok(i) | Err(i) => i,
80 };
81 self.worktrees.insert(i, handle);
82 }
83
84 cx.emit(WorktreeStoreEvent::WorktreeAdded(worktree.clone()));
85
86 let handle_id = worktree.entity_id();
87 cx.observe_release(worktree, move |_, worktree, cx| {
88 cx.emit(WorktreeStoreEvent::WorktreeRemoved(
89 handle_id,
90 worktree.id(),
91 ));
92 })
93 .detach();
94 }
95
96 pub fn remove_worktree(&mut self, id_to_remove: WorktreeId, cx: &mut ModelContext<Self>) {
97 self.worktrees.retain(|worktree| {
98 if let Some(worktree) = worktree.upgrade() {
99 if worktree.read(cx).id() == id_to_remove {
100 cx.emit(WorktreeStoreEvent::WorktreeRemoved(
101 worktree.entity_id(),
102 id_to_remove,
103 ));
104 false
105 } else {
106 true
107 }
108 } else {
109 false
110 }
111 });
112 }
113
114 pub fn set_worktrees_reordered(&mut self, worktrees_reordered: bool) {
115 self.worktrees_reordered = worktrees_reordered;
116 }
117
118 pub fn set_worktrees_from_proto(
119 &mut self,
120 worktrees: Vec<proto::WorktreeMetadata>,
121 replica_id: ReplicaId,
122 remote_id: u64,
123 client: AnyProtoClient,
124 cx: &mut ModelContext<Self>,
125 ) -> Result<()> {
126 let mut old_worktrees_by_id = self
127 .worktrees
128 .drain(..)
129 .filter_map(|worktree| {
130 let worktree = worktree.upgrade()?;
131 Some((worktree.read(cx).id(), worktree))
132 })
133 .collect::<HashMap<_, _>>();
134
135 for worktree in worktrees {
136 if let Some(old_worktree) =
137 old_worktrees_by_id.remove(&WorktreeId::from_proto(worktree.id))
138 {
139 self.worktrees.push(WorktreeHandle::Strong(old_worktree));
140 } else {
141 self.add(
142 &Worktree::remote(remote_id, replica_id, worktree, client.clone(), cx),
143 cx,
144 );
145 }
146 }
147
148 Ok(())
149 }
150
151 pub fn move_worktree(
152 &mut self,
153 source: WorktreeId,
154 destination: WorktreeId,
155 cx: &mut ModelContext<Self>,
156 ) -> Result<()> {
157 if source == destination {
158 return Ok(());
159 }
160
161 let mut source_index = None;
162 let mut destination_index = None;
163 for (i, worktree) in self.worktrees.iter().enumerate() {
164 if let Some(worktree) = worktree.upgrade() {
165 let worktree_id = worktree.read(cx).id();
166 if worktree_id == source {
167 source_index = Some(i);
168 if destination_index.is_some() {
169 break;
170 }
171 } else if worktree_id == destination {
172 destination_index = Some(i);
173 if source_index.is_some() {
174 break;
175 }
176 }
177 }
178 }
179
180 let source_index =
181 source_index.with_context(|| format!("Missing worktree for id {source}"))?;
182 let destination_index =
183 destination_index.with_context(|| format!("Missing worktree for id {destination}"))?;
184
185 if source_index == destination_index {
186 return Ok(());
187 }
188
189 let worktree_to_move = self.worktrees.remove(source_index);
190 self.worktrees.insert(destination_index, worktree_to_move);
191 self.worktrees_reordered = true;
192 cx.emit(WorktreeStoreEvent::WorktreeOrderChanged);
193 cx.notify();
194 Ok(())
195 }
196
197 pub fn disconnected_from_host(&mut self, cx: &mut AppContext) {
198 for worktree in &self.worktrees {
199 if let Some(worktree) = worktree.upgrade() {
200 worktree.update(cx, |worktree, _| {
201 if let Some(worktree) = worktree.as_remote_mut() {
202 worktree.disconnected_from_host();
203 }
204 });
205 }
206 }
207 }
208
209 pub fn set_shared(&mut self, is_shared: bool, cx: &mut ModelContext<Self>) {
210 self.is_shared = is_shared;
211
212 // When shared, retain all worktrees
213 if is_shared {
214 for worktree_handle in self.worktrees.iter_mut() {
215 match worktree_handle {
216 WorktreeHandle::Strong(_) => {}
217 WorktreeHandle::Weak(worktree) => {
218 if let Some(worktree) = worktree.upgrade() {
219 *worktree_handle = WorktreeHandle::Strong(worktree);
220 }
221 }
222 }
223 }
224 }
225 // When not shared, only retain the visible worktrees
226 else {
227 for worktree_handle in self.worktrees.iter_mut() {
228 if let WorktreeHandle::Strong(worktree) = worktree_handle {
229 let is_visible = worktree.update(cx, |worktree, _| {
230 worktree.stop_observing_updates();
231 worktree.is_visible()
232 });
233 if !is_visible {
234 *worktree_handle = WorktreeHandle::Weak(worktree.downgrade());
235 }
236 }
237 }
238 }
239 }
240
241 pub async fn handle_create_project_entry(
242 this: Model<Self>,
243 envelope: TypedEnvelope<proto::CreateProjectEntry>,
244 mut cx: AsyncAppContext,
245 ) -> Result<proto::ProjectEntryResponse> {
246 let worktree = this.update(&mut cx, |this, cx| {
247 let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id);
248 this.worktree_for_id(worktree_id, cx)
249 .ok_or_else(|| anyhow!("worktree not found"))
250 })??;
251 Worktree::handle_create_entry(worktree, envelope.payload, cx).await
252 }
253
254 pub async fn handle_rename_project_entry(
255 this: Model<Self>,
256 envelope: TypedEnvelope<proto::RenameProjectEntry>,
257 mut cx: AsyncAppContext,
258 ) -> Result<proto::ProjectEntryResponse> {
259 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
260 let worktree = this.update(&mut cx, |this, cx| {
261 this.worktree_for_entry(entry_id, cx)
262 .ok_or_else(|| anyhow!("worktree not found"))
263 })??;
264 Worktree::handle_rename_entry(worktree, envelope.payload, cx).await
265 }
266
267 pub async fn handle_copy_project_entry(
268 this: Model<Self>,
269 envelope: TypedEnvelope<proto::CopyProjectEntry>,
270 mut cx: AsyncAppContext,
271 ) -> Result<proto::ProjectEntryResponse> {
272 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
273 let worktree = this.update(&mut cx, |this, cx| {
274 this.worktree_for_entry(entry_id, cx)
275 .ok_or_else(|| anyhow!("worktree not found"))
276 })??;
277 Worktree::handle_copy_entry(worktree, envelope.payload, cx).await
278 }
279
280 pub async fn handle_delete_project_entry(
281 this: Model<Self>,
282 envelope: TypedEnvelope<proto::DeleteProjectEntry>,
283 mut cx: AsyncAppContext,
284 ) -> Result<proto::ProjectEntryResponse> {
285 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
286 let worktree = this.update(&mut cx, |this, cx| {
287 this.worktree_for_entry(entry_id, cx)
288 .ok_or_else(|| anyhow!("worktree not found"))
289 })??;
290 Worktree::handle_delete_entry(worktree, envelope.payload, cx).await
291 }
292
293 pub async fn handle_expand_project_entry(
294 this: Model<Self>,
295 envelope: TypedEnvelope<proto::ExpandProjectEntry>,
296 mut cx: AsyncAppContext,
297 ) -> Result<proto::ExpandProjectEntryResponse> {
298 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
299 let worktree = this
300 .update(&mut cx, |this, cx| this.worktree_for_entry(entry_id, cx))?
301 .ok_or_else(|| anyhow!("invalid request"))?;
302 Worktree::handle_expand_entry(worktree, envelope.payload, cx).await
303 }
304}
305
306#[derive(Clone)]
307enum WorktreeHandle {
308 Strong(Model<Worktree>),
309 Weak(WeakModel<Worktree>),
310}
311
312impl WorktreeHandle {
313 fn upgrade(&self) -> Option<Model<Worktree>> {
314 match self {
315 WorktreeHandle::Strong(handle) => Some(handle.clone()),
316 WorktreeHandle::Weak(handle) => handle.upgrade(),
317 }
318 }
319}