Skip to main content

web_sys_async_io/
reader.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{ready, Poll},
5};
6
7use wasm_bindgen_futures::JsFuture;
8
9#[derive(Debug, Default)]
10pub enum Op {
11    #[default]
12    Idle,
13    ReadPending(JsFuture),
14    ConsumingReadBuffer {
15        read_buffer: js_sys::Uint8Array,
16        already_read: usize,
17    },
18}
19
20#[derive(Debug)]
21pub struct Reader {
22    pub inner: web_sys::ReadableStreamByobReader,
23    pub op: Op,
24    pub internal_buf: Option<js_sys::ArrayBuffer>,
25}
26
27impl Reader {
28    pub fn new(inner: web_sys::ReadableStreamByobReader) -> Self {
29        Self {
30            inner,
31            op: Op::default(),
32            internal_buf: None,
33        }
34    }
35
36    pub fn with_buf(
37        inner: web_sys::ReadableStreamByobReader,
38        internal_buf: js_sys::ArrayBuffer,
39    ) -> Self {
40        Self {
41            inner,
42            op: Op::default(),
43            internal_buf: Some(internal_buf),
44        }
45    }
46}
47
48impl tokio::io::AsyncRead for Reader {
49    fn poll_read(
50        mut self: Pin<&mut Self>,
51        cx: &mut std::task::Context<'_>,
52        buf: &mut tokio::io::ReadBuf<'_>,
53    ) -> Poll<std::io::Result<()>> {
54        // A read into a buffer with no remaining capacity must complete
55        // immediately without requesting more data from the stream.
56        if buf.remaining() == 0 {
57            return Poll::Ready(Ok(()));
58        }
59
60        match self.op {
61            Op::ReadPending(ref mut fut) => {
62                let result = ready!(Pin::new(fut).poll(cx));
63
64                let read_result = match result {
65                    Ok(val) => val,
66                    Err(err) => return Poll::Ready(Err(super::js_value_to_io_error(err))),
67                };
68                let read_result: crate::sys::ReadableStreamByobReaderValue = read_result.into();
69
70                let value = read_result.value();
71                // No value indicates an error condition or end of stream.
72                let Some(internal_buf_view) = value else {
73                    self.op = Op::Idle;
74                    return Poll::Ready(Ok(()));
75                };
76
77                self.op = Op::ConsumingReadBuffer {
78                    read_buffer: internal_buf_view,
79                    already_read: 0,
80                };
81
82                self.poll_read(cx, buf)
83            }
84            Op::ConsumingReadBuffer {
85                ref mut read_buffer,
86                already_read,
87            } => {
88                let remaining_size = read_buffer.byte_length() as usize - already_read;
89
90                let buf_remaining_size = buf.remaining();
91                let copy_size = remaining_size.min(buf_remaining_size);
92
93                let write_slice = buf.initialize_unfilled_to(copy_size);
94                let source_view = js_sys::Uint8Array::new_with_byte_offset_and_length(
95                    &read_buffer.buffer(),
96                    already_read as u32,
97                    copy_size as u32,
98                );
99                source_view.copy_to(&mut write_slice[..copy_size]);
100                buf.advance(copy_size);
101
102                if remaining_size <= buf_remaining_size {
103                    // The buffer returned is actually the same buffer we passed
104                    // earlier when we called the `read_with_array_buffer_view`
105                    // under the hood - despite it being an entirely new JS object.
106                    // We now have to assume the ownership of the buffer and
107                    // properly keep for the next time.
108                    self.internal_buf = Some(read_buffer.buffer());
109                    self.op = Op::Idle;
110                    Poll::Ready(Ok(()))
111                } else {
112                    self.op = Op::ConsumingReadBuffer {
113                        read_buffer: read_buffer.clone(),
114                        already_read: already_read + copy_size,
115                    };
116                    Poll::Ready(Ok(()))
117                }
118            }
119            Op::Idle => {
120                let requested_size = buf.remaining().try_into().unwrap();
121                let internal_buf = self
122                    .internal_buf
123                    .take()
124                    .filter(|internal_buf| {
125                        let actual_size = internal_buf.byte_length();
126                        debug_assert!(actual_size > 0);
127                        actual_size >= requested_size
128                    })
129                    .unwrap_or_else(|| js_sys::ArrayBuffer::new(requested_size));
130                let internal_buf_view = js_sys::Uint8Array::new(&internal_buf);
131                // Despite this not being properly indicated at the type system,
132                // the `read_with_array_buffer_view` fn is actually supposed
133                // to be taking the buffer by value - as it takes the ownership of
134                // the buffer and the old JS reference to it is no longer valid.
135                let fut =
136                    JsFuture::from(self.inner.read_with_array_buffer_view(&internal_buf_view));
137                self.op = Op::ReadPending(fut);
138                self.poll_read(cx, buf)
139            }
140        }
141    }
142}