bevy_asset/io/
memory.rs

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

bevy_asset/io/
memory.rs
