Support reexporting from anywhere by wrapping in a declarative macro (#46)

* Refactor Scrub

* Wrap macro in a declarative macro and use $crate

As a downside we get slightly worse error messages as the filename and
line number will point to the macro definition site instead of the call
site.

Author: Kestrer

Cargo.toml

@@ -2,4 +2,6 @@
 members = [
     "async-stream",
     "async-stream-impl",
+    "async-stream/tests/reexporter",
+    "async-stream/tests/test-reexport",
 ]

async-stream-impl/src/lib.rs

@@ -1,46 +1,51 @@
 use proc_macro::TokenStream;
-use proc_macro2::{Delimiter, Group, TokenStream as TokenStream2, TokenTree};
+use proc_macro2::{Group, TokenStream as TokenStream2, TokenTree};
 use quote::quote;
+use syn::parse::Parser;
 use syn::visit_mut::VisitMut;
 
-struct Scrub {
-    is_xforming: bool,
+struct Scrub<'a> {
+    /// Whether the stream is a try stream.
     is_try: bool,
+    /// The unit expression, `()`.
     unit: Box<syn::Expr>,
-    num_yield: u32,
+    has_yielded: bool,
+    crate_path: &'a TokenStream2,
 }
 
-fn parse_input(input: TokenStream) -> syn::Result<Vec<syn::Stmt>> {
-    let input = replace_for_await(input.into());
-    // syn does not provide a way to parse `Vec<Stmt>` directly from `TokenStream`,
-    // so wrap input in a brace and then parse it as a block.
-    let input = TokenStream2::from(TokenTree::Group(Group::new(Delimiter::Brace, input)));
-    let syn::Block { stmts, .. } = syn::parse2(input)?;
-
-    Ok(stmts)
+fn parse_input(input: TokenStream) -> syn::Result<(TokenStream2, Vec<syn::Stmt>)> {
+    let mut input = TokenStream2::from(input).into_iter();
+    let crate_path = match input.next().unwrap() {
+        TokenTree::Group(group) => group.stream(),
+        _ => panic!(),
+    };
+    let stmts = syn::Block::parse_within.parse2(replace_for_await(input))?;
+    Ok((crate_path, stmts))
 }
 
-impl VisitMut for Scrub {
-    fn visit_expr_mut(&mut self, i: &mut syn::Expr) {
-        if !self.is_xforming {
-            syn::visit_mut::visit_expr_mut(self, i);
-            return;
+impl<'a> Scrub<'a> {
+    fn new(is_try: bool, crate_path: &'a TokenStream2) -> Self {
+        Self {
+            is_try,
+            unit: syn::parse_quote!(()),
+            has_yielded: false,
+            crate_path,
         }
+    }
+}
 
+impl VisitMut for Scrub<'_> {
+    fn visit_expr_mut(&mut self, i: &mut syn::Expr) {
         match i {
             syn::Expr::Yield(yield_expr) => {
-                self.num_yield += 1;
+                self.has_yielded = true;
 
-                let value_expr = if let Some(ref e) = yield_expr.expr {
-                    e
-                } else {
-                    &self.unit
-                };
+                let value_expr = yield_expr.expr.as_ref().unwrap_or(&self.unit);
 
                 // let ident = &self.yielder;
 
                 *i = if self.is_try {
-                    syn::parse_quote! { __yield_tx.send(Ok(#value_expr)).await }
+                    syn::parse_quote! { __yield_tx.send(::core::result::Result::Ok(#value_expr)).await }
                 } else {
                     syn::parse_quote! { __yield_tx.send(#value_expr).await }
                 };
@@ -52,19 +57,16 @@ impl VisitMut for Scrub {
 
                 *i = syn::parse_quote! {
                     match #e {
-                        Ok(v) => v,
-                        Err(e) => {
-                            __yield_tx.send(Err(e.into())).await;
+                        ::core::result::Result::Ok(v) => v,
+                        ::core::result::Result::Err(e) => {
+                            __yield_tx.send(::core::result::Result::Err(e.into())).await;
                             return;
                         }
                     }
                 };
             }
             syn::Expr::Closure(_) | syn::Expr::Async(_) => {
-                let prev = self.is_xforming;
-                self.is_xforming = false;
-                syn::visit_mut::visit_expr_mut(self, i);
-                self.is_xforming = prev;
+                // Don't transform inner closures or async blocks.
             }
             syn::Expr::ForLoop(expr) => {
                 syn::visit_mut::visit_expr_for_loop_mut(self, expr);
@@ -87,14 +89,15 @@ impl VisitMut for Scrub {
                     return;
                 }
 
+                let crate_path = self.crate_path;
                 *i = syn::parse_quote! {{
                     let mut __pinned = #expr;
                     let mut __pinned = unsafe {
                         ::core::pin::Pin::new_unchecked(&mut __pinned)
                     };
                     #label
                     loop {
-                        let #pat = match ::async_stream::reexport::next(&mut __pinned).await {
+                        let #pat = match #crate_path::reexport::next(&mut __pinned).await {
                             ::core::option::Option::Some(e) => e,
                             ::core::option::Option::None => break,
                         };
@@ -106,143 +109,81 @@ impl VisitMut for Scrub {
         }
     }
 
-    fn visit_item_mut(&mut self, i: &mut syn::Item) {
-        let prev = self.is_xforming;
-        self.is_xforming = false;
-        syn::visit_mut::visit_item_mut(self, i);
-        self.is_xforming = prev;
+    fn visit_item_mut(&mut self, _: &mut syn::Item) {
+        // Don't transform inner items.
     }
 }
 
-/// Asynchronous stream
-///
-/// See [crate](index.html) documentation for more details.
-///
-/// # Examples
-///
-/// ```rust
-/// use async_stream::stream;
-///
-/// use futures_util::pin_mut;
-/// use futures_util::stream::StreamExt;
-///
-/// #[tokio::main]
-/// async fn main() {
-///     let s = stream! {
-///         for i in 0..3 {
-///             yield i;
-///         }
-///     };
-///
-///     pin_mut!(s); // needed for iteration
-///
-///     while let Some(value) = s.next().await {
-///         println!("got {}", value);
-///     }
-/// }
-/// ```
+/// The first token tree in the stream must be a group containing the path to the `async-stream`
+/// crate.
 #[proc_macro]
-pub fn stream(input: TokenStream) -> TokenStream {
-    let mut stmts = match parse_input(input) {
+#[doc(hidden)]
+pub fn stream_inner(input: TokenStream) -> TokenStream {
+    let (crate_path, mut stmts) = match parse_input(input) {
         Ok(x) => x,
         Err(e) => return e.to_compile_error().into(),
     };
 
-    let mut scrub = Scrub {
-        is_xforming: true,
-        is_try: false,
-        unit: syn::parse_quote!(()),
-        num_yield: 0,
-    };
+    let mut scrub = Scrub::new(false, &crate_path);
 
-    for mut stmt in &mut stmts[..] {
+    for mut stmt in &mut stmts {
         scrub.visit_stmt_mut(&mut stmt);
     }
 
-    let dummy_yield = if scrub.num_yield == 0 {
+    let dummy_yield = if scrub.has_yielded {
+        None
+    } else {
         Some(quote!(if false {
             __yield_tx.send(()).await;
         }))
-    } else {
-        None
     };
 
     quote!({
-        let (mut __yield_tx, __yield_rx) = ::async_stream::yielder::pair();
-        ::async_stream::AsyncStream::new(__yield_rx, async move {
+        let (mut __yield_tx, __yield_rx) = #crate_path::yielder::pair();
+        #crate_path::AsyncStream::new(__yield_rx, async move {
             #dummy_yield
             #(#stmts)*
         })
     })
     .into()
 }
 
-/// Asynchronous fallible stream
-///
-/// See [crate](index.html) documentation for more details.
-///
-/// # Examples
-///
-/// ```rust
-/// use tokio::net::{TcpListener, TcpStream};
-///
-/// use async_stream::try_stream;
-/// use futures_core::stream::Stream;
-///
-/// use std::io;
-/// use std::net::SocketAddr;
-///
-/// fn bind_and_accept(addr: SocketAddr)
-///     -> impl Stream<Item = io::Result<TcpStream>>
-/// {
-///     try_stream! {
-///         let mut listener = TcpListener::bind(addr).await?;
-///
-///         loop {
-///             let (stream, addr) = listener.accept().await?;
-///             println!("received on {:?}", addr);
-///             yield stream;
-///         }
-///     }
-/// }
-/// ```
+/// The first token tree in the stream must be a group containing the path to the `async-stream`
+/// crate.
 #[proc_macro]
-pub fn try_stream(input: TokenStream) -> TokenStream {
-    let mut stmts = match parse_input(input) {
+#[doc(hidden)]
+pub fn try_stream_inner(input: TokenStream) -> TokenStream {
+    let (crate_path, mut stmts) = match parse_input(input) {
         Ok(x) => x,
         Err(e) => return e.to_compile_error().into(),
     };
 
-    let mut scrub = Scrub {
-        is_xforming: true,
-        is_try: true,
-        unit: syn::parse_quote!(()),
-        num_yield: 0,
-    };
+    let mut scrub = Scrub::new(true, &crate_path);
 
-    for mut stmt in &mut stmts[..] {
+    for mut stmt in &mut stmts {
         scrub.visit_stmt_mut(&mut stmt);
     }
 
-    let dummy_yield = if scrub.num_yield == 0 {
+    let dummy_yield = if scrub.has_yielded {
+        None
+    } else {
         Some(quote!(if false {
             __yield_tx.send(()).await;
         }))
-    } else {
-        None
     };
 
     quote!({
-        let (mut __yield_tx, __yield_rx) = ::async_stream::yielder::pair();
-        ::async_stream::AsyncStream::new(__yield_rx, async move {
+        let (mut __yield_tx, __yield_rx) = #crate_path::yielder::pair();
+        #crate_path::AsyncStream::new(__yield_rx, async move {
             #dummy_yield
             #(#stmts)*
         })
     })
     .into()
 }
 
-fn replace_for_await(input: TokenStream2) -> TokenStream2 {
+/// Replace `for await` with `#[await] for`, which will be later transformed into a `next` loop.
+fn replace_for_await(input: impl IntoIterator<Item = TokenTree>) -> TokenStream2 {
     let mut input = input.into_iter().peekable();
     let mut tokens = Vec::new();
 

async-stream/src/lib.rs

@@ -162,11 +162,82 @@ mod next;
 #[doc(hidden)]
 pub mod yielder;
 
-// Used by the macro, but not intended to be accessed publically.
+// Used by the macro, but not intended to be accessed publicly.
 #[doc(hidden)]
 pub use crate::async_stream::AsyncStream;
 
-pub use async_stream_impl::{stream, try_stream};
+#[doc(hidden)]
+pub use async_stream_impl;
+
+/// Asynchronous stream
+///
+/// See [crate](index.html) documentation for more details.
+///
+/// # Examples
+///
+/// ```
+/// use async_stream::stream;
+///
+/// use futures_util::pin_mut;
+/// use futures_util::stream::StreamExt;
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let s = stream! {
+///         for i in 0..3 {
+///             yield i;
+///         }
+///     };
+///
+///     pin_mut!(s); // needed for iteration
+///
+///     while let Some(value) = s.next().await {
+///         println!("got {}", value);
+///     }
+/// }
+/// ```
+#[macro_export]
+macro_rules! stream {
+    ($($tt:tt)*) => {
+        $crate::async_stream_impl::stream_inner!(($crate) $($tt)*)
+    }
+}
+
+/// Asynchronous fallible stream
+///
+/// See [crate](index.html) documentation for more details.
+///
+/// # Examples
+///
+/// ```
+/// use tokio::net::{TcpListener, TcpStream};
+///
+/// use async_stream::try_stream;
+/// use futures_core::stream::Stream;
+///
+/// use std::io;
+/// use std::net::SocketAddr;
+///
+/// fn bind_and_accept(addr: SocketAddr)
+///     -> impl Stream<Item = io::Result<TcpStream>>
+/// {
+///     try_stream! {
+///         let mut listener = TcpListener::bind(addr).await?;
+///
+///         loop {
+///             let (stream, addr) = listener.accept().await?;
+///             println!("received on {:?}", addr);
+///             yield stream;
+///         }
+///     }
+/// }
+/// ```
+#[macro_export]
+macro_rules! try_stream {
+    ($($tt:tt)*) => {
+        $crate::async_stream_impl::try_stream_inner!(($crate) $($tt)*)
+    }
+}
 
 #[doc(hidden)]
 pub mod reexport {

async-stream/tests/reexporter/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "reexporter"
+version = "0.0.0"
+authors = []
+edition = "2018"
+publish = false
+
+[dependencies]
+async-stream = { path = "../.." }

async-stream/tests/reexporter/src/lib.rs

@@ -0,0 +1,3 @@
+//! A crate to test reexporting the `stream!` and `try_stream!` macros.
+
+pub use async_stream::{stream, try_stream};