Import async-stream-impl 0.3.4upstream/0.3.4upstream

5 files changed, 427 insertions, 0 deletions

diff --git a/.cargo_vcs_info.json b/.cargo_vcs_info.json
new file mode 100644
index 0000000..f6454dd
--- /dev/null
+++ b/.cargo_vcs_info.json
@@ -0,0 +1,6 @@
+{
+  "git": {
+    "sha1": "bf41b9645fe39b8865da2f25edc286eb42d49ec8"
+  },
+  "path_in_vcs": "async-stream-impl"
+}
\ No newline at end of file
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..4b74dfb
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,46 @@
+# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO
+#
+# When uploading crates to the registry Cargo will automatically
+# "normalize" Cargo.toml files for maximal compatibility
+# with all versions of Cargo and also rewrite `path` dependencies
+# to registry (e.g., crates.io) dependencies.
+#
+# If you are reading this file be aware that the original Cargo.toml
+# will likely look very different (and much more reasonable).
+# See Cargo.toml.orig for the original contents.
+
+[package]
+edition = "2018"
+rust-version = "1.45"
+name = "async-stream-impl"
+version = "0.3.4"
+authors = ["Carl Lerche <me@carllerche.com>"]
+description = "proc macros for async-stream crate"
+license = "MIT"
+repository = "https://github.com/tokio-rs/async-stream"
+
+[lib]
+proc-macro = true
+
+[dependencies.proc-macro2]
+version = "1"
+
+[dependencies.quote]
+version = "1"
+
+[dependencies.syn]
+version = "1"
+features = [
+    "full",
+    "visit-mut",
+]
+
+[dev-dependencies.futures-core]
+version = "0.3"
+
+[dev-dependencies.futures-util]
+version = "0.3"
+
+[dev-dependencies.tokio]
+version = "1"
+features = ["full"]
diff --git a/Cargo.toml.orig b/Cargo.toml.orig
new file mode 100644
index 0000000..fc0a8e3
--- /dev/null
+++ b/Cargo.toml.orig
@@ -0,0 +1,23 @@
+[package]
+name = "async-stream-impl"
+version = "0.3.4"
+edition = "2018"
+rust-version = "1.45"
+license = "MIT"
+authors = ["Carl Lerche <me@carllerche.com>"]
+description = "proc macros for async-stream crate"
+repository = "https://github.com/tokio-rs/async-stream"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro2 = "1"
+syn = { version = "1", features = ["full", "visit-mut"] }
+quote = "1"
+
+[dev-dependencies]
+async-stream = { path = "../async-stream" }
+futures-core = "0.3"
+futures-util = "0.3"
+tokio = { version = "1", features = ["full"] }
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..8cbd7d6
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,51 @@
+Copyright (c) 2019 Carl Lerche
+
+Permission is hereby granted, free of charge, to any
+person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the
+Software without restriction, including without
+limitation the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software
+is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice
+shall be included in all copies or substantial portions
+of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.
+
+Copyright (c) 2018 David Tolnay
+
+Permission is hereby granted, free of charge, to any
+person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the
+Software without restriction, including without
+limitation the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software
+is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice
+shall be included in all copies or substantial portions
+of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..082f734
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,301 @@
+use proc_macro::TokenStream;
+use proc_macro2::{Group, TokenStream as TokenStream2, TokenTree};
+use quote::quote;
+use syn::parse::{Parse, ParseStream, Parser, Result};
+use syn::visit_mut::VisitMut;
+
+struct Scrub<'a> {
+    /// Whether the stream is a try stream.
+    is_try: bool,
+    /// The unit expression, `()`.
+    unit: Box<syn::Expr>,
+    has_yielded: bool,
+    crate_path: &'a TokenStream2,
+}
+
+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<'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,
+        }
+    }
+}
+
+struct Partial<T>(T, TokenStream2);
+
+impl<T: Parse> Parse for Partial<T> {
+    fn parse(input: ParseStream) -> Result<Self> {
+        Ok(Partial(input.parse()?, input.parse()?))
+    }
+}
+
+fn visit_token_stream_impl(
+    visitor: &mut Scrub<'_>,
+    tokens: TokenStream2,
+    modified: &mut bool,
+    out: &mut TokenStream2,
+) {
+    use quote::ToTokens;
+    use quote::TokenStreamExt;
+
+    let mut tokens = tokens.into_iter().peekable();
+    while let Some(tt) = tokens.next() {
+        match tt {
+            TokenTree::Ident(i) if i == "yield" => {
+                let stream = std::iter::once(TokenTree::Ident(i)).chain(tokens).collect();
+                match syn::parse2(stream) {
+                    Ok(Partial(yield_expr, rest)) => {
+                        let mut expr = syn::Expr::Yield(yield_expr);
+                        visitor.visit_expr_mut(&mut expr);
+                        expr.to_tokens(out);
+                        *modified = true;
+                        tokens = rest.into_iter().peekable();
+                    }
+                    Err(e) => {
+                        out.append_all(&mut e.to_compile_error().into_iter());
+                        *modified = true;
+                        return;
+                    }
+                }
+            }
+            TokenTree::Ident(i) if i == "stream" || i == "try_stream" => {
+                out.append(TokenTree::Ident(i));
+                match tokens.peek() {
+                    Some(TokenTree::Punct(p)) if p.as_char() == '!' => {
+                        out.extend(tokens.next()); // !
+                        if let Some(TokenTree::Group(_)) = tokens.peek() {
+                            out.extend(tokens.next()); // { .. } or [ .. ] or ( .. )
+                        }
+                    }
+                    _ => {}
+                }
+            }
+            TokenTree::Group(group) => {
+                let mut content = group.stream();
+                *modified |= visitor.visit_token_stream(&mut content);
+                let mut new = Group::new(group.delimiter(), content);
+                new.set_span(group.span());
+                out.append(new);
+            }
+            other => out.append(other),
+        }
+    }
+}
+
+impl Scrub<'_> {
+    fn visit_token_stream(&mut self, tokens: &mut TokenStream2) -> bool {
+        let (mut out, mut modified) = (TokenStream2::new(), false);
+        visit_token_stream_impl(self, tokens.clone(), &mut modified, &mut out);
+
+        if modified {
+            *tokens = out;
+        }
+
+        modified
+    }
+}
+
+impl VisitMut for Scrub<'_> {
+    fn visit_expr_mut(&mut self, i: &mut syn::Expr) {
+        match i {
+            syn::Expr::Yield(yield_expr) => {
+                self.has_yielded = true;
+
+                syn::visit_mut::visit_expr_yield_mut(self, yield_expr);
+
+                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(::core::result::Result::Ok(#value_expr)).await }
+                } else {
+                    syn::parse_quote! { __yield_tx.send(#value_expr).await }
+                };
+            }
+            syn::Expr::Try(try_expr) => {
+                syn::visit_mut::visit_expr_try_mut(self, try_expr);
+                // let ident = &self.yielder;
+                let e = &try_expr.expr;
+
+                *i = syn::parse_quote! {
+                    match #e {
+                        ::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(_) => {
+                // Don't transform inner closures or async blocks.
+            }
+            syn::Expr::ForLoop(expr) => {
+                syn::visit_mut::visit_expr_for_loop_mut(self, expr);
+                // TODO: Should we allow other attributes?
+                if expr.attrs.len() != 1 || !expr.attrs[0].path.is_ident("await") {
+                    return;
+                }
+                let syn::ExprForLoop {
+                    attrs,
+                    label,
+                    pat,
+                    expr,
+                    body,
+                    ..
+                } = expr;
+
+                let attr = attrs.pop().unwrap();
+                if let Err(e) = syn::parse2::<syn::parse::Nothing>(attr.tokens) {
+                    *i = syn::parse2(e.to_compile_error()).unwrap();
+                    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 #crate_path::__private::next(&mut __pinned).await {
+                            ::core::option::Option::Some(e) => e,
+                            ::core::option::Option::None => break,
+                        };
+                        #body
+                    }
+                }}
+            }
+            _ => syn::visit_mut::visit_expr_mut(self, i),
+        }
+    }
+
+    fn visit_macro_mut(&mut self, mac: &mut syn::Macro) {
+        let mac_ident = mac.path.segments.last().map(|p| &p.ident);
+        if mac_ident.map_or(false, |i| i == "stream" || i == "try_stream") {
+            return;
+        }
+
+        self.visit_token_stream(&mut mac.tokens);
+    }
+
+    fn visit_item_mut(&mut self, i: &mut syn::Item) {
+        // Recurse into macros but otherwise don't transform inner items.
+        if let syn::Item::Macro(i) = i {
+            self.visit_macro_mut(&mut i.mac);
+        }
+    }
+}
+
+/// The first token tree in the stream must be a group containing the path to the `async-stream`
+/// crate.
+#[proc_macro]
+#[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::new(false, &crate_path);
+
+    for stmt in &mut stmts {
+        scrub.visit_stmt_mut(stmt);
+    }
+
+    let dummy_yield = if scrub.has_yielded {
+        None
+    } else {
+        Some(quote!(if false {
+            __yield_tx.send(()).await;
+        }))
+    };
+
+    quote!({
+        let (mut __yield_tx, __yield_rx) = unsafe { #crate_path::__private::yielder::pair() };
+        #crate_path::__private::AsyncStream::new(__yield_rx, async move {
+            #dummy_yield
+            #(#stmts)*
+        })
+    })
+    .into()
+}
+
+/// The first token tree in the stream must be a group containing the path to the `async-stream`
+/// crate.
+#[proc_macro]
+#[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::new(true, &crate_path);
+
+    for stmt in &mut stmts {
+        scrub.visit_stmt_mut(stmt);
+    }
+
+    let dummy_yield = if scrub.has_yielded {
+        None
+    } else {
+        Some(quote!(if false {
+            __yield_tx.send(()).await;
+        }))
+    };
+
+    quote!({
+        let (mut __yield_tx, __yield_rx) = unsafe { #crate_path::__private::yielder::pair() };
+        #crate_path::__private::AsyncStream::new(__yield_rx, async move {
+            #dummy_yield
+            #(#stmts)*
+        })
+    })
+    .into()
+}
+
+/// 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();
+
+    while let Some(token) = input.next() {
+        match token {
+            TokenTree::Ident(ident) => {
+                match input.peek() {
+                    Some(TokenTree::Ident(next)) if ident == "for" && next == "await" => {
+                        tokens.extend(quote!(#[#next]));
+                        let _ = input.next();
+                    }
+                    _ => {}
+                }
+                tokens.push(ident.into());
+            }
+            TokenTree::Group(group) => {
+                let stream = replace_for_await(group.stream());
+                let mut new_group = Group::new(group.delimiter(), stream);
+                new_group.set_span(group.span());
+                tokens.push(new_group.into());
+            }
+            _ => tokens.push(token),
+        }
+    }
+
+    tokens.into_iter().collect()
+}