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