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