bevy_asset/io/
memory.rs

1use crate::io::{AssetReader, AssetReaderError, PathStream, Reader};
2use alloc::sync::Arc;
3use bevy_utils::HashMap;
4use core::{pin::Pin, task::Poll};
5use futures_io::AsyncRead;
6use futures_lite::{ready, Stream};
7use parking_lot::RwLock;
8use std::path::{Path, PathBuf};
9
10use super::AsyncSeekForward;
11
12#[derive(Default, Debug)]
13struct DirInternal {
14    assets: HashMap<Box<str>, Data>,
15    metadata: HashMap<Box<str>, Data>,
16    dirs: HashMap<Box<str>, Dir>,
17    path: PathBuf,
18}
19
20/// A clone-able (internally Arc-ed) / thread-safe "in memory" filesystem.
21/// This is built for [`MemoryAssetReader`] and is primarily intended for unit tests.
22#[derive(Default, Clone, Debug)]
23pub struct Dir(Arc<RwLock<DirInternal>>);
24
25impl Dir {
26    /// Creates a new [`Dir`] for the given `path`.
27    pub fn new(path: PathBuf) -> Self {
28        Self(Arc::new(RwLock::new(DirInternal {
29            path,
30            ..Default::default()
31        })))
32    }
33
34    pub fn insert_asset_text(&self, path: &Path, asset: &str) {
35        self.insert_asset(path, asset.as_bytes().to_vec());
36    }
37
38    pub fn insert_meta_text(&self, path: &Path, asset: &str) {
39        self.insert_meta(path, asset.as_bytes().to_vec());
40    }
41
42    pub fn insert_asset(&self, path: &Path, value: impl Into<Value>) {
43        let mut dir = self.clone();
44        if let Some(parent) = path.parent() {
45            dir = self.get_or_insert_dir(parent);
46        }
47        dir.0.write().assets.insert(
48            path.file_name().unwrap().to_string_lossy().into(),
49            Data {
50                value: value.into(),
51                path: path.to_owned(),
52            },
53        );
54    }
55
56    /// Removes the stored asset at `path` and returns the `Data` stored if found and otherwise `None`.
57    pub fn remove_asset(&self, path: &Path) -> Option<Data> {
58        let mut dir = self.clone();
59        if let Some(parent) = path.parent() {
60            dir = self.get_or_insert_dir(parent);
61        }
62        let key: Box<str> = path.file_name().unwrap().to_string_lossy().into();
63        let data = dir.0.write().assets.remove(&key);
64        data
65    }
66
67    pub fn insert_meta(&self, path: &Path, value: impl Into<Value>) {
68        let mut dir = self.clone();
69        if let Some(parent) = path.parent() {
70            dir = self.get_or_insert_dir(parent);
71        }
72        dir.0.write().metadata.insert(
73            path.file_name().unwrap().to_string_lossy().into(),
74            Data {
75                value: value.into(),
76                path: path.to_owned(),
77            },
78        );
79    }
80
81    pub fn get_or_insert_dir(&self, path: &Path) -> Dir {
82        let mut dir = self.clone();
83        let mut full_path = PathBuf::new();
84        for c in path.components() {
85            full_path.push(c);
86            let name = c.as_os_str().to_string_lossy().into();
87            dir = {
88                let dirs = &mut dir.0.write().dirs;
89                dirs.entry(name)
90                    .or_insert_with(|| Dir::new(full_path.clone()))
91                    .clone()
92            };
93        }
94
95        dir
96    }
97
98    pub fn get_dir(&self, path: &Path) -> Option<Dir> {
99        let mut dir = self.clone();
100        for p in path.components() {
101            let component = p.as_os_str().to_str().unwrap();
102            let next_dir = dir.0.read().dirs.get(component)?.clone();
103            dir = next_dir;
104        }
105        Some(dir)
106    }
107
108    pub fn get_asset(&self, path: &Path) -> Option<Data> {
109        let mut dir = self.clone();
110        if let Some(parent) = path.parent() {
111            dir = dir.get_dir(parent)?;
112        }
113
114        path.file_name()
115            .and_then(|f| dir.0.read().assets.get(f.to_str().unwrap()).cloned())
116    }
117
118    pub fn get_metadata(&self, path: &Path) -> Option<Data> {
119        let mut dir = self.clone();
120        if let Some(parent) = path.parent() {
121            dir = dir.get_dir(parent)?;
122        }
123
124        path.file_name()
125            .and_then(|f| dir.0.read().metadata.get(f.to_str().unwrap()).cloned())
126    }
127
128    pub fn path(&self) -> PathBuf {
129        self.0.read().path.to_owned()
130    }
131}
132
133pub struct DirStream {
134    dir: Dir,
135    index: usize,
136    dir_index: usize,
137}
138
139impl DirStream {
140    fn new(dir: Dir) -> Self {
141        Self {
142            dir,
143            index: 0,
144            dir_index: 0,
145        }
146    }
147}
148
149impl Stream for DirStream {
150    type Item = PathBuf;
151
152    fn poll_next(
153        self: Pin<&mut Self>,
154        _cx: &mut core::task::Context<'_>,
155    ) -> Poll<Option<Self::Item>> {
156        let this = self.get_mut();
157        let dir = this.dir.0.read();
158
159        let dir_index = this.dir_index;
160        if let Some(dir_path) = dir
161            .dirs
162            .keys()
163            .nth(dir_index)
164            .map(|d| dir.path.join(d.as_ref()))
165        {
166            this.dir_index += 1;
167            Poll::Ready(Some(dir_path))
168        } else {
169            let index = this.index;
170            this.index += 1;
171            Poll::Ready(dir.assets.values().nth(index).map(|d| d.path().to_owned()))
172        }
173    }
174}
175
176/// In-memory [`AssetReader`] implementation.
177/// This is primarily intended for unit tests.
178#[derive(Default, Clone)]
179pub struct MemoryAssetReader {
180    pub root: Dir,
181}
182
183/// Asset data stored in a [`Dir`].
184#[derive(Clone, Debug)]
185pub struct Data {
186    path: PathBuf,
187    value: Value,
188}
189
190/// Stores either an allocated vec of bytes or a static array of bytes.
191#[derive(Clone, Debug)]
192pub enum Value {
193    Vec(Arc<Vec<u8>>),
194    Static(&'static [u8]),
195}
196
197impl Data {
198    fn path(&self) -> &Path {
199        &self.path
200    }
201    fn value(&self) -> &[u8] {
202        match &self.value {
203            Value::Vec(vec) => vec,
204            Value::Static(value) => value,
205        }
206    }
207}
208
209impl From<Vec<u8>> for Value {
210    fn from(value: Vec<u8>) -> Self {
211        Self::Vec(Arc::new(value))
212    }
213}
214
215impl From<&'static [u8]> for Value {
216    fn from(value: &'static [u8]) -> Self {
217        Self::Static(value)
218    }
219}
220
221impl<const N: usize> From<&'static [u8; N]> for Value {
222    fn from(value: &'static [u8; N]) -> Self {
223        Self::Static(value)
224    }
225}
226
227struct DataReader {
228    data: Data,
229    bytes_read: usize,
230}
231
232impl AsyncRead for DataReader {
233    fn poll_read(
234        mut self: Pin<&mut Self>,
235        cx: &mut core::task::Context<'_>,
236        buf: &mut [u8],
237    ) -> Poll<futures_io::Result<usize>> {
238        if self.bytes_read >= self.data.value().len() {
239            Poll::Ready(Ok(0))
240        } else {
241            let n =
242                ready!(Pin::new(&mut &self.data.value()[self.bytes_read..]).poll_read(cx, buf))?;
243            self.bytes_read += n;
244            Poll::Ready(Ok(n))
245        }
246    }
247}
248
249impl AsyncSeekForward for DataReader {
250    fn poll_seek_forward(
251        mut self: Pin<&mut Self>,
252        _cx: &mut core::task::Context<'_>,
253        offset: u64,
254    ) -> Poll<std::io::Result<u64>> {
255        let result = self
256            .bytes_read
257            .try_into()
258            .map(|bytes_read: u64| bytes_read + offset);
259
260        if let Ok(new_pos) = result {
261            self.bytes_read = new_pos as _;
262            Poll::Ready(Ok(new_pos as _))
263        } else {
264            Poll::Ready(Err(std::io::Error::new(
265                std::io::ErrorKind::InvalidInput,
266                "seek position is out of range",
267            )))
268        }
269    }
270}
271
272impl Reader for DataReader {
273    fn read_to_end<'a>(
274        &'a mut self,
275        buf: &'a mut Vec<u8>,
276    ) -> stackfuture::StackFuture<'a, std::io::Result<usize>, { super::STACK_FUTURE_SIZE }> {
277        stackfuture::StackFuture::from(async {
278            if self.bytes_read >= self.data.value().len() {
279                Ok(0)
280            } else {
281                buf.extend_from_slice(&self.data.value()[self.bytes_read..]);
282                let n = self.data.value().len() - self.bytes_read;
283                self.bytes_read = self.data.value().len();
284                Ok(n)
285            }
286        })
287    }
288}
289
290impl AssetReader for MemoryAssetReader {
291    async fn read<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
292        self.root
293            .get_asset(path)
294            .map(|data| DataReader {
295                data,
296                bytes_read: 0,
297            })
298            .ok_or_else(|| AssetReaderError::NotFound(path.to_path_buf()))
299    }
300
301    async fn read_meta<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
302        self.root
303            .get_metadata(path)
304            .map(|data| DataReader {
305                data,
306                bytes_read: 0,
307            })
308            .ok_or_else(|| AssetReaderError::NotFound(path.to_path_buf()))
309    }
310
311    async fn read_directory<'a>(
312        &'a self,
313        path: &'a Path,
314    ) -> Result<Box<PathStream>, AssetReaderError> {
315        self.root
316            .get_dir(path)
317            .map(|dir| {
318                let stream: Box<PathStream> = Box::new(DirStream::new(dir));
319                stream
320            })
321            .ok_or_else(|| AssetReaderError::NotFound(path.to_path_buf()))
322    }
323
324    async fn is_directory<'a>(&'a self, path: &'a Path) -> Result<bool, AssetReaderError> {
325        Ok(self.root.get_dir(path).is_some())
326    }
327}
328
329#[cfg(test)]
330pub mod test {
331    use super::Dir;
332    use std::path::Path;
333
334    #[test]
335    fn memory_dir() {
336        let dir = Dir::default();
337        let a_path = Path::new("a.txt");
338        let a_data = "a".as_bytes().to_vec();
339        let a_meta = "ameta".as_bytes().to_vec();
340
341        dir.insert_asset(a_path, a_data.clone());
342        let asset = dir.get_asset(a_path).unwrap();
343        assert_eq!(asset.path(), a_path);
344        assert_eq!(asset.value(), a_data);
345
346        dir.insert_meta(a_path, a_meta.clone());
347        let meta = dir.get_metadata(a_path).unwrap();
348        assert_eq!(meta.path(), a_path);
349        assert_eq!(meta.value(), a_meta);
350
351        let b_path = Path::new("x/y/b.txt");
352        let b_data = "b".as_bytes().to_vec();
353        let b_meta = "meta".as_bytes().to_vec();
354        dir.insert_asset(b_path, b_data.clone());
355        dir.insert_meta(b_path, b_meta.clone());
356
357        let asset = dir.get_asset(b_path).unwrap();
358        assert_eq!(asset.path(), b_path);
359        assert_eq!(asset.value(), b_data);
360
361        let meta = dir.get_metadata(b_path).unwrap();
362        assert_eq!(meta.path(), b_path);
363        assert_eq!(meta.value(), b_meta);
364    }
365}
bevy_asset/io/memory.rs

bevy_asset/io/
memory.rs
