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 worktree.read(cx).id() != id_to_remove
100 } else {
101 false
102 }
103 });
104 }
105
106 pub fn set_worktrees_reordered(&mut self, worktrees_reordered: bool) {
107 self.worktrees_reordered = worktrees_reordered;
108 }
109
110 pub fn set_worktrees_from_proto(
111 &mut self,
112 worktrees: Vec<proto::WorktreeMetadata>,
113 replica_id: ReplicaId,
114 remote_id: u64,
115 client: AnyProtoClient,
116 cx: &mut ModelContext<Self>,
117 ) -> Result<()> {
118 let mut old_worktrees_by_id = self
119 .worktrees
120 .drain(..)
121 .filter_map(|worktree| {
122 let worktree = worktree.upgrade()?;
123 Some((worktree.read(cx).id(), worktree))
124 })
125 .collect::<HashMap<_, _>>();
126
127 for worktree in worktrees {
128 if let Some(old_worktree) =
129 old_worktrees_by_id.remove(&WorktreeId::from_proto(worktree.id))
130 {
131 self.worktrees.push(WorktreeHandle::Strong(old_worktree));
132 } else {
133 self.add(
134 &Worktree::remote(remote_id, replica_id, worktree, client.clone(), cx),
135 cx,
136 );
137 }
138 }
139
140 Ok(())
141 }
142
143 pub fn move_worktree(
144 &mut self,
145 source: WorktreeId,
146 destination: WorktreeId,
147 cx: &mut ModelContext<Self>,
148 ) -> Result<()> {
149 if source == destination {
150 return Ok(());
151 }
152
153 let mut source_index = None;
154 let mut destination_index = None;
155 for (i, worktree) in self.worktrees.iter().enumerate() {
156 if let Some(worktree) = worktree.upgrade() {
157 let worktree_id = worktree.read(cx).id();
158 if worktree_id == source {
159 source_index = Some(i);
160 if destination_index.is_some() {
161 break;
162 }
163 } else if worktree_id == destination {
164 destination_index = Some(i);
165 if source_index.is_some() {
166 break;
167 }
168 }
169 }
170 }
171
172 let source_index =
173 source_index.with_context(|| format!("Missing worktree for id {source}"))?;
174 let destination_index =
175 destination_index.with_context(|| format!("Missing worktree for id {destination}"))?;
176
177 if source_index == destination_index {
178 return Ok(());
179 }
180
181 let worktree_to_move = self.worktrees.remove(source_index);
182 self.worktrees.insert(destination_index, worktree_to_move);
183 self.worktrees_reordered = true;
184 cx.emit(WorktreeStoreEvent::WorktreeOrderChanged);
185 cx.notify();
186 Ok(())
187 }
188
189 pub fn disconnected_from_host(&mut self, cx: &mut AppContext) {
190 for worktree in &self.worktrees {
191 if let Some(worktree) = worktree.upgrade() {
192 worktree.update(cx, |worktree, _| {
193 if let Some(worktree) = worktree.as_remote_mut() {
194 worktree.disconnected_from_host();
195 }
196 });
197 }
198 }
199 }
200
201 pub fn set_shared(&mut self, is_shared: bool, cx: &mut ModelContext<Self>) {
202 self.is_shared = is_shared;
203
204 // When shared, retain all worktrees
205 if is_shared {
206 for worktree_handle in self.worktrees.iter_mut() {
207 match worktree_handle {
208 WorktreeHandle::Strong(_) => {}
209 WorktreeHandle::Weak(worktree) => {
210 if let Some(worktree) = worktree.upgrade() {
211 *worktree_handle = WorktreeHandle::Strong(worktree);
212 }
213 }
214 }
215 }
216 }
217 // When not shared, only retain the visible worktrees
218 else {
219 for worktree_handle in self.worktrees.iter_mut() {
220 if let WorktreeHandle::Strong(worktree) = worktree_handle {
221 let is_visible = worktree.update(cx, |worktree, _| {
222 worktree.stop_observing_updates();
223 worktree.is_visible()
224 });
225 if !is_visible {
226 *worktree_handle = WorktreeHandle::Weak(worktree.downgrade());
227 }
228 }
229 }
230 }
231 }
232
233 pub async fn handle_create_project_entry(
234 this: Model<Self>,
235 envelope: TypedEnvelope<proto::CreateProjectEntry>,
236 mut cx: AsyncAppContext,
237 ) -> Result<proto::ProjectEntryResponse> {
238 let worktree = this.update(&mut cx, |this, cx| {
239 let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id);
240 this.worktree_for_id(worktree_id, cx)
241 .ok_or_else(|| anyhow!("worktree not found"))
242 })??;
243 Worktree::handle_create_entry(worktree, envelope.payload, cx).await
244 }
245
246 pub async fn handle_rename_project_entry(
247 this: Model<Self>,
248 envelope: TypedEnvelope<proto::RenameProjectEntry>,
249 mut cx: AsyncAppContext,
250 ) -> Result<proto::ProjectEntryResponse> {
251 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
252 let worktree = this.update(&mut cx, |this, cx| {
253 this.worktree_for_entry(entry_id, cx)
254 .ok_or_else(|| anyhow!("worktree not found"))
255 })??;
256 Worktree::handle_rename_entry(worktree, envelope.payload, cx).await
257 }
258
259 pub async fn handle_copy_project_entry(
260 this: Model<Self>,
261 envelope: TypedEnvelope<proto::CopyProjectEntry>,
262 mut cx: AsyncAppContext,
263 ) -> Result<proto::ProjectEntryResponse> {
264 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
265 let worktree = this.update(&mut cx, |this, cx| {
266 this.worktree_for_entry(entry_id, cx)
267 .ok_or_else(|| anyhow!("worktree not found"))
268 })??;
269 Worktree::handle_copy_entry(worktree, envelope.payload, cx).await
270 }
271
272 pub async fn handle_delete_project_entry(
273 this: Model<Self>,
274 envelope: TypedEnvelope<proto::DeleteProjectEntry>,
275 mut cx: AsyncAppContext,
276 ) -> Result<proto::ProjectEntryResponse> {
277 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
278 let worktree = this.update(&mut cx, |this, cx| {
279 this.worktree_for_entry(entry_id, cx)
280 .ok_or_else(|| anyhow!("worktree not found"))
281 })??;
282 Worktree::handle_delete_entry(worktree, envelope.payload, cx).await
283 }
284
285 pub async fn handle_expand_project_entry(
286 this: Model<Self>,
287 envelope: TypedEnvelope<proto::ExpandProjectEntry>,
288 mut cx: AsyncAppContext,
289 ) -> Result<proto::ExpandProjectEntryResponse> {
290 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
291 let worktree = this
292 .update(&mut cx, |this, cx| this.worktree_for_entry(entry_id, cx))?
293 .ok_or_else(|| anyhow!("invalid request"))?;
294 Worktree::handle_expand_entry(worktree, envelope.payload, cx).await
295 }
296}
297
298#[derive(Clone)]
299enum WorktreeHandle {
300 Strong(Model<Worktree>),
301 Weak(WeakModel<Worktree>),
302}
303
304impl WorktreeHandle {
305 fn upgrade(&self) -> Option<Model<Worktree>> {
306 match self {
307 WorktreeHandle::Strong(handle) => Some(handle.clone()),
308 WorktreeHandle::Weak(handle) => handle.upgrade(),
309 }
310 }
311}