Skip to main content

xet_runtime/utils/
async_iterator.rs

1use async_trait::async_trait;
2
3#[async_trait]
4pub trait AsyncIterator<E: Send + Sync + 'static>: Send + Sync {
5    type Item: Send + Sync;
6
7    /// The traditional next method for iterators, with a Result and Error
8    /// type.  Returns None when everything is done.
9    async fn next(&mut self) -> Result<Option<Self::Item>, E>;
10}
11
12#[async_trait]
13pub trait BatchedAsyncIterator<E: Send + Sync + 'static>: AsyncIterator<E> {
14    /// Return a block of items.  If the stream is done, then an empty vector is returned;
15    /// otherwise, at least one item is returned.
16    ///
17    /// If given, max_num dictates the maximum number of items to return.  If None, then all
18    /// available items are returned.
19    async fn next_batch(&mut self, max_num: Option<usize>) -> Result<Vec<Self::Item>, E>;
20
21    /// Returns the number of items remaining in the stream
22    /// if known, and None otherwise.  Returns Some(0) if
23    /// there are no items remaining.
24    fn items_remaining(&self) -> Option<usize>;
25}
26
27#[async_trait]
28impl<E: Send + Sync + 'static> AsyncIterator<E> for Vec<u8> {
29    type Item = Vec<u8>;
30
31    async fn next(&mut self) -> Result<Option<Self::Item>, E> {
32        if self.is_empty() {
33            Ok(None)
34        } else {
35            Ok(Some(std::mem::take(self)))
36        }
37    }
38}