Move public entry point functions to beginning of module

Author: pepijnve

datafusion/execution/src/async_stream.rs

@@ -25,22 +25,129 @@ use std::sync::{Arc};
 use std::task::{Context, Poll};
 use parking_lot::Mutex;
 
+/// Creates a [`Stream`] from an async generator function.
+///
+/// The `generator` closure receives an [`Emitter<T>`] and runs as an async
+/// block. Each `emitter.emit(value).await` call suspends the generator and
+/// produces the next item in the stream. The stream ends when the generator
+/// future resolves.
+///
+/// # Example
+///
+/// ```
+/// use datafusion_execution::async_stream;
+/// use futures::StreamExt;
+///
+/// # #[tokio::main(flavor = "current_thread")]
+/// # async fn main() {
+/// let stream = async_stream(|mut emitter| async move {
+///     for i in 0_i32..3 {
+///         emitter.emit(i).await;
+///     }
+/// });
+///
+/// let values: Vec<i32> = stream.collect().await;
+/// assert_eq!(values, vec![0, 1, 2]);
+/// # }
+/// ```
+pub fn async_stream<T, F: Future<Output = ()>>(
+    generator: impl FnOnce(Emitter<T>) -> F,
+) -> impl FusedStream<Item = T> {
+    let (emitter, receiver) = tx_rx();
+    AsyncStream::new(receiver, generator(emitter))
+}
+
+/// Creates a fallible [`Stream`] from an async generator function.
+///
+/// The `generator` closure receives a [`TryEmitter<T, E>`] and runs as an
+/// async block that returns `Result<(), E>`. Each `emitter.emit(value).await`
+/// call suspends the generator and produces `Ok(value)` as the next stream
+/// item. The `?` operator can be used inside the generator to short-circuit on
+/// errors: the error is emitted as the final `Err(e)` item and the stream
+/// ends. The stream also ends when the generator future resolves to `Ok(())`.
+///
+/// # Example
+///
+/// ```
+/// use datafusion_execution::async_try_stream;
+/// use futures::StreamExt;
+///
+/// # #[tokio::main(flavor = "current_thread")]
+/// # async fn main() {
+/// let stream = async_try_stream(|mut emitter| async move {
+///     emitter.emit(1_i32).await;
+///     emitter.emit(2_i32).await;
+///     Err::<(), _>("something went wrong")?;
+///     emitter.emit(3_i32).await; // never reached
+///     Ok(())
+/// });
+///
+/// let values: Vec<Result<i32, &str>> = stream.collect().await;
+/// assert_eq!(values, vec![Ok(1), Ok(2), Err("something went wrong")]);
+/// # }
+/// ```
+pub fn async_try_stream<T, E, F: Future<Output = Result<(), E>>>(
+    generator: impl FnOnce(TryEmitter<T, E>) -> F,
+) -> impl FusedStream<Item = Result<T, E>> {
+    let (try_emitter, mut emitter, receiver) = try_tx_rx::<T, E>();
+    AsyncStream::new(receiver, async move {
+        if let Err(e) = generator(try_emitter).await {
+            emitter.emit(Err(e)).await
+        }
+    })
+}
+
+/// Creates an `Emitter`/`Receiver` pair
+fn tx_rx<T>() -> (Emitter<T>, Receiver<T>) {
+    let slot = Arc::new(Mutex::new(None));
+    (
+        Emitter {
+            slot: Arc::clone(&slot),
+        },
+        Receiver { slot },
+    )
+}
+
+/// Creates an `TryEmitter`/`Emitter`/`Receiver` triplet
+#[expect(
+    clippy::type_complexity,
+    reason = "three-element tuple is clearer than an alias here"
+)]
+fn try_tx_rx<T, E>() -> (
+    TryEmitter<T, E>,
+    Emitter<Result<T, E>>,
+    Receiver<Result<T, E>>,
+) {
+    let slot = Arc::new(Mutex::new(None));
+    (
+        TryEmitter {
+            slot: Arc::clone(&slot),
+        },
+        Emitter {
+            slot: Arc::clone(&slot),
+        },
+        Receiver { slot },
+    )
+}
+
+type SlotRef<T> = Arc<Mutex<Option<T>>>;
+
 /// A handle for emitting values from an [`async_stream`] generator.
 ///
 /// The generator closure receives an `Emitter<T>` as its argument.
 pub struct Emitter<T> {
-    slot: Arc<Mutex<Option<T>>>,
+    slot: SlotRef<T>,
 }
 
 /// A handle for emitting values from an [`async_try_stream`] generator.
 ///
 /// The generator closure receives a `TryEmitter<T, E>` as its argument.
 pub struct TryEmitter<T, E> {
-    slot: Arc<Mutex<Option<Result<T, E>>>>,
+    slot: SlotRef<T>,
 }
 
 struct Receiver<T> {
-    slot: Arc<Mutex<Option<T>>>,
+    slot: SlotRef<T>,
 }
 
 impl<T> Emitter<T> {
@@ -177,109 +284,6 @@ where
     }
 }
 
