Snap for 8564071 from d4bfaa8ce0356656942203fbdc58919a92882127 to mainline-os-statsd-releaseaml_sta_331910000aml_sta_331811000aml_sta_331711010aml_sta_331610000aml_sta_331511000aml_sta_331410000aml_sta_331311000aml_sta_331010010aml_sta_330910000android13-mainline-os-statsd-release

Change-Id: Ic211f0223d28eaa433e9c9cacfcb1e1540fd6591

8 files changed, 174 insertions, 147 deletions

diff --git a/.cargo_vcs_info.json b/.cargo_vcs_info.json
new file mode 100644
index 0000000..620b2af
--- /dev/null
+++ b/.cargo_vcs_info.json
@@ -0,0 +1,5 @@
+{
+  "git": {
+    "sha1": "b28da881695e3c66e1782f0c2c330d4e162eb7c2"
+  }
+}
diff --git a/Android.bp b/Android.bp
index a212b3b..7e1ac13 100644
--- a/Android.bp
+++ b/Android.bp
@@ -1,4 +1,5 @@
-// This file is generated by cargo2android.py --run --dependencies --device.
+// This file is generated by cargo2android.py --config cargo2android.json.
+// Do not modify this file as changes will be overridden on upgrade.
 
 package {
     default_applicable_licenses: [
@@ -22,6 +23,8 @@ license {
 rust_proc_macro {
     name: "libasync_stream_impl",
     crate_name: "async_stream_impl",
+    cargo_env_compat: true,
+    cargo_pkg_version: "0.3.2",
     srcs: ["src/lib.rs"],
     edition: "2018",
     rustlibs: [
@@ -30,9 +33,3 @@ rust_proc_macro {
         "libsyn",
     ],
 }
-
-// dependent_library ["feature_list"]
-//   proc-macro2-1.0.24 "default,proc-macro"
-//   quote-1.0.8 "default,proc-macro"
-//   syn-1.0.60 "clone-impls,default,derive,extra-traits,full,parsing,printing,proc-macro,quote,visit-mut"
-//   unicode-xid-0.2.1 "default"
diff --git a/Cargo.toml b/Cargo.toml
index df59eae..44a0cb9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,11 +13,11 @@
 [package]
 edition = "2018"
 name = "async-stream-impl"
-version = "0.3.0"
+version = "0.3.2"
 authors = ["Carl Lerche <me@carllerche.com>"]
 description = "proc macros for async-stream crate"
 homepage = "https://github.com/tokio-rs/async-stream"
-documentation = "https://docs.rs/async-stream-impl/0.3.0/async-stream-impl"
+documentation = "https://docs.rs/async-stream-impl"
 license = "MIT"
 repository = "https://github.com/tokio-rs/async-stream"
 
@@ -39,5 +39,5 @@ version = "0.3"
 version = "0.3"
 
 [dev-dependencies.tokio]
-version = "0.2"
+version = "1"
 features = ["full"]
diff --git a/Cargo.toml.orig b/Cargo.toml.orig
index dd6d9d9..4317fdf 100644
--- a/Cargo.toml.orig
+++ b/Cargo.toml.orig
@@ -1,11 +1,11 @@
 [package]
 name = "async-stream-impl"
-version = "0.3.0"
+version = "0.3.2"
 edition = "2018"
 license = "MIT"
 authors = ["Carl Lerche <me@carllerche.com>"]
 description = "proc macros for async-stream crate"
-documentation = "https://docs.rs/async-stream-impl/0.3.0/async-stream-impl"
+documentation = "https://docs.rs/async-stream-impl"
 homepage = "https://github.com/tokio-rs/async-stream"
 repository = "https://github.com/tokio-rs/async-stream"
 
@@ -18,7 +18,7 @@ syn = { version = "1", features = ["extra-traits", "full", "visit-mut"]}
 quote = "1"
 
 [dev-dependencies]
-# async-stream = { version = "0.3.0", path = "../async-stream" }
+async-stream = { path = "../async-stream" }
 futures-core = "0.3"
 futures-util = "0.3"
-tokio = { version = "0.2", features = ["full"] }
+tokio = { version = "1", features = ["full"] }
diff --git a/METADATA b/METADATA
index 8e2e26d..f08d3d2 100644
--- a/METADATA
+++ b/METADATA
@@ -7,13 +7,13 @@ third_party {
   }
   url {
     type: ARCHIVE
-    value: "https://static.crates.io/crates/async-stream-impl/async-stream-impl-0.3.0.crate"
+    value: "https://static.crates.io/crates/async-stream-impl/async-stream-impl-0.3.2.crate"
   }
-  version: "0.3.0"
+  version: "0.3.2"
   license_type: NOTICE
   last_upgrade_date {
     year: 2021
-    month: 2
-    day: 8
+    month: 6
+    day: 21
   }
 }
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 8fccfe0..dfc3524 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -1,17 +1,11 @@
 // Generated by update_crate_tests.py for tests that depend on this crate.
 {
-  "presubmit": [
+  "imports": [
     {
-      "name": "tokio-test_device_test_tests_block_on"
+      "path": "external/rust/crates/tokio"
     },
     {
-      "name": "tokio-test_device_test_tests_io"
-    },
-    {
-      "name": "tokio-test_device_test_src_lib"
-    },
-    {
-      "name": "tokio-test_device_test_tests_macros"
+      "path": "external/rust/crates/tokio-test"
     }
   ]
 }
diff --git a/cargo2android.json b/cargo2android.json
new file mode 100644
index 0000000..bf78496
--- /dev/null
+++ b/cargo2android.json
@@ -0,0 +1,4 @@
+{
+  "device": true,
+  "run": true
+}
\ No newline at end of file
diff --git a/src/lib.rs b/src/lib.rs
index 9081ecf..81d3752 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,47 +1,126 @@
 extern crate proc_macro;
 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::{Parse, ParseStream, Parser, Result};
 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)?;
+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))
+}
 
-    Ok(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,
+        }
+    }
 }
 
-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;
+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.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 }
                 };
@@ -53,19 +132,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);
@@ -88,6 +164,7 @@ impl VisitMut for Scrub {
                     return;
                 }
 
+                let crate_path = self.crate_path;
                 *i = syn::parse_quote! {{
                     let mut __pinned = #expr;
                     let mut __pinned = unsafe {
@@ -95,7 +172,7 @@ impl VisitMut for Scrub {
                     };
                     #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,
                         };
@@ -107,70 +184,50 @@ impl VisitMut for Scrub {
         }
     }
 
+    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) {
-        let prev = self.is_xforming;
-        self.is_xforming = false;
-        syn::visit_mut::visit_item_mut(self, i);
-        self.is_xforming = prev;
+        // Recurse into macros but otherwise don't transform inner items.
+        if let syn::Item::Macro(i) = i {
+            self.visit_macro_mut(&mut i.mac);
+        }
     }
 }
 
-/// 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)*
         })
@@ -178,64 +235,33 @@ pub fn stream(input: TokenStream) -> TokenStream {
     .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)*
         })
@@ -243,7 +269,8 @@ pub fn try_stream(input: TokenStream) -> TokenStream {
     .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();