From eb44db44373bc7a268c193951ffdfc19e64c1234 Mon Sep 17 00:00:00 2001 From: BeastLe9enD Date: Sat, 30 Mar 2024 23:26:30 +0100 Subject: [PATCH] Add `AsyncSeek` trait to `Reader` to be able to seek inside asset loaders (#12547) # Objective For some asset loaders, it can be useful not to read the entire asset file and just read a specific region of a file. For this, we need a way to seek at a specific position inside the file ## Solution I added support for `AsyncSeek` to `Reader`. In my case, I want to only read a part of a file, and for that I need to seek to a specific point. ## Migration Guide Every custom reader (which previously only needed the `AsyncRead` trait implemented) now also needs to implement the `AsyncSeek` trait to add the seek capability. --- .../bevy_asset/src/io/file/sync_file_asset.rs | 16 ++- crates/bevy_asset/src/io/memory.rs | 43 ++++++- crates/bevy_asset/src/io/mod.rs | 112 +++++++++++++++++- crates/bevy_asset/src/io/processor_gated.rs | 14 ++- crates/bevy_asset/src/processor/process.rs | 4 +- 5 files changed, 182 insertions(+), 7 deletions(-) diff --git a/crates/bevy_asset/src/io/file/sync_file_asset.rs b/crates/bevy_asset/src/io/file/sync_file_asset.rs index 4264721501..83eacd5649 100644 --- a/crates/bevy_asset/src/io/file/sync_file_asset.rs +++ b/crates/bevy_asset/src/io/file/sync_file_asset.rs @@ -1,4 +1,4 @@ -use futures_io::{AsyncRead, AsyncWrite}; +use futures_io::{AsyncRead, AsyncSeek, AsyncWrite}; use futures_lite::Stream; use crate::io::{ @@ -8,7 +8,7 @@ use crate::io::{ use std::{ fs::{read_dir, File}, - io::{Read, Write}, + io::{Read, Seek, Write}, path::{Path, PathBuf}, pin::Pin, task::Poll, @@ -30,6 +30,18 @@ impl AsyncRead for FileReader { } } +impl AsyncSeek for FileReader { + fn poll_seek( + self: Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + pos: std::io::SeekFrom, + ) -> Poll> { + let this = self.get_mut(); + let seek = this.0.seek(pos); + Poll::Ready(seek) + } +} + struct FileWriter(File); impl AsyncWrite for FileWriter { diff --git a/crates/bevy_asset/src/io/memory.rs b/crates/bevy_asset/src/io/memory.rs index 563086f7b0..782df9e3ef 100644 --- a/crates/bevy_asset/src/io/memory.rs +++ b/crates/bevy_asset/src/io/memory.rs @@ -1,8 +1,9 @@ use crate::io::{AssetReader, AssetReaderError, PathStream, Reader}; use bevy_utils::HashMap; -use futures_io::AsyncRead; +use futures_io::{AsyncRead, AsyncSeek}; use futures_lite::{ready, Stream}; use parking_lot::RwLock; +use std::io::SeekFrom; use std::{ path::{Path, PathBuf}, pin::Pin, @@ -236,6 +237,46 @@ impl AsyncRead for DataReader { } } +impl AsyncSeek for DataReader { + fn poll_seek( + mut self: Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + pos: SeekFrom, + ) -> Poll> { + let result = match pos { + SeekFrom::Start(offset) => offset.try_into(), + SeekFrom::End(offset) => self + .data + .value() + .len() + .try_into() + .map(|len: i64| len - offset), + SeekFrom::Current(offset) => self + .bytes_read + .try_into() + .map(|bytes_read: i64| bytes_read + offset), + }; + + if let Ok(new_pos) = result { + if new_pos < 0 { + Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "seek position is out of range", + ))) + } else { + self.bytes_read = new_pos as _; + + Poll::Ready(Ok(new_pos as _)) + } + } else { + Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "seek position is out of range", + ))) + } + } +} + impl AssetReader for MemoryAssetReader { async fn read<'a>(&'a self, path: &'a Path) -> Result>, AssetReaderError> { self.root diff --git a/crates/bevy_asset/src/io/mod.rs b/crates/bevy_asset/src/io/mod.rs index 766e536455..1a1c680ecd 100644 --- a/crates/bevy_asset/src/io/mod.rs +++ b/crates/bevy_asset/src/io/mod.rs @@ -22,8 +22,10 @@ pub use futures_lite::{AsyncReadExt, AsyncWriteExt}; pub use source::*; use bevy_utils::{BoxedFuture, ConditionalSendFuture}; -use futures_io::{AsyncRead, AsyncWrite}; +use futures_io::{AsyncRead, AsyncSeek, AsyncWrite}; use futures_lite::{ready, Stream}; +use std::io::SeekFrom; +use std::task::Context; use std::{ path::{Path, PathBuf}, pin::Pin, @@ -55,7 +57,11 @@ impl From for AssetReaderError { } } -pub type Reader<'a> = dyn AsyncRead + Unpin + Send + Sync + 'a; +pub trait AsyncReadAndSeek: AsyncRead + AsyncSeek {} + +impl AsyncReadAndSeek for T {} + +pub type Reader<'a> = dyn AsyncReadAndSeek + Unpin + Send + Sync + 'a; /// Performs read operations on an asset storage. [`AssetReader`] exposes a "virtual filesystem" /// API, where asset bytes and asset metadata bytes are both stored and accessible for a given @@ -442,6 +448,108 @@ impl AsyncRead for VecReader { } } +impl AsyncSeek for VecReader { + fn poll_seek( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + pos: SeekFrom, + ) -> Poll> { + let result = match pos { + SeekFrom::Start(offset) => offset.try_into(), + SeekFrom::End(offset) => self.bytes.len().try_into().map(|len: i64| len - offset), + SeekFrom::Current(offset) => self + .bytes_read + .try_into() + .map(|bytes_read: i64| bytes_read + offset), + }; + + if let Ok(new_pos) = result { + if new_pos < 0 { + Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "seek position is out of range", + ))) + } else { + self.bytes_read = new_pos as _; + + Poll::Ready(Ok(new_pos as _)) + } + } else { + Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "seek position is out of range", + ))) + } + } +} + +/// An [`AsyncRead`] implementation capable of reading a [`&[u8]`]. +pub struct SliceReader<'a> { + bytes: &'a [u8], + bytes_read: usize, +} + +impl<'a> SliceReader<'a> { + /// Create a new [`SliceReader`] for `bytes`. + pub fn new(bytes: &'a [u8]) -> Self { + Self { + bytes, + bytes_read: 0, + } + } +} + +impl<'a> AsyncRead for SliceReader<'a> { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll> { + if self.bytes_read >= self.bytes.len() { + Poll::Ready(Ok(0)) + } else { + let n = ready!(Pin::new(&mut &self.bytes[self.bytes_read..]).poll_read(cx, buf))?; + self.bytes_read += n; + Poll::Ready(Ok(n)) + } + } +} + +impl<'a> AsyncSeek for SliceReader<'a> { + fn poll_seek( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + pos: SeekFrom, + ) -> Poll> { + let result = match pos { + SeekFrom::Start(offset) => offset.try_into(), + SeekFrom::End(offset) => self.bytes.len().try_into().map(|len: i64| len - offset), + SeekFrom::Current(offset) => self + .bytes_read + .try_into() + .map(|bytes_read: i64| bytes_read + offset), + }; + + if let Ok(new_pos) = result { + if new_pos < 0 { + Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "seek position is out of range", + ))) + } else { + self.bytes_read = new_pos as _; + + Poll::Ready(Ok(new_pos as _)) + } + } else { + Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "seek position is out of range", + ))) + } + } +} + /// Appends `.meta` to the given path. pub(crate) fn get_meta_path(path: &Path) -> PathBuf { let mut meta_path = path.to_path_buf(); diff --git a/crates/bevy_asset/src/io/processor_gated.rs b/crates/bevy_asset/src/io/processor_gated.rs index 5f3bfe1467..f3d517b6b9 100644 --- a/crates/bevy_asset/src/io/processor_gated.rs +++ b/crates/bevy_asset/src/io/processor_gated.rs @@ -5,7 +5,9 @@ use crate::{ }; use async_lock::RwLockReadGuardArc; use bevy_utils::tracing::trace; -use futures_io::AsyncRead; +use futures_io::{AsyncRead, AsyncSeek}; +use std::io::SeekFrom; +use std::task::Poll; use std::{path::Path, pin::Pin, sync::Arc}; use super::ErasedAssetReader; @@ -139,3 +141,13 @@ impl<'a> AsyncRead for TransactionLockedReader<'a> { Pin::new(&mut self.reader).poll_read(cx, buf) } } + +impl<'a> AsyncSeek for TransactionLockedReader<'a> { + fn poll_seek( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + pos: SeekFrom, + ) -> Poll> { + Pin::new(&mut self.reader).poll_seek(cx, pos) + } +} diff --git a/crates/bevy_asset/src/processor/process.rs b/crates/bevy_asset/src/processor/process.rs index ce18bb53da..442c9451c8 100644 --- a/crates/bevy_asset/src/processor/process.rs +++ b/crates/bevy_asset/src/processor/process.rs @@ -1,3 +1,4 @@ +use crate::io::SliceReader; use crate::{ io::{ AssetReaderError, AssetWriterError, MissingAssetWriterError, @@ -343,12 +344,13 @@ impl<'a> ProcessContext<'a> { let server = &self.processor.server; let loader_name = std::any::type_name::(); let loader = server.get_asset_loader_with_type_name(loader_name).await?; + let mut reader = SliceReader::new(self.asset_bytes); let loaded_asset = server .load_with_meta_loader_and_reader( self.path, Box::new(meta), &*loader, - &mut self.asset_bytes, + &mut reader, false, true, )