From 2b10c4861938197551de6625a1c190d436d3c074 Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sun, 10 Sep 2023 23:45:24 -0400
Subject: [PATCH] Update io-uring to use streem for file bytes

---
 src/file.rs   | 134 +++++++++++++-------------------------------------
 src/stream.rs |   1 +
 2 files changed, 36 insertions(+), 99 deletions(-)

diff --git a/src/file.rs b/src/file.rs
index a8b0fc1..a2408df 100644
--- a/src/file.rs
+++ b/src/file.rs
@@ -113,10 +113,7 @@ mod io_uring {
     use std::{
         convert::TryInto,
         fs::Metadata,
-        future::Future,
         path::{Path, PathBuf},
-        pin::Pin,
-        task::{Context, Poll},
     };
     use streem::IntoStreamer;
     use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
@@ -319,20 +316,7 @@ mod io_uring {
             from_start: Option<u64>,
             len: Option<u64>,
         ) -> Result<impl Stream<Item = std::io::Result<Bytes>>, FileError> {
-            let size = self.metadata().await?.len();
-
-            let cursor = from_start.unwrap_or(0);
-            let size = len.unwrap_or(size - cursor) + cursor;
-
-            Ok(BytesStream {
-                state: ReadFileState::File {
-                    file: Some(self),
-                    bytes: Some(BytesMut::new()),
-                },
-                size,
-                cursor,
-                callback: read_file,
-            })
+            Ok(bytes_stream(self, from_start, len))
         }
 
         async fn read_at<T: IoBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T> {
@@ -344,98 +328,50 @@ mod io_uring {
         }
     }
 
-    pin_project_lite::pin_project! {
-        struct BytesStream<F, Fut> {
-            #[pin]
-            state: ReadFileState<Fut>,
-            size: u64,
-            cursor: u64,
-            #[pin]
-            callback: F,
-        }
-    }
-
-    pin_project_lite::pin_project! {
-        #[project = ReadFileStateProj]
-        #[project_replace = ReadFileStateProjReplace]
-        enum ReadFileState<Fut> {
-            File {
-                file: Option<File>,
-                bytes: Option<BytesMut>,
-            },
-            Future {
-                #[pin]
-                fut: Fut,
-            },
-        }
-    }
-
-    async fn read_file(
+    fn bytes_stream(
         file: File,
-        buf: BytesMut,
-        cursor: u64,
-    ) -> (File, BufResult<usize, BytesMut>) {
-        let buf_res = file.read_at(buf, cursor).await;
+        from_start: Option<u64>,
+        len: Option<u64>,
+    ) -> impl Stream<Item = std::io::Result<Bytes>> {
+        streem::try_from_fn(|yielder| async move {
+            let file_size = file.metadata().await?.len();
 
-        (file, buf_res)
-    }
+            let mut cursor = from_start.unwrap_or(0);
+            let remaining_size = file_size.saturating_sub(cursor);
+            let read_until = len.unwrap_or(remaining_size) + cursor;
 
-    impl<F, Fut> Stream for BytesStream<F, Fut>
-    where
-        F: Fn(File, BytesMut, u64) -> Fut,
-        Fut: Future<Output = (File, BufResult<usize, BytesMut>)> + 'static,
-    {
-        type Item = std::io::Result<Bytes>;
+            let mut bytes = BytesMut::new();
 
-        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-            let mut this = self.as_mut().project();
+            loop {
+                let max_size = read_until.saturating_sub(cursor);
 
-            match this.state.as_mut().project() {
-                ReadFileStateProj::File { file, bytes } => {
-                    let cursor = *this.cursor;
-                    let max_size = *this.size - *this.cursor;
-
-                    if max_size == 0 {
-                        return Poll::Ready(None);
-                    }
-
-                    let capacity = max_size.min(65_356) as usize;
-                    let mut bytes = bytes.take().unwrap();
-                    let file = file.take().unwrap();
-
-                    if bytes.capacity() < capacity {
-                        bytes.reserve(capacity - bytes.capacity());
-                    }
-
-                    let fut = (this.callback)(file, bytes, cursor);
-
-                    this.state.project_replace(ReadFileState::Future { fut });
-                    self.poll_next(cx)
+                if max_size == 0 {
+                    break;
                 }
-                ReadFileStateProj::Future { fut } => match fut.poll(cx) {
-                    Poll::Pending => Poll::Pending,
-                    Poll::Ready((file, (Ok(n), mut buf))) => {
-                        let bytes = buf.split_off(n);
 
-                        this.state.project_replace(ReadFileState::File {
-                            file: Some(file),
-                            bytes: Some(bytes),
-                        });
+                let capacity = max_size.min(65_356) as usize;
 
-                        let n: u64 = match n.try_into() {
-                            Ok(n) => n,
-                            Err(_) => {
-                                return Poll::Ready(Some(Err(std::io::ErrorKind::Other.into())))
-                            }
-                        };
-                        *this.cursor += n;
+                if bytes.capacity() < capacity {
+                    bytes.reserve(capacity - bytes.capacity());
+                }
 
-                        Poll::Ready(Some(Ok(buf.into())))
-                    }
-                    Poll::Ready((_, (Err(e), _))) => Poll::Ready(Some(Err(e))),
-                },
+                let (result, mut buf_) = file.read_at(bytes, cursor).await;
+
+                let n = match result {
+                    Ok(n) => n,
+                    Err(e) => return Err(e),
+                };
+
+                bytes = buf_.split_off(n);
+
+                let n: u64 = n.try_into().map_err(|_| std::io::ErrorKind::Other)?;
+                cursor += n;
+
+                yielder.yield_ok(buf_.into()).await;
             }
-        }
+
+            Ok(())
+        })
     }
 
     #[cfg(test)]
diff --git a/src/stream.rs b/src/stream.rs
index f7d4449..9e617aa 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -58,6 +58,7 @@ where
     })
 }
 
+#[cfg(not(feature = "io-uring"))]
 pub(crate) fn map_ok<S, T1, T2, E, F>(stream: S, f: F) -> impl Stream<Item = Result<T2, E>>
 where
     S: Stream<Item = Result<T1, E>>,