From da7886fddcb3462143e02fa3add1b90a1c262d80 Mon Sep 17 00:00:00 2001
From: Aumetra Weisman <aumetra@cryptolab.net>
Date: Wed, 26 Jun 2024 23:38:46 +0200
Subject: [PATCH] Fix unsoundness issues #106 and #107

---
 async-stream-impl/src/lib.rs | 27 ++++++++++++++++++++-------
 async-stream/src/yielder.rs  |  2 +-
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/async-stream-impl/src/lib.rs b/async-stream-impl/src/lib.rs
index d3eab1a..0881922 100644
--- a/async-stream-impl/src/lib.rs
+++ b/async-stream-impl/src/lib.rs
@@ -120,10 +120,19 @@ impl VisitMut for Scrub<'_> {
 
                 // let ident = &self.yielder;
 
-                *i = if self.is_try {
-                    syn::parse_quote! { __yield_tx.send(::core::result::Result::Ok(#value_expr)).await }
+                let yield_expr = if self.is_try {
+                    quote! { __yield_tx.send(::core::result::Result::Ok(#value_expr)).await }
                 } else {
-                    syn::parse_quote! { __yield_tx.send(#value_expr).await }
+                    quote! { __yield_tx.send(#value_expr).await }
+                };
+                *i = syn::parse_quote! {
+                    {
+                        #[allow(unreachable_code)]
+                        if false {
+                            break 'check_stream_scope (loop {});
+                        }
+                        #yield_expr
+                    }
                 };
             }
             syn::Expr::Try(try_expr) => {
@@ -225,8 +234,10 @@ pub fn stream_inner(input: TokenStream) -> TokenStream {
     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)*
+            'check_stream_scope: {
+                #dummy_yield
+                #(#stmts)*
+            }
         })
     })
     .into()
@@ -259,8 +270,10 @@ pub fn try_stream_inner(input: TokenStream) -> TokenStream {
     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)*
+            'check_stream_scope: {
+                #dummy_yield
+                #(#stmts)*
+            }
         })
     })
     .into()
diff --git a/async-stream/src/yielder.rs b/async-stream/src/yielder.rs
index 597e1c9..764c983 100644
--- a/async-stream/src/yielder.rs
+++ b/async-stream/src/yielder.rs
@@ -7,7 +7,7 @@ use std::task::{Context, Poll};
 
 #[derive(Debug)]
 pub struct Sender<T> {
-    _p: PhantomData<T>,
+    _p: PhantomData<fn(T) -> T>,
 }
 
 #[derive(Debug)]