bevy_asset/io/
memory.rs

1use crate::io::{
2    AssetReader, AssetReaderError, AssetWriter, AssetWriterError, PathStream, Reader,
3    ReaderRequiredFeatures,
4};
5use alloc::{borrow::ToOwned, boxed::Box, sync::Arc, vec, vec::Vec};
6use bevy_platform::{
7    collections::HashMap,
8    sync::{PoisonError, RwLock},
9};
10use core::{pin::Pin, task::Poll};
11use futures_io::{AsyncRead, AsyncWrite};
12use futures_lite::Stream;
13use std::{
14    io::{Error, ErrorKind, SeekFrom},
15    path::{Path, PathBuf},
16};
17
18use super::AsyncSeek;
19
20#[derive(Default, Debug)]
21struct DirInternal {
22    assets: HashMap<Box<str>, Data>,
23    metadata: HashMap<Box<str>, Data>,
24    dirs: HashMap<Box<str>, Dir>,
25    path: PathBuf,
26}
27
28/// A clone-able (internally Arc-ed) / thread-safe "in memory" filesystem.
29/// This is built for [`MemoryAssetReader`] and is primarily intended for unit tests.
30#[derive(Default, Clone, Debug)]
31pub struct Dir(Arc<RwLock<DirInternal>>);
32
33impl Dir {
34    /// Creates a new [`Dir`] for the given `path`.
35    pub fn new(path: PathBuf) -> Self {
36        Self(Arc::new(RwLock::new(DirInternal {
37            path,
38            ..Default::default()
39        })))
40    }
41
42    pub fn insert_asset_text(&self, path: &Path, asset: &str) {
43        self.insert_asset(path, asset.as_bytes().to_vec());
44    }
45
46    pub fn insert_meta_text(&self, path: &Path, asset: &str) {
47        self.insert_meta(path, asset.as_bytes().to_vec());
48    }
49
50    pub fn insert_asset(&self, path: &Path, value: impl Into<Value>) {
51        let mut dir = self.clone();
52        if let Some(parent) = path.parent() {
53            dir = self.get_or_insert_dir(parent);
54        }
55        dir.0
56            .write()
57            .unwrap_or_else(PoisonError::into_inner)
58            .assets
59            .insert(
60                path.file_name().unwrap().to_string_lossy().into(),
61                Data {
62                    value: value.into(),
63                    path: path.to_owned(),
64                },
65            );
66    }
67
68    /// Removes the stored asset at `path`.
69    ///
70    /// Returns the [`Data`] stored if found, [`None`] otherwise.
71    pub fn remove_asset(&self, path: &Path) -> Option<Data> {
72        let mut dir = self.clone();
73        if let Some(parent) = path.parent() {
74            dir = self.get_or_insert_dir(parent);
75        }
76        let key: Box<str> = path.file_name().unwrap().to_string_lossy().into();
77        dir.0
78            .write()
79            .unwrap_or_else(PoisonError::into_inner)
80            .assets
81            .remove(&key)
82    }
83
84    pub fn insert_meta(&self, path: &Path, value: impl Into<Value>) {
85        let mut dir = self.clone();
86        if let Some(parent) = path.parent() {
87            dir = self.get_or_insert_dir(parent);
88        }
89        dir.0
90            .write()
91            .unwrap_or_else(PoisonError::into_inner)
92            .metadata
93            .insert(
94                path.file_name().unwrap().to_string_lossy().into(),
95                Data {
96                    value: value.into(),
97                    path: path.to_owned(),
98                },
99            );
100    }
101
102    /// Removes the stored metadata at `path`.
103    ///
104    /// Returns the [`Data`] stored if found, [`None`] otherwise.
105    pub fn remove_metadata(&self, path: &Path) -> Option<Data> {
106        let mut dir = self.clone();
107        if let Some(parent) = path.parent() {
108            dir = self.get_or_insert_dir(parent);
109        }
110        let key: Box<str> = path.file_name().unwrap().to_string_lossy().into();
111        dir.0
112            .write()
113            .unwrap_or_else(PoisonError::into_inner)
114            .metadata
115            .remove(&key)
116    }
117
118    pub fn get_or_insert_dir(&self, path: &Path) -> Dir {
119        let mut dir = self.clone();
120        let mut full_path = PathBuf::new();
121        for c in path.components() {
122            full_path.push(c);
123            let name = c.as_os_str().to_string_lossy().into();
124            dir = {
125                let dirs = &mut dir.0.write().unwrap_or_else(PoisonError::into_inner).dirs;
126                dirs.entry(name)
127                    .or_insert_with(|| Dir::new(full_path.clone()))
128                    .clone()
129            };
130        }
131
132        dir
133    }
134
135    /// Removes the dir at `path`.
136    ///
137    /// Returns the [`Dir`] stored if found, [`None`] otherwise.
138    pub fn remove_dir(&self, path: &Path) -> Option<Dir> {
139        let mut dir = self.clone();
140        if let Some(parent) = path.parent() {
141            dir = self.get_or_insert_dir(parent);
142        }
143        let key: Box<str> = path.file_name().unwrap().to_string_lossy().into();
144        dir.0
145            .write()
146            .unwrap_or_else(PoisonError::into_inner)
147            .dirs
148            .remove(&key)
149    }
150
151    pub fn get_dir(&self, path: &Path) -> Option<Dir> {
152        let mut dir = self.clone();
153        for p in path.components() {
154            let component = p.as_os_str().to_str().unwrap();
155            let next_dir = dir
156                .0
157                .read()
158                .unwrap_or_else(PoisonError::into_inner)
159                .dirs
160                .get(component)?
161                .clone();
162            dir = next_dir;
163        }
164        Some(dir)
165    }
166
167    pub fn get_asset(&self, path: &Path) -> Option<Data> {
168        let mut dir = self.clone();
169        if let Some(parent) = path.parent() {
170            dir = dir.get_dir(parent)?;
171        }
172
173        path.file_name().and_then(|f| {
174            dir.0
175                .read()
176                .unwrap_or_else(PoisonError::into_inner)
177                .assets
178                .get(f.to_str().unwrap())
179                .cloned()
180        })
181    }
182
183    pub fn get_metadata(&self, path: &Path) -> Option<Data> {
184        let mut dir = self.clone();
185        if let Some(parent) = path.parent() {
186            dir = dir.get_dir(parent)?;
187        }
188
189        path.file_name().and_then(|f| {
190            dir.0
191                .read()
192                .unwrap_or_else(PoisonError::into_inner)
193                .metadata
194                .get(f.to_str().unwrap())
195                .cloned()
196        })
197    }
198
199    pub fn path(&self) -> PathBuf {
200        self.0
201            .read()
202            .unwrap_or_else(PoisonError::into_inner)
203            .path
204            .to_owned()
205    }
206}
207
208pub struct DirStream {
209    dir: Dir,
210    index: usize,
211    dir_index: usize,
212}
213
214impl DirStream {
215    fn new(dir: Dir) -> Self {
216        Self {
217            dir,
218            index: 0,
219            dir_index: 0,
220        }
221    }
222}
223
224impl Stream for DirStream {
225    type Item = PathBuf;
226
227    fn poll_next(
228        self: Pin<&mut Self>,
229        _cx: &mut core::task::Context<'_>,
230    ) -> Poll<Option<Self::Item>> {
231        let this = self.get_mut();
232        let dir = this.dir.0.read().unwrap_or_else(PoisonError::into_inner);
233
234        let dir_index = this.dir_index;
235        if let Some(dir_path) = dir
236            .dirs
237            .keys()
238            .nth(dir_index)
239            .map(|d| dir.path.join(d.as_ref()))
240        {
241            this.dir_index += 1;
242            Poll::Ready(Some(dir_path))
243        } else {
244            let index = this.index;
245            this.index += 1;
246            Poll::Ready(dir.assets.values().nth(index).map(|d| d.path().to_owned()))
247        }
248    }
249}
250
251/// In-memory [`AssetReader`] implementation.
252/// This is primarily intended for unit tests.
253#[derive(Default, Clone)]
254pub struct MemoryAssetReader {
255    pub root: Dir,
256}
257
258/// In-memory [`AssetWriter`] implementation.
259///
260/// This is primarily intended for unit tests.
261#[derive(Default, Clone)]
262pub struct MemoryAssetWriter {
263    pub root: Dir,
264}
265
266/// Asset data stored in a [`Dir`].
267#[derive(Clone, Debug)]
268pub struct Data {
269    path: PathBuf,
270    value: Value,
271}
272
273/// Stores either an allocated vec of bytes or a static array of bytes.
274#[derive(Clone, Debug)]
275pub enum Value {
276    Vec(Arc<Vec<u8>>),
277    Static(&'static [u8]),
278}
279
280impl Data {
281    /// The path that this data was written to.
282    pub fn path(&self) -> &Path {
283        &self.path
284    }
285
286    /// The value in bytes that was written here.
287    pub fn value(&self) -> &[u8] {
288        match &self.value {
289            Value::Vec(vec) => vec,
290            Value::Static(value) => value,
291        }
292    }
293}
294
295impl From<Vec<u8>> for Value {
296    fn from(value: Vec<u8>) -> Self {
297        Self::Vec(Arc::new(value))
298    }
299}
300
301impl From<&'static [u8]> for Value {
302    fn from(value: &'static [u8]) -> Self {
303        Self::Static(value)
304    }
305}
306
307impl<const N: usize> From<&'static [u8; N]> for Value {
308    fn from(value: &'static [u8; N]) -> Self {
309        Self::Static(value)
310    }
311}
312
313struct DataReader {
314    data: Data,
315    bytes_read: usize,
316}
317
318impl AsyncRead for DataReader {
319    fn poll_read(
320        self: Pin<&mut Self>,
321        _cx: &mut core::task::Context<'_>,
322        buf: &mut [u8],
323    ) -> Poll<futures_io::Result<usize>> {
324        // Get the mut borrow to avoid trying to borrow the pin itself multiple times.
325        let this = self.get_mut();
326        Poll::Ready(Ok(crate::io::slice_read(
327            this.data.value(),
328            &mut this.bytes_read,
329            buf,
330        )))
331    }
332}
333
334impl AsyncSeek for DataReader {
335    fn poll_seek(
336        self: Pin<&mut Self>,
337        _cx: &mut core::task::Context<'_>,
338        pos: SeekFrom,
339    ) -> Poll<std::io::Result<u64>> {
340        // Get the mut borrow to avoid trying to borrow the pin itself multiple times.
341        let this = self.get_mut();
342        Poll::Ready(crate::io::slice_seek(
343            this.data.value(),
344            &mut this.bytes_read,
345            pos,
346        ))
347    }
348}
349
350impl Reader for DataReader {
351    fn read_to_end<'a>(
352        &'a mut self,
353        buf: &'a mut Vec<u8>,
354    ) -> stackfuture::StackFuture<'a, std::io::Result<usize>, { super::STACK_FUTURE_SIZE }> {
355        crate::io::read_to_end(self.data.value(), &mut self.bytes_read, buf)
356    }
357}
358
359impl AssetReader for MemoryAssetReader {
360    async fn read<'a>(
361        &'a self,
362        path: &'a Path,
363        _required_features: ReaderRequiredFeatures,
364    ) -> Result<impl Reader + 'a, AssetReaderError> {
365        self.root
366            .get_asset(path)
367            .map(|data| DataReader {
368                data,
369                bytes_read: 0,
370            })
371            .ok_or_else(|| AssetReaderError::NotFound(path.to_path_buf()))
372    }
373
374    async fn read_meta<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
375        self.root
376            .get_metadata(path)
377            .map(|data| DataReader {
378                data,
379                bytes_read: 0,
380            })
381            .ok_or_else(|| AssetReaderError::NotFound(path.to_path_buf()))
382    }
383
384    async fn read_directory<'a>(
385        &'a self,
386        path: &'a Path,
387    ) -> Result<Box<PathStream>, AssetReaderError> {
388        self.root
389            .get_dir(path)
390            .map(|dir| {
391                let stream: Box<PathStream> = Box::new(DirStream::new(dir));
392                stream
393            })
394            .ok_or_else(|| AssetReaderError::NotFound(path.to_path_buf()))
395    }
396
397    async fn is_directory<'a>(&'a self, path: &'a Path) -> Result<bool, AssetReaderError> {
398        Ok(self.root.get_dir(path).is_some())
399    }
400}
401
402/// A writer that writes into [`Dir`], buffering internally until flushed/closed.
403struct DataWriter {
404    /// The dir to write to.
405    dir: Dir,
406    /// The path to write to.
407    path: PathBuf,
408    /// The current buffer of data.
409    ///
410    /// This will include data that has been flushed already.
411    current_data: Vec<u8>,
412    /// Whether to write to the data or to the meta.
413    is_meta_writer: bool,
414}
415
416impl AsyncWrite for DataWriter {
417    fn poll_write(
418        self: Pin<&mut Self>,
419        _: &mut core::task::Context<'_>,
420        buf: &[u8],
421    ) -> Poll<std::io::Result<usize>> {
422        self.get_mut().current_data.extend_from_slice(buf);
423        Poll::Ready(Ok(buf.len()))
424    }
425
426    fn poll_flush(
427        self: Pin<&mut Self>,
428        _: &mut core::task::Context<'_>,
429    ) -> Poll<std::io::Result<()>> {
430        // Write the data to our fake disk. This means we will repeatedly reinsert the asset.
431        if self.is_meta_writer {
432            self.dir.insert_meta(&self.path, self.current_data.clone());
433        } else {
434            self.dir.insert_asset(&self.path, self.current_data.clone());
435        }
436        Poll::Ready(Ok(()))
437    }
438
439    fn poll_close(
440        self: Pin<&mut Self>,
441        cx: &mut core::task::Context<'_>,
442    ) -> Poll<std::io::Result<()>> {
443        // A flush will just write the data to Dir, which is all we need to do for close.
444        self.poll_flush(cx)
445    }
446}
447
448impl AssetWriter for MemoryAssetWriter {
449    async fn write<'a>(&'a self, path: &'a Path) -> Result<Box<super::Writer>, AssetWriterError> {
450        Ok(Box::new(DataWriter {
451            dir: self.root.clone(),
452            path: path.to_owned(),
453            current_data: vec![],
454            is_meta_writer: false,
455        }))
456    }
457
458    async fn write_meta<'a>(
459        &'a self,
460        path: &'a Path,
461    ) -> Result<Box<super::Writer>, AssetWriterError> {
462        Ok(Box::new(DataWriter {
463            dir: self.root.clone(),
464            path: path.to_owned(),
465            current_data: vec![],
466            is_meta_writer: true,
467        }))
468    }
469
470    async fn remove<'a>(&'a self, path: &'a Path) -> Result<(), AssetWriterError> {
471        if self.root.remove_asset(path).is_none() {
472            return Err(AssetWriterError::Io(Error::new(
473                ErrorKind::NotFound,
474                "no such file",
475            )));
476        }
477        Ok(())
478    }
479
480    async fn remove_meta<'a>(&'a self, path: &'a Path) -> Result<(), AssetWriterError> {
481        self.root.remove_metadata(path);
482        Ok(())
483    }
484
485    async fn rename<'a>(
486        &'a self,
487        old_path: &'a Path,
488        new_path: &'a Path,
489    ) -> Result<(), AssetWriterError> {
490        let Some(old_asset) = self.root.get_asset(old_path) else {
491            return Err(AssetWriterError::Io(Error::new(
492                ErrorKind::NotFound,
493                "no such file",
494            )));
495        };
496        self.root.insert_asset(new_path, old_asset.value);
497        // Remove the asset after instead of before since otherwise there'd be a moment where the
498        // Dir is unlocked and missing both the old and new paths. This just prevents race
499        // conditions.
500        self.root.remove_asset(old_path);
501        Ok(())
502    }
503
504    async fn rename_meta<'a>(
505        &'a self,
506        old_path: &'a Path,
507        new_path: &'a Path,
508    ) -> Result<(), AssetWriterError> {
509        let Some(old_meta) = self.root.get_metadata(old_path) else {
510            return Err(AssetWriterError::Io(Error::new(
511                ErrorKind::NotFound,
512                "no such file",
513            )));
514        };
515        self.root.insert_meta(new_path, old_meta.value);
516        // Remove the meta after instead of before since otherwise there'd be a moment where the
517        // Dir is unlocked and missing both the old and new paths. This just prevents race
518        // conditions.
519        self.root.remove_metadata(old_path);
520        Ok(())
521    }
522
523    async fn create_directory<'a>(&'a self, path: &'a Path) -> Result<(), AssetWriterError> {
524        // Just pretend we're on a file system that doesn't consider directory re-creation a
525        // failure.
526        self.root.get_or_insert_dir(path);
527        Ok(())
528    }
529
530    async fn remove_directory<'a>(&'a self, path: &'a Path) -> Result<(), AssetWriterError> {
531        if self.root.remove_dir(path).is_none() {
532            return Err(AssetWriterError::Io(Error::new(
533                ErrorKind::NotFound,
534                "no such dir",
535            )));
536        }
537        Ok(())
538    }
539
540    async fn remove_empty_directory<'a>(&'a self, path: &'a Path) -> Result<(), AssetWriterError> {
541        let Some(dir) = self.root.get_dir(path) else {
542            return Err(AssetWriterError::Io(Error::new(
543                ErrorKind::NotFound,
544                "no such dir",
545            )));
546        };
547
548        let dir = dir.0.read().unwrap();
549        if !dir.assets.is_empty() || !dir.metadata.is_empty() || !dir.dirs.is_empty() {
550            return Err(AssetWriterError::Io(Error::new(
551                ErrorKind::DirectoryNotEmpty,
552                "not empty",
553            )));
554        }
555
556        self.root.remove_dir(path);
557        Ok(())
558    }
559
560    async fn remove_assets_in_directory<'a>(
561        &'a self,
562        path: &'a Path,
563    ) -> Result<(), AssetWriterError> {
564        let Some(dir) = self.root.get_dir(path) else {
565            return Err(AssetWriterError::Io(Error::new(
566                ErrorKind::NotFound,
567                "no such dir",
568            )));
569        };
570
571        let mut dir = dir.0.write().unwrap();
572        dir.assets.clear();
573        dir.dirs.clear();
574        dir.metadata.clear();
575        Ok(())
576    }
577}
578
579#[cfg(test)]
580pub mod test {
581    use super::Dir;
582    use std::path::Path;
583
584    #[test]
585    fn memory_dir() {
586        let dir = Dir::default();
587        let a_path = Path::new("a.txt");
588        let a_data = "a".as_bytes().to_vec();
589        let a_meta = "ameta".as_bytes().to_vec();
590
591        dir.insert_asset(a_path, a_data.clone());
592        let asset = dir.get_asset(a_path).unwrap();
593        assert_eq!(asset.path(), a_path);
594        assert_eq!(asset.value(), a_data);
595
596        dir.insert_meta(a_path, a_meta.clone());
597        let meta = dir.get_metadata(a_path).unwrap();
598        assert_eq!(meta.path(), a_path);
599        assert_eq!(meta.value(), a_meta);
600
601        let b_path = Path::new("x/y/b.txt");
602        let b_data = "b".as_bytes().to_vec();
603        let b_meta = "meta".as_bytes().to_vec();
604        dir.insert_asset(b_path, b_data.clone());
605        dir.insert_meta(b_path, b_meta.clone());
606
607        let asset = dir.get_asset(b_path).unwrap();
608        assert_eq!(asset.path(), b_path);
609        assert_eq!(asset.value(), b_data);
610
611        let meta = dir.get_metadata(b_path).unwrap();
612        assert_eq!(meta.path(), b_path);
613        assert_eq!(meta.value(), b_meta);
614    }
615}
bevy_asset/io/memory.rs

bevy_asset/io/
memory.rs