-fn tx_rx<T>() -> (Emitter<T>, Receiver<T>) {
-    let slot = Arc::new(Mutex::new(None));
-    (
-        Emitter {
-            slot: Arc::clone(&slot),
-        },
-        Receiver { slot },
-    )
-}
-
-/// Creates a [`Stream`] from an async generator function.
-///
-/// The `generator` closure receives an [`Emitter<T>`] and runs as an async
-/// block. Each `emitter.emit(value).await` call suspends the generator and
-/// produces the next item in the stream. The stream ends when the generator
-/// future resolves.
-///
-/// # Example
-///
-/// ```
-/// use datafusion_execution::async_stream;
-/// use futures::StreamExt;
-///
-/// # #[tokio::main(flavor = "current_thread")]
-/// # async fn main() {
-/// let stream = async_stream(|mut emitter| async move {
-///     for i in 0_i32..3 {
-///         emitter.emit(i).await;
-///     }
-/// });
-///
-/// let values: Vec<i32> = stream.collect().await;
-/// assert_eq!(values, vec![0, 1, 2]);
-/// # }
-/// ```
-pub fn async_stream<T, F: Future<Output = ()>>(
-    generator: impl FnOnce(Emitter<T>) -> F,
-) -> impl FusedStream<Item = T> {
-    let (emitter, receiver) = tx_rx();
-    AsyncStream::new(receiver, generator(emitter))
-}
-
-#[expect(
-    clippy::type_complexity,
-    reason = "three-element tuple is clearer than an alias here"
-)]
-fn try_tx_rx<T, E>() -> (
-    TryEmitter<T, E>,
-    Emitter<Result<T, E>>,
-    Receiver<Result<T, E>>,
-) {
-    let slot = Arc::new(Mutex::new(None));
-    (
-        TryEmitter {
-            slot: Arc::clone(&slot),
-        },
-        Emitter {
-            slot: Arc::clone(&slot),
-        },
-        Receiver { slot },
-    )
-}
-
-/// Creates a fallible [`Stream`] from an async generator function.
-///
-/// The `generator` closure receives a [`TryEmitter<T, E>`] and runs as an
-/// async block that returns `Result<(), E>`. Each `emitter.emit(value).await`
-/// call suspends the generator and produces `Ok(value)` as the next stream
-/// item. The `?` operator can be used inside the generator to short-circuit on
-/// errors: the error is emitted as the final `Err(e)` item and the stream
-/// ends. The stream also ends when the generator future resolves to `Ok(())`.
-///
-/// # Example
-///
-/// ```
-/// use datafusion_execution::async_try_stream;
-/// use futures::StreamExt;
-///
-/// # #[tokio::main(flavor = "current_thread")]
-/// # async fn main() {
-/// let stream = async_try_stream(|mut emitter| async move {
-///     emitter.emit(1_i32).await;
-///     emitter.emit(2_i32).await;
-///     Err::<(), _>("something went wrong")?;
-///     emitter.emit(3_i32).await; // never reached
-///     Ok(())
-/// });
-///
-/// let values: Vec<Result<i32, &str>> = stream.collect().await;
-/// assert_eq!(values, vec![Ok(1), Ok(2), Err("something went wrong")]);
-/// # }
-/// ```
-pub fn async_try_stream<T, E, F: Future<Output = Result<(), E>>>(
-    generator: impl FnOnce(TryEmitter<T, E>) -> F,
-) -> impl FusedStream<Item = Result<T, E>> {
-    let (try_emitter, mut emitter, receiver) = try_tx_rx::<T, E>();
-    AsyncStream::new(receiver, async move {
-        if let Err(e) = generator(try_emitter).await {
-            emitter.emit(Err(e)).await
-        }
-    })
-}
-
 #[cfg(test)]
 mod test {
     use crate::async_stream::Emitter;