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#[derive(Default, Clone, Debug)]
25pub struct Dir(Arc<RwLock<DirInternal>>);
26
27impl Dir {
28 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#[derive(Default, Clone)]
170pub struct MemoryAssetReader {
171 pub root: Dir,
172}
173
174#[derive(Clone, Debug)]
176pub struct Data {
177 path: PathBuf,
178 value: Value,
179}
180
181#[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}