bevy_asset/io/
memory.rs

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

bevy_asset/io/
memory.rs
