bevy_asset/io/
memory.rs

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

bevy_asset/io/
memory.rs
