1use crate::{
2 worktree_store::{WorktreeStore, WorktreeStoreEvent},
3 Project, ProjectEntryId, ProjectPath,
4};
5use anyhow::{Context as _, Result};
6use collections::{hash_map, HashMap, HashSet};
7use futures::{channel::oneshot, StreamExt};
8use gpui::{
9 hash, prelude::*, AppContext, EventEmitter, Img, Model, ModelContext, Subscription, Task,
10 WeakModel,
11};
12use language::{DiskState, File};
13use rpc::{AnyProtoClient, ErrorExt as _};
14use std::ffi::OsStr;
15use std::num::NonZeroU64;
16use std::path::Path;
17use std::sync::Arc;
18use util::ResultExt;
19use worktree::{LoadedBinaryFile, PathChange, Worktree};
20
21#[derive(Clone, Copy, Debug, Hash, PartialEq, PartialOrd, Ord, Eq)]
22pub struct ImageId(NonZeroU64);
23
24impl std::fmt::Display for ImageId {
25 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26 write!(f, "{}", self.0)
27 }
28}
29
30impl From<NonZeroU64> for ImageId {
31 fn from(id: NonZeroU64) -> Self {
32 ImageId(id)
33 }
34}
35
36pub enum ImageItemEvent {
37 ReloadNeeded,
38 Reloaded,
39 FileHandleChanged,
40}
41
42impl EventEmitter<ImageItemEvent> for ImageItem {}
43
44pub enum ImageStoreEvent {
45 ImageAdded(Model<ImageItem>),
46}
47
48impl EventEmitter<ImageStoreEvent> for ImageStore {}
49
50pub struct ImageItem {
51 pub id: ImageId,
52 pub file: Arc<dyn File>,
53 pub image: Arc<gpui::Image>,
54 reload_task: Option<Task<()>>,
55}
56
57impl ImageItem {
58 pub fn project_path(&self, cx: &AppContext) -> ProjectPath {
59 ProjectPath {
60 worktree_id: self.file.worktree_id(cx),
61 path: self.file.path().clone(),
62 }
63 }
64
65 pub fn path(&self) -> &Arc<Path> {
66 self.file.path()
67 }
68
69 fn file_updated(&mut self, new_file: Arc<dyn File>, cx: &mut ModelContext<Self>) {
70 let mut file_changed = false;
71
72 let old_file = self.file.as_ref();
73 if new_file.path() != old_file.path() {
74 file_changed = true;
75 }
76
77 let old_state = old_file.disk_state();
78 let new_state = new_file.disk_state();
79 if old_state != new_state {
80 file_changed = true;
81 if matches!(new_state, DiskState::Present { .. }) {
82 cx.emit(ImageItemEvent::ReloadNeeded)
83 }
84 }
85
86 self.file = new_file;
87 if file_changed {
88 cx.emit(ImageItemEvent::FileHandleChanged);
89 cx.notify();
90 }
91 }
92
93 fn reload(&mut self, cx: &mut ModelContext<Self>) -> Option<oneshot::Receiver<()>> {
94 let local_file = self.file.as_local()?;
95 let (tx, rx) = futures::channel::oneshot::channel();
96
97 let content = local_file.load_bytes(cx);
98 self.reload_task = Some(cx.spawn(|this, mut cx| async move {
99 if let Some(image) = content
100 .await
101 .context("Failed to load image content")
102 .and_then(create_gpui_image)
103 .log_err()
104 {
105 this.update(&mut cx, |this, cx| {
106 this.image = image;
107 cx.emit(ImageItemEvent::Reloaded);
108 })
109 .log_err();
110 }
111 _ = tx.send(());
112 }));
113 Some(rx)
114 }
115}
116
117impl crate::Item for ImageItem {
118 fn try_open(
119 project: &Model<Project>,
120 path: &ProjectPath,
121 cx: &mut AppContext,
122 ) -> Option<Task<gpui::Result<Model<Self>>>> {
123 let path = path.clone();
124 let project = project.clone();
125
126 let ext = path
127 .path
128 .extension()
129 .and_then(OsStr::to_str)
130 .map(str::to_lowercase)
131 .unwrap_or_default();
132 let ext = ext.as_str();
133
134 // Only open the item if it's a binary image (no SVGs, etc.)
135 // Since we do not have a way to toggle to an editor
136 if Img::extensions().contains(&ext) && !ext.contains("svg") {
137 Some(cx.spawn(|mut cx| async move {
138 project
139 .update(&mut cx, |project, cx| project.open_image(path, cx))?
140 .await
141 }))
142 } else {
143 None
144 }
145 }
146
147 fn entry_id(&self, _: &AppContext) -> Option<ProjectEntryId> {
148 worktree::File::from_dyn(Some(&self.file))?.entry_id
149 }
150
151 fn project_path(&self, cx: &AppContext) -> Option<ProjectPath> {
152 Some(self.project_path(cx).clone())
153 }
154}
155
156trait ImageStoreImpl {
157 fn open_image(
158 &self,
159 path: Arc<Path>,
160 worktree: Model<Worktree>,
161 cx: &mut ModelContext<ImageStore>,
162 ) -> Task<Result<Model<ImageItem>>>;
163
164 fn reload_images(
165 &self,
166 images: HashSet<Model<ImageItem>>,
167 cx: &mut ModelContext<ImageStore>,
168 ) -> Task<Result<()>>;
169
170 fn as_local(&self) -> Option<Model<LocalImageStore>>;
171}
172
173struct RemoteImageStore {}
174
175struct LocalImageStore {
176 local_image_ids_by_path: HashMap<ProjectPath, ImageId>,
177 local_image_ids_by_entry_id: HashMap<ProjectEntryId, ImageId>,
178 image_store: WeakModel<ImageStore>,
179 _subscription: Subscription,
180}
181
182pub struct ImageStore {
183 state: Box<dyn ImageStoreImpl>,
184 opened_images: HashMap<ImageId, WeakModel<ImageItem>>,
185 worktree_store: Model<WorktreeStore>,
186 #[allow(clippy::type_complexity)]
187 loading_images_by_path: HashMap<
188 ProjectPath,
189 postage::watch::Receiver<Option<Result<Model<ImageItem>, Arc<anyhow::Error>>>>,
190 >,
191}
192
193impl ImageStore {
194 pub fn local(worktree_store: Model<WorktreeStore>, cx: &mut ModelContext<Self>) -> Self {
195 let this = cx.weak_model();
196 Self {
197 state: Box::new(cx.new_model(|cx| {
198 let subscription = cx.subscribe(
199 &worktree_store,
200 |this: &mut LocalImageStore, _, event, cx| {
201 if let WorktreeStoreEvent::WorktreeAdded(worktree) = event {
202 this.subscribe_to_worktree(worktree, cx);
203 }
204 },
205 );
206
207 LocalImageStore {
208 local_image_ids_by_path: Default::default(),
209 local_image_ids_by_entry_id: Default::default(),
210 image_store: this,
211 _subscription: subscription,
212 }
213 })),
214 opened_images: Default::default(),
215 loading_images_by_path: Default::default(),
216 worktree_store,
217 }
218 }
219
220 pub fn remote(
221 worktree_store: Model<WorktreeStore>,
222 _upstream_client: AnyProtoClient,
223 _remote_id: u64,
224 cx: &mut ModelContext<Self>,
225 ) -> Self {
226 Self {
227 state: Box::new(cx.new_model(|_| RemoteImageStore {})),
228 opened_images: Default::default(),
229 loading_images_by_path: Default::default(),
230 worktree_store,
231 }
232 }
233
234 pub fn images(&self) -> impl '_ + Iterator<Item = Model<ImageItem>> {
235 self.opened_images
236 .values()
237 .filter_map(|image| image.upgrade())
238 }
239
240 pub fn get(&self, image_id: ImageId) -> Option<Model<ImageItem>> {
241 self.opened_images
242 .get(&image_id)
243 .and_then(|image| image.upgrade())
244 }
245
246 pub fn get_by_path(&self, path: &ProjectPath, cx: &AppContext) -> Option<Model<ImageItem>> {
247 self.images()
248 .find(|image| &image.read(cx).project_path(cx) == path)
249 }
250
251 pub fn open_image(
252 &mut self,
253 project_path: ProjectPath,
254 cx: &mut ModelContext<Self>,
255 ) -> Task<Result<Model<ImageItem>>> {
256 let existing_image = self.get_by_path(&project_path, cx);
257 if let Some(existing_image) = existing_image {
258 return Task::ready(Ok(existing_image));
259 }
260
261 let Some(worktree) = self
262 .worktree_store
263 .read(cx)
264 .worktree_for_id(project_path.worktree_id, cx)
265 else {
266 return Task::ready(Err(anyhow::anyhow!("no such worktree")));
267 };
268
269 let loading_watch = match self.loading_images_by_path.entry(project_path.clone()) {
270 // If the given path is already being loaded, then wait for that existing
271 // task to complete and return the same image.
272 hash_map::Entry::Occupied(e) => e.get().clone(),
273
274 // Otherwise, record the fact that this path is now being loaded.
275 hash_map::Entry::Vacant(entry) => {
276 let (mut tx, rx) = postage::watch::channel();
277 entry.insert(rx.clone());
278
279 let project_path = project_path.clone();
280 let load_image = self
281 .state
282 .open_image(project_path.path.clone(), worktree, cx);
283
284 cx.spawn(move |this, mut cx| async move {
285 let load_result = load_image.await;
286 *tx.borrow_mut() = Some(this.update(&mut cx, |this, _cx| {
287 // Record the fact that the image is no longer loading.
288 this.loading_images_by_path.remove(&project_path);
289 let image = load_result.map_err(Arc::new)?;
290 Ok(image)
291 })?);
292 anyhow::Ok(())
293 })
294 .detach();
295 rx
296 }
297 };
298
299 cx.background_executor().spawn(async move {
300 Self::wait_for_loading_image(loading_watch)
301 .await
302 .map_err(|e| e.cloned())
303 })
304 }
305
306 pub async fn wait_for_loading_image(
307 mut receiver: postage::watch::Receiver<
308 Option<Result<Model<ImageItem>, Arc<anyhow::Error>>>,
309 >,
310 ) -> Result<Model<ImageItem>, Arc<anyhow::Error>> {
311 loop {
312 if let Some(result) = receiver.borrow().as_ref() {
313 match result {
314 Ok(image) => return Ok(image.to_owned()),
315 Err(e) => return Err(e.to_owned()),
316 }
317 }
318 receiver.next().await;
319 }
320 }
321
322 pub fn reload_images(
323 &self,
324 images: HashSet<Model<ImageItem>>,
325 cx: &mut ModelContext<ImageStore>,
326 ) -> Task<Result<()>> {
327 if images.is_empty() {
328 return Task::ready(Ok(()));
329 }
330
331 self.state.reload_images(images, cx)
332 }
333
334 fn add_image(
335 &mut self,
336 image: Model<ImageItem>,
337 cx: &mut ModelContext<ImageStore>,
338 ) -> Result<()> {
339 let image_id = image.read(cx).id;
340
341 self.opened_images.insert(image_id, image.downgrade());
342
343 cx.subscribe(&image, Self::on_image_event).detach();
344 cx.emit(ImageStoreEvent::ImageAdded(image));
345 Ok(())
346 }
347
348 fn on_image_event(
349 &mut self,
350 image: Model<ImageItem>,
351 event: &ImageItemEvent,
352 cx: &mut ModelContext<Self>,
353 ) {
354 match event {
355 ImageItemEvent::FileHandleChanged => {
356 if let Some(local) = self.state.as_local() {
357 local.update(cx, |local, cx| {
358 local.image_changed_file(image, cx);
359 })
360 }
361 }
362 _ => {}
363 }
364 }
365}
366
367impl ImageStoreImpl for Model<LocalImageStore> {
368 fn open_image(
369 &self,
370 path: Arc<Path>,
371 worktree: Model<Worktree>,
372 cx: &mut ModelContext<ImageStore>,
373 ) -> Task<Result<Model<ImageItem>>> {
374 let this = self.clone();
375
376 let load_file = worktree.update(cx, |worktree, cx| {
377 worktree.load_binary_file(path.as_ref(), cx)
378 });
379 cx.spawn(move |image_store, mut cx| async move {
380 let LoadedBinaryFile { file, content } = load_file.await?;
381 let image = create_gpui_image(content)?;
382
383 let model = cx.new_model(|cx| ImageItem {
384 id: cx.entity_id().as_non_zero_u64().into(),
385 file: file.clone(),
386 image,
387 reload_task: None,
388 })?;
389
390 let image_id = cx.read_model(&model, |model, _| model.id)?;
391
392 this.update(&mut cx, |this, cx| {
393 image_store.update(cx, |image_store, cx| {
394 image_store.add_image(model.clone(), cx)
395 })??;
396 this.local_image_ids_by_path.insert(
397 ProjectPath {
398 worktree_id: file.worktree_id(cx),
399 path: file.path.clone(),
400 },
401 image_id,
402 );
403
404 if let Some(entry_id) = file.entry_id {
405 this.local_image_ids_by_entry_id.insert(entry_id, image_id);
406 }
407
408 anyhow::Ok(())
409 })??;
410
411 Ok(model)
412 })
413 }
414
415 fn reload_images(
416 &self,
417 images: HashSet<Model<ImageItem>>,
418 cx: &mut ModelContext<ImageStore>,
419 ) -> Task<Result<()>> {
420 cx.spawn(move |_, mut cx| async move {
421 for image in images {
422 if let Some(rec) = image.update(&mut cx, |image, cx| image.reload(cx))? {
423 rec.await?
424 }
425 }
426 Ok(())
427 })
428 }
429
430 fn as_local(&self) -> Option<Model<LocalImageStore>> {
431 Some(self.clone())
432 }
433}
434
435impl LocalImageStore {
436 fn subscribe_to_worktree(&mut self, worktree: &Model<Worktree>, cx: &mut ModelContext<Self>) {
437 cx.subscribe(worktree, |this, worktree, event, cx| {
438 if worktree.read(cx).is_local() {
439 match event {
440 worktree::Event::UpdatedEntries(changes) => {
441 this.local_worktree_entries_changed(&worktree, changes, cx);
442 }
443 _ => {}
444 }
445 }
446 })
447 .detach();
448 }
449
450 fn local_worktree_entries_changed(
451 &mut self,
452 worktree_handle: &Model<Worktree>,
453 changes: &[(Arc<Path>, ProjectEntryId, PathChange)],
454 cx: &mut ModelContext<Self>,
455 ) {
456 let snapshot = worktree_handle.read(cx).snapshot();
457 for (path, entry_id, _) in changes {
458 self.local_worktree_entry_changed(*entry_id, path, worktree_handle, &snapshot, cx);
459 }
460 }
461
462 fn local_worktree_entry_changed(
463 &mut self,
464 entry_id: ProjectEntryId,
465 path: &Arc<Path>,
466 worktree: &Model<worktree::Worktree>,
467 snapshot: &worktree::Snapshot,
468 cx: &mut ModelContext<Self>,
469 ) -> Option<()> {
470 let project_path = ProjectPath {
471 worktree_id: snapshot.id(),
472 path: path.clone(),
473 };
474 let image_id = match self.local_image_ids_by_entry_id.get(&entry_id) {
475 Some(&image_id) => image_id,
476 None => self.local_image_ids_by_path.get(&project_path).copied()?,
477 };
478
479 let image = self
480 .image_store
481 .update(cx, |image_store, _| {
482 if let Some(image) = image_store.get(image_id) {
483 Some(image)
484 } else {
485 image_store.opened_images.remove(&image_id);
486 None
487 }
488 })
489 .ok()
490 .flatten();
491 let image = if let Some(image) = image {
492 image
493 } else {
494 self.local_image_ids_by_path.remove(&project_path);
495 self.local_image_ids_by_entry_id.remove(&entry_id);
496 return None;
497 };
498
499 image.update(cx, |image, cx| {
500 let Some(old_file) = worktree::File::from_dyn(Some(&image.file)) else {
501 return;
502 };
503 if old_file.worktree != *worktree {
504 return;
505 }
506
507 let snapshot_entry = old_file
508 .entry_id
509 .and_then(|entry_id| snapshot.entry_for_id(entry_id))
510 .or_else(|| snapshot.entry_for_path(old_file.path.as_ref()));
511
512 let new_file = if let Some(entry) = snapshot_entry {
513 worktree::File {
514 disk_state: match entry.mtime {
515 Some(mtime) => DiskState::Present { mtime },
516 None => old_file.disk_state,
517 },
518 is_local: true,
519 entry_id: Some(entry.id),
520 path: entry.path.clone(),
521 worktree: worktree.clone(),
522 is_private: entry.is_private,
523 }
524 } else {
525 worktree::File {
526 disk_state: DiskState::Deleted,
527 is_local: true,
528 entry_id: old_file.entry_id,
529 path: old_file.path.clone(),
530 worktree: worktree.clone(),
531 is_private: old_file.is_private,
532 }
533 };
534
535 if new_file == *old_file {
536 return;
537 }
538
539 if new_file.path != old_file.path {
540 self.local_image_ids_by_path.remove(&ProjectPath {
541 path: old_file.path.clone(),
542 worktree_id: old_file.worktree_id(cx),
543 });
544 self.local_image_ids_by_path.insert(
545 ProjectPath {
546 worktree_id: new_file.worktree_id(cx),
547 path: new_file.path.clone(),
548 },
549 image_id,
550 );
551 }
552
553 if new_file.entry_id != old_file.entry_id {
554 if let Some(entry_id) = old_file.entry_id {
555 self.local_image_ids_by_entry_id.remove(&entry_id);
556 }
557 if let Some(entry_id) = new_file.entry_id {
558 self.local_image_ids_by_entry_id.insert(entry_id, image_id);
559 }
560 }
561
562 image.file_updated(Arc::new(new_file), cx);
563 });
564 None
565 }
566
567 fn image_changed_file(&mut self, image: Model<ImageItem>, cx: &mut AppContext) -> Option<()> {
568 let file = worktree::File::from_dyn(Some(&image.read(cx).file))?;
569
570 let image_id = image.read(cx).id;
571 if let Some(entry_id) = file.entry_id {
572 match self.local_image_ids_by_entry_id.get(&entry_id) {
573 Some(_) => {
574 return None;
575 }
576 None => {
577 self.local_image_ids_by_entry_id.insert(entry_id, image_id);
578 }
579 }
580 };
581 self.local_image_ids_by_path.insert(
582 ProjectPath {
583 worktree_id: file.worktree_id(cx),
584 path: file.path.clone(),
585 },
586 image_id,
587 );
588
589 Some(())
590 }
591}
592
593fn create_gpui_image(content: Vec<u8>) -> anyhow::Result<Arc<gpui::Image>> {
594 let format = image::guess_format(&content)?;
595
596 Ok(Arc::new(gpui::Image {
597 id: hash(&content),
598 format: match format {
599 image::ImageFormat::Png => gpui::ImageFormat::Png,
600 image::ImageFormat::Jpeg => gpui::ImageFormat::Jpeg,
601 image::ImageFormat::WebP => gpui::ImageFormat::Webp,
602 image::ImageFormat::Gif => gpui::ImageFormat::Gif,
603 image::ImageFormat::Bmp => gpui::ImageFormat::Bmp,
604 image::ImageFormat::Tiff => gpui::ImageFormat::Tiff,
605 _ => Err(anyhow::anyhow!("Image format not supported"))?,
606 },
607 bytes: content,
608 }))
609}
610
611impl ImageStoreImpl for Model<RemoteImageStore> {
612 fn open_image(
613 &self,
614 _path: Arc<Path>,
615 _worktree: Model<Worktree>,
616 _cx: &mut ModelContext<ImageStore>,
617 ) -> Task<Result<Model<ImageItem>>> {
618 Task::ready(Err(anyhow::anyhow!(
619 "Opening images from remote is not supported"
620 )))
621 }
622
623 fn reload_images(
624 &self,
625 _images: HashSet<Model<ImageItem>>,
626 _cx: &mut ModelContext<ImageStore>,
627 ) -> Task<Result<()>> {
628 Task::ready(Err(anyhow::anyhow!(
629 "Reloading images from remote is not supported"
630 )))
631 }
632
633 fn as_local(&self) -> Option<Model<LocalImageStore>> {
634 None
635 }
636}
637
638#[cfg(test)]
639mod tests {
640 use super::*;
641 use fs::FakeFs;
642 use gpui::TestAppContext;
643 use serde_json::json;
644 use settings::SettingsStore;
645 use std::path::PathBuf;
646
647 pub fn init_test(cx: &mut TestAppContext) {
648 if std::env::var("RUST_LOG").is_ok() {
649 env_logger::try_init().ok();
650 }
651
652 cx.update(|cx| {
653 let settings_store = SettingsStore::test(cx);
654 cx.set_global(settings_store);
655 language::init(cx);
656 Project::init_settings(cx);
657 });
658 }
659
660 #[gpui::test]
661 async fn test_image_not_loaded_twice(cx: &mut TestAppContext) {
662 init_test(cx);
663 let fs = FakeFs::new(cx.executor());
664
665 fs.insert_tree("/root", json!({})).await;
666 // Create a png file that consists of a single white pixel
667 fs.insert_file(
668 "/root/image_1.png",
669 vec![
670 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48,
671 0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00,
672 0x00, 0x1F, 0x15, 0xC4, 0x89, 0x00, 0x00, 0x00, 0x0A, 0x49, 0x44, 0x41, 0x54, 0x78,
673 0x9C, 0x63, 0x00, 0x01, 0x00, 0x00, 0x05, 0x00, 0x01, 0x0D, 0x0A, 0x2D, 0xB4, 0x00,
674 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82,
675 ],
676 )
677 .await;
678
679 let project = Project::test(fs, ["/root".as_ref()], cx).await;
680
681 let worktree_id =
682 cx.update(|cx| project.read(cx).worktrees(cx).next().unwrap().read(cx).id());
683
684 let project_path = ProjectPath {
685 worktree_id,
686 path: PathBuf::from("image_1.png").into(),
687 };
688
689 let (task1, task2) = project.update(cx, |project, cx| {
690 (
691 project.open_image(project_path.clone(), cx),
692 project.open_image(project_path.clone(), cx),
693 )
694 });
695
696 let image1 = task1.await.unwrap();
697 let image2 = task2.await.unwrap();
698
699 assert_eq!(image1, image2);
700 }
701}