bevy_asset/io/
memory.rs

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