Allow Reader to implement AsyncSeek and provide a way for loaders to "ask" for a compatible reader. (#22104)

# Objective

- In #14194, we replaced `AsyncSeek` with `AsyncSeekForward`. This
allowed us to support more kinds of sources, but the cost was that
loaders were more limited in what they could do with the reader.
- Archive-like things generally use offsets relative to the start of the
file. `AsyncSeekForward` made it difficult to deal with this (or
impossible in some cases).
- This is a step towards #21641.
- Fixes #12880.

## Solution

- Allow `AssetLoader` to say which "features" of a reader it needs.
- Pass that list of features to `AssetReader` so it can decide how to
handle it.
- Add a new `AssetReaderError::UnsupportedFeature`.

This design is kind of "weak" - there's no guarantee that a loader that
requests `AnySeek` will get a reader that actually implements
`AsyncSeek`. Or on the other side, there's no guarantee that a loader
actually requests the features that it uses. However, in practice it's
likely enough: errors are likely to guide users to the correct
situation. In the future, we could perhaps have a "sieve reader", which
blocks any features the loader didn't explicitly request. Perhaps this
is a debug only feature, or something that can be toggled.

## Testing

- Ran the modified examples. They still seem to work!

---------

Co-authored-by: Carter Anderson <[email protected]>

Author: andriyDev

crates/bevy_asset/src/io/android.rs

@@ -1,4 +1,7 @@
-use crate::io::{get_meta_path, AssetReader, AssetReaderError, PathStream, Reader, VecReader};
+use crate::io::{
+    get_meta_path, AssetReader, AssetReaderError, PathStream, Reader, ReaderRequiredFeatures,
+    VecReader,
+};
 use alloc::{borrow::ToOwned, boxed::Box, ffi::CString, vec::Vec};
 use futures_lite::stream;
 use std::path::Path;
@@ -16,7 +19,11 @@ use std::path::Path;
 pub struct AndroidAssetReader;
 
 impl AssetReader for AndroidAssetReader {
-    async fn read<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
+    async fn read<'a>(
+        &'a self,
+        path: &'a Path,
+        _required_features: ReaderRequiredFeatures,
+    ) -> Result<impl Reader + 'a, AssetReaderError> {
         let asset_manager = bevy_android::ANDROID_APP
             .get()
             .expect("Bevy must be setup with the #[bevy_main] macro on Android")

crates/bevy_asset/src/io/file/file_asset.rs

@@ -1,40 +1,23 @@
 use crate::io::{
-    get_meta_path, AssetReader, AssetReaderError, AssetWriter, AssetWriterError, AsyncSeekForward,
-    PathStream, Reader, Writer,
+    get_meta_path, AssetReader, AssetReaderError, AssetWriter, AssetWriterError, PathStream,
+    Reader, ReaderRequiredFeatures, Writer,
 };
 use async_fs::{read_dir, File};
-use futures_io::AsyncSeek;
 use futures_lite::StreamExt;
 
 use alloc::{borrow::ToOwned, boxed::Box};
-use core::{pin::Pin, task, task::Poll};
 use std::path::Path;
 
 use super::{FileAssetReader, FileAssetWriter};
 
-impl AsyncSeekForward for File {
-    fn poll_seek_forward(
-        mut self: Pin<&mut Self>,
-        cx: &mut task::Context<'_>,
-        offset: u64,
-    ) -> Poll<futures_io::Result<u64>> {
-        let offset: Result<i64, _> = offset.try_into();
-
-        if let Ok(offset) = offset {
-            Pin::new(&mut self).poll_seek(cx, futures_io::SeekFrom::Current(offset))
-        } else {
-            Poll::Ready(Err(std::io::Error::new(
-                std::io::ErrorKind::InvalidInput,
-                "seek position is out of range",
-            )))
-        }
-    }
-}
-
 impl Reader for File {}
 
 impl AssetReader for FileAssetReader {
-    async fn read<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
+    async fn read<'a>(
+        &'a self,
+        path: &'a Path,
+        _required_features: ReaderRequiredFeatures,
+    ) -> Result<impl Reader + 'a, AssetReaderError> {
         let full_path = self.root_path.join(path);
         File::open(&full_path).await.map_err(|e| {
             if e.kind() == std::io::ErrorKind::NotFound {

crates/bevy_asset/src/io/file/sync_file_asset.rs

@@ -2,15 +2,15 @@ use futures_io::{AsyncRead, AsyncWrite};
 use futures_lite::Stream;
 
 use crate::io::{
-    get_meta_path, AssetReader, AssetReaderError, AssetWriter, AssetWriterError, AsyncSeekForward,
-    PathStream, Reader, Writer,
+    get_meta_path, AssetReader, AssetReaderError, AssetWriter, AssetWriterError, AsyncSeek,
+    PathStream, Reader, ReaderRequiredFeatures, Writer,
 };
 
 use alloc::{borrow::ToOwned, boxed::Box, vec::Vec};
 use core::{pin::Pin, task::Poll};
 use std::{
     fs::{read_dir, File},
-    io::{Read, Seek, Write},
+    io::{Read, Seek, SeekFrom, Write},
     path::{Path, PathBuf},
 };
 
@@ -30,17 +30,13 @@ impl AsyncRead for FileReader {
     }
 }
 
-impl AsyncSeekForward for FileReader {
-    fn poll_seek_forward(
-        self: Pin<&mut Self>,
+impl AsyncSeek for FileReader {
+    fn poll_seek(
+        mut self: Pin<&mut Self>,
         _cx: &mut core::task::Context<'_>,
-        offset: u64,
+        pos: SeekFrom,
     ) -> Poll<std::io::Result<u64>> {
-        let this = self.get_mut();
-        let current = this.0.stream_position()?;
-        let seek = this.0.seek(std::io::SeekFrom::Start(current + offset));
-
-        Poll::Ready(seek)
+        Poll::Ready(self.0.seek(pos))
     }
 }
 
@@ -99,7 +95,11 @@ impl Stream for DirReader {
 }
 
 impl AssetReader for FileAssetReader {
-    async fn read<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
+    async fn read<'a>(
+        &'a self,
+        path: &'a Path,
+        _required_features: ReaderRequiredFeatures,
+    ) -> Result<impl Reader + 'a, AssetReaderError> {
         let full_path = self.root_path.join(path);
         match File::open(&full_path) {
             Ok(file) => Ok(FileReader(file)),

crates/bevy_asset/src/io/gated.rs

@@ -1,4 +1,4 @@
-use crate::io::{AssetReader, AssetReaderError, PathStream, Reader};
+use crate::io::{AssetReader, AssetReaderError, PathStream, Reader, ReaderRequiredFeatures};
 use alloc::{boxed::Box, sync::Arc};
 use async_channel::{Receiver, Sender};
 use bevy_platform::{collections::HashMap, sync::RwLock};
@@ -55,7 +55,11 @@ impl<R: AssetReader> GatedReader<R> {
 }
 
 impl<R: AssetReader> AssetReader for GatedReader<R> {
-    async fn read<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
+    async fn read<'a>(
+        &'a self,
+        path: &'a Path,
+        required_features: ReaderRequiredFeatures,
+    ) -> Result<impl Reader + 'a, AssetReaderError> {
         let receiver = {
             let mut gates = self.gates.write().unwrap_or_else(PoisonError::into_inner);
             let gates = gates
@@ -64,7 +68,7 @@ impl<R: AssetReader> AssetReader for GatedReader<R> {
             gates.1.clone()
         };
         receiver.recv().await.unwrap();
-        let result = self.reader.read(path).await?;
+        let result = self.reader.read(path, required_features).await?;
         Ok(result)
     }
 

crates/bevy_asset/src/io/memory.rs

@@ -1,18 +1,21 @@
-use crate::io::{AssetReader, AssetReaderError, AssetWriter, AssetWriterError, PathStream, Reader};
+use crate::io::{
+    AssetReader, AssetReaderError, AssetWriter, AssetWriterError, PathStream, Reader,
+    ReaderRequiredFeatures,
+};
 use alloc::{borrow::ToOwned, boxed::Box, sync::Arc, vec, vec::Vec};
 use bevy_platform::{
     collections::HashMap,
     sync::{PoisonError, RwLock},
 };
 use core::{pin::Pin, task::Poll};
 use futures_io::{AsyncRead, AsyncWrite};
-use futures_lite::{ready, Stream};
+use futures_lite::Stream;
 use std::{
-    io::{Error, ErrorKind},
+    io::{Error, ErrorKind, SeekFrom},
     path::{Path, PathBuf},
 };
 
-use super::AsyncSeekForward;
+use super::AsyncSeek;
 
 #[derive(Default, Debug)]
 struct DirInternal {
@@ -314,41 +317,33 @@ struct DataReader {
 
 impl AsyncRead for DataReader {
     fn poll_read(
-        mut self: Pin<&mut Self>,
-        cx: &mut core::task::Context<'_>,
+        self: Pin<&mut Self>,
+        _cx: &mut core::task::Context<'_>,
         buf: &mut [u8],
     ) -> Poll<futures_io::Result<usize>> {
-        if self.bytes_read >= self.data.value().len() {
-            Poll::Ready(Ok(0))
-        } else {
-            let n =
-                ready!(Pin::new(&mut &self.data.value()[self.bytes_read..]).poll_read(cx, buf))?;
-            self.bytes_read += n;
-            Poll::Ready(Ok(n))
-        }
+        // Get the mut borrow to avoid trying to borrow the pin itself multiple times.
+        let this = self.get_mut();
+        Poll::Ready(Ok(crate::io::slice_read(
+            this.data.value(),
+            &mut this.bytes_read,
+            buf,
+        )))
     }
 }
 
-impl AsyncSeekForward for DataReader {
-    fn poll_seek_forward(
-        mut self: Pin<&mut Self>,
+impl AsyncSeek for DataReader {
+    fn poll_seek(
+        self: Pin<&mut Self>,
         _cx: &mut core::task::Context<'_>,
-        offset: u64,
+        pos: SeekFrom,
     ) -> Poll<std::io::Result<u64>> {
-        let result = self
-            .bytes_read
-            .try_into()
-            .map(|bytes_read: u64| bytes_read + offset);
-
-        if let Ok(new_pos) = result {
-            self.bytes_read = new_pos as _;
-            Poll::Ready(Ok(new_pos as _))
-        } else {
-            Poll::Ready(Err(Error::new(
-                ErrorKind::InvalidInput,
-                "seek position is out of range",
-            )))
-        }
+        // Get the mut borrow to avoid trying to borrow the pin itself multiple times.
+        let this = self.get_mut();
+        Poll::Ready(crate::io::slice_seek(
+            this.data.value(),
+            &mut this.bytes_read,
+            pos,
+        ))
     }
 }
 
@@ -357,21 +352,16 @@ impl Reader for DataReader {
         &'a mut self,
         buf: &'a mut Vec<u8>,
     ) -> stackfuture::StackFuture<'a, std::io::Result<usize>, { super::STACK_FUTURE_SIZE }> {
-        stackfuture::StackFuture::from(async {
-            if self.bytes_read >= self.data.value().len() {
-                Ok(0)
-            } else {
-                buf.extend_from_slice(&self.data.value()[self.bytes_read..]);
-                let n = self.data.value().len() - self.bytes_read;
-                self.bytes_read = self.data.value().len();
-                Ok(n)
-            }
-        })
+        crate::io::read_to_end(self.data.value(), &mut self.bytes_read, buf)
     }
 }
 
 impl AssetReader for MemoryAssetReader {
-    async fn read<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
+    async fn read<'a>(
+        &'a self,
+        path: &'a Path,
+        _required_features: ReaderRequiredFeatures,
+    ) -> Result<impl Reader + 'a, AssetReaderError> {
         self.root
             .get_asset(path)
             .map(|data| DataReader {