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