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