comfy_core/
asset_loader.rs

1use image::RgbaImage;
2
3use crate::*;
4use std::sync::mpsc::{Receiver, Sender};
5
6pub fn init_asset_source(
7    dir: &'static include_dir::Dir<'static>,
8    base_path: fn(&str) -> String,
9) {
10    ASSETS.borrow_mut().asset_loader.asset_source =
11        Some(AssetSource { dir, base_path });
12}
13
14pub struct AssetLoader {
15    #[cfg(not(target_arch = "wasm32"))]
16    pub thread_pool: rayon::ThreadPool,
17    pub wgpu_load_queue: Arc<Mutex<Option<TextureLoadQueue>>>,
18
19    pub texture_load_queue: Vec<(String, String)>,
20    pub sound_load_queue: Vec<(String, String)>,
21
22    pub sounds: Arc<Mutex<HashMap<Sound, StaticSoundData>>>,
23    pub sound_send: Arc<Mutex<Sender<SoundAssetData>>>,
24    pub sound_recv: Arc<Mutex<Receiver<SoundAssetData>>>,
25
26    pub texture_data_send: Arc<Mutex<Sender<TextureAssetData>>>,
27    pub texture_data_recv: Arc<Mutex<Receiver<TextureAssetData>>>,
28
29    pub asset_source: Option<AssetSource>,
30
31    pub data_load_queue: Vec<AssetData>,
32
33    pub pending_textures: HashSet<String>,
34}
35
36impl AssetLoader {
37    pub fn new(sounds: Arc<Mutex<HashMap<Sound, StaticSoundData>>>) -> Self {
38        let current_queue = Arc::new(Mutex::new(None::<TextureLoadQueue>));
39
40        let (sound_send, sound_recv) =
41            std::sync::mpsc::channel::<SoundAssetData>();
42
43        let (texture_data_send, texture_data_recv) =
44            std::sync::mpsc::channel::<TextureAssetData>();
45
46        let texture_data_send = Arc::new(Mutex::new(texture_data_send));
47        let texture_data_recv = Arc::new(Mutex::new(texture_data_recv));
48
49        Self {
50            wgpu_load_queue: current_queue,
51            #[cfg(not(target_arch = "wasm32"))]
52            thread_pool: rayon::ThreadPoolBuilder::new().build().unwrap(),
53
54            texture_load_queue: Vec::new(),
55            sound_load_queue: Vec::new(),
56
57            sounds,
58            sound_send: Arc::new(Mutex::new(sound_send)),
59            sound_recv: Arc::new(Mutex::new(sound_recv)),
60
61            texture_data_send,
62            texture_data_recv,
63
64            asset_source: None,
65
66            data_load_queue: Vec::new(),
67
68            pending_textures: Default::default(),
69        }
70    }
71
72    pub fn queue_load_sounds(&mut self, sounds: Vec<(String, String)>) {
73        inc_assets_queued(sounds.len());
74        self.sound_load_queue.extend(sounds);
75    }
76
77    pub fn queue_load_textures(&mut self, textures: Vec<(String, String)>) {
78        inc_assets_queued(textures.len());
79
80        for (key, _path) in textures.iter() {
81            self.pending_textures.insert(key.clone());
82        }
83
84        self.texture_load_queue.extend(textures)
85    }
86
87    pub fn parse_texture_byte_queue(
88        &mut self,
89        texture_image_map: Arc<Mutex<HashMap<TextureHandle, Arc<RgbaImage>>>>,
90    ) {
91        let _span = span!("parse_texture_byte_queue");
92
93        while let Ok(request) = self.texture_data_recv.lock().try_recv() {
94            let image_map = texture_image_map.clone();
95            let wgpu_load_queue = self.wgpu_load_queue.clone();
96
97            let process_image = move || {
98                let image = image::load_from_memory(&request.bytes);
99
100                let item = match image {
101                    Ok(image) => {
102                        image_map
103                            .lock()
104                            .insert(request.handle, Arc::new(image.to_rgba8()));
105
106                        inc_assets_loaded(1);
107
108                        LoadedImage {
109                            path: request.path,
110                            handle: request.handle,
111                            image,
112                        }
113                    }
114                    Err(err) => {
115                        error!("Failed to load {} ... {}", request.path, err);
116                        return;
117                    }
118                };
119
120                wgpu_load_queue.lock().get_or_insert_with(Vec::new).push(item);
121            };
122
123            #[cfg(target_arch = "wasm32")]
124            process_image();
125
126            #[cfg(not(target_arch = "wasm32"))]
127            self.thread_pool.spawn(process_image);
128        }
129    }
130
131    pub fn sound_tick(&mut self) {
132        let _span = span!("sound_tick");
133
134        while let Ok(item) = self.sound_recv.lock().try_recv() {
135            let sounds = self.sounds.clone();
136
137            let sound_loop = move || {
138                match StaticSoundData::from_cursor(
139                    std::io::Cursor::new(item.bytes),
140                    StaticSoundSettings::default(),
141                ) {
142                    Ok(sound) => {
143                        trace!("Sound {}", item.path);
144                        sounds.lock().insert(item.handle, sound);
145                        inc_assets_loaded(1);
146                    }
147                    Err(err) => {
148                        error!(
149                            "Failed to parse sound at {}: {:?}",
150                            item.path, err
151                        );
152                    }
153                }
154            };
155
156            #[cfg(target_arch = "wasm32")]
157            sound_loop();
158
159            #[cfg(not(target_arch = "wasm32"))]
160            self.thread_pool.spawn(sound_loop);
161        }
162    }
163
164    pub fn load_textures_to_memory(
165        &mut self,
166        loaded_textures: HashSet<TextureHandle>,
167        textures: &mut HashMap<String, TextureHandle>,
168    ) {
169        let _span = span!("load_textures_to_memory");
170
171        if let Some(asset_source) = self.asset_source.as_ref() {
172            for (key, relative_path) in self
173                .texture_load_queue
174                .drain(..)
175                .filter(|(key, _relative_path)| {
176                    !loaded_textures.contains(&texture_id_unchecked(key))
177                })
178            {
179                let handle = texture_id_unchecked(&key);
180
181                textures.insert(key, handle);
182
183                if let Ok(bytes) = asset_source.load_single_item(&relative_path)
184                {
185                    let texture_data =
186                        TextureAssetData { path: relative_path, handle, bytes };
187
188                    self.texture_data_send.lock().send(texture_data).log_err();
189                } else {
190                    report_error(
191                        format!("{:?}", handle),
192                        format!("Error loading {}", relative_path),
193                    );
194                }
195            }
196        } else {
197            assert!(
198                self.texture_load_queue.is_empty(),
199                "AssetSource must be initialized before textures are loaded"
200            );
201        }
202    }
203
204    pub fn load_sounds_to_memory(
205        &mut self,
206        sound_ids: &mut HashMap<String, Sound>,
207    ) {
208        let _span = span!("load_sounds_to_memory");
209
210        if let Some(asset_source) = self.asset_source.as_ref() {
211            for (key, relative_path) in self.sound_load_queue.drain(..) {
212                let handle = Sound::from_path(&key);
213
214                if self.sounds.lock().contains_key(&handle) {
215                    continue;
216                }
217
218                sound_ids.insert(key.to_string(), handle);
219
220                if let Ok(bytes) = asset_source.load_single_item(&relative_path)
221                {
222                    let item =
223                        SoundAssetData { path: relative_path, handle, bytes };
224
225                    self.sound_send.lock().send(item).log_err();
226                } else {
227                    error!("Error loading {}", relative_path);
228                    continue;
229                }
230            }
231        } else {
232            assert!(
233                self.sound_load_queue.is_empty(),
234                "AssetSource must be initialized before sounds are loaded"
235            );
236        }
237    }
238}
239
240pub enum AssetData {
241    Sound(SoundAssetData),
242    Texture(TextureAssetData),
243}
244
245pub struct SoundAssetData {
246    pub path: String,
247    pub handle: Sound,
248    pub bytes: Vec<u8>,
249}
250
251pub struct TextureAssetData {
252    pub path: String,
253    pub handle: TextureHandle,
254    pub bytes: Vec<u8>,
255}
256
257pub struct LoadedImage {
258    pub path: String,
259    pub handle: TextureHandle,
260    pub image: image::DynamicImage,
261}