Rust in Everything
基于 async-stream 0.3.6

从零写一个最小 stream!

主线位置：两段都自己走一遍。宏展开一段：yield v 改写成 __tx.send(v).await，整块包进 MiniStream::new(rx, async move { .. })；运行时一段：poll_next 在栈上放一个 Option<T>，把地址写进线程局部变量，SendFut::poll 把值写过去，返回 Pending。

flowchart TB
  src["stream! { ... yield v; ... }"] --> sc["Scrub：yield 改写"]
  sc --> exp["MiniStream::new(rx, async move { '__mini_scope: { .. } })"]
  exp --> pn["poll_next：slot = None，STORE = &mut slot"]
  pn --> fut["poll async 块"]
  fut --> sf["SendFut::poll：写进 slot，Pending"]
  sf --> back["恢复 STORE；slot 有值就 Ready(Some(v))"]

本页只讲思路和关键代码。工程是一个三个包的 workspace：mini（运行时，107 行，去掉空行和注释 82 行）、mini-macro（proc-macro，91 行，去掉空行和注释 79 行）、demo（几个跑示例的二进制）。只支持 yield：没有 ?、没有 for await、没有 try_stream!。

# mini-macro/Cargo.toml
[lib]
proc-macro = true

[dependencies]
proc-macro2 = "1"
quote = "1"
syn = { version = "2", features = ["full", "visit-mut"] }

# mini/Cargo.toml
[dependencies]
futures-core = "0.3"
mini-macro = { path = "../mini-macro" }

# demo/Cargo.toml
[dependencies]
futures = "0.3"
mini = { path = "../mini" }

实际解析到 syn 2.0.119、quote 1.0.47、proc-macro2 1.0.107、futures 0.3.34；demo 用 futures::executor::block_on，不需要 tokio。M2 Max、rustc 1.99.0。

第一步：运行时，先手写展开结果

运行时就是 yielder 和 poll_next 两章的内容，照着写一遍。两端和线程局部变量：

thread_local!(static STORE: Cell<*mut ()> = const { Cell::new(ptr::null_mut()) });

/// 发送端：零大小，对 T 不变（防 #107）。
pub struct Sender<T> {
  _p: PhantomData<fn(T) -> T>,
}

/// 接收端：零大小，只用来把 T 钉在 MiniStream 上。
pub struct Receiver<T> {
  _p: PhantomData<T>,
}

/// 只许宏调用：两端必须来自同一次 pair，send 只能在这个流的 poll_next 里被 poll。
#[doc(hidden)]
pub unsafe fn pair<T>() -> (Sender<T>, Receiver<T>) {
  (Sender { _p: PhantomData }, Receiver { _p: PhantomData })
}

Sender 一开始就写成 PhantomData<fn(T) -> T>，对 T 不变，这是 0.3.6 对 #107 的修法，没理由再踩一遍（取舍与对比）。pair 本身没有不安全的操作，标成 unsafe 是为了表达「调用方要负责的前提」，和 async-stream 一样。

发送的 future：

impl<T> Future for SendFut<T> {
  type Output = ();

  fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
    if self.value.is_none() {
      return Poll::Ready(());
    }
    STORE.with(|cell| {
      let slot = cell.get() as *mut Option<T>;
      // SAFETY: 宏保证此刻 STORE 指向同一个流的 poll_next 栈上的 Option<T>
      match unsafe { slot.as_mut() } {
        Some(slot) if slot.is_none() => *slot = self.value.take(),
        Some(_) => {}
        None => panic!("send polled outside of poll_next"),
      }
    });
    Poll::Pending
  }
}

和 Send::poll 逐行对应：值已交出就 Ready；否则写进空槽、返回 Pending；不碰 _cx。impl<T> Unpin for SendFut<T> {} 也照抄，理由一样：从不 pin 住 T。

接收端。async-stream 用 Enter 守卫恢复 STORE，这里写一个只有 Drop 的小结构体，作用相同：

/// drop 时把 STORE 恢复成进入前的值；poll 中途 panic 也会恢复。
struct Restore(*mut ());

impl Drop for Restore {
  fn drop(&mut self) {
    STORE.with(|cell| cell.set(self.0));
  }
}

impl<T, F: Future<Output = ()>> Stream for MiniStream<T, F> {
  type Item = T;

  fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
    // SAFETY: fut 是结构性 pin 的字段，这里不移动它，MiniStream 也没有手写 Unpin
    let this = unsafe { self.get_unchecked_mut() };
    if this.done {
      return Poll::Ready(None);
    }
    let mut slot: Option<T> = None;
    let prev = STORE.with(|cell| cell.replace(&mut slot as *mut Option<T> as *mut ()));
    let restore = Restore(prev);
    let res = unsafe { Pin::new_unchecked(&mut this.fut) }.poll(cx);
    drop(restore);
    this.done = res.is_ready();
    match slot {
      Some(v) => Poll::Ready(Some(v)),
      None if this.done => Poll::Ready(None),
      None => Poll::Pending,
    }
  }
}

和 AsyncStream::poll_next 的差别只有一处：async-stream 用 pin-project-lite 生成字段投影，这里不想多一个依赖，手写了 get_unchecked_mut 和 Pin::new_unchecked。这两句 unsafe 成立的前提写在注释里：fut 从不被移出，MiniStream 没有手写 Unpin，所以 async 块是 !Unpin 时 MiniStream 也是 !Unpin。drop(restore) 放在读 slot 之前，保证读槽时 STORE 已经不再指向它。

有了运行时，先不写宏，手写一份「展开结果」：

fn main() {
  // 手写「展开结果」：for i in 0..3 { yield i * 10; }
  let (mut tx, rx) = unsafe { pair::<u32>() };
  let s = MiniStream::new(rx, async move {
    for i in 0..3 {
      println!("  [block] send {}", i * 10);
      tx.send(i * 10).await;
    }
    println!("  [block] end");
  });
  let items: Vec<u32> = block_on(s.collect());
  println!("{items:?}");
}

  [block] send 0
  [block] send 10
  [block] send 20
  [block] end
[0, 10, 20]

运行时这条路通了。这一步先做的好处是：后面宏出了问题，可以确定问题在改写，不在运行时。

第二步：proc-macro，改写 yield

改写器只覆盖三个方法，结构和 Scrub 一样，但每个都更简单：

impl VisitMut for Scrub {
  fn visit_expr_mut(&mut self, expr: &mut syn::Expr) {
    match expr {
      syn::Expr::Yield(y) => {
        self.has_yielded = true;
        visit_mut::visit_expr_yield_mut(self, y);
        let value = match y.expr.take() {
          Some(v) => *v,
          None => syn::parse_quote!(()),
        };
        let check = if self.check_scope {
          quote! {
            #[allow(unreachable_code)]
            if false {
              break '__mini_scope (loop {});
            }
          }
        } else {
          quote! {}
        };
        *expr = syn::parse_quote! {{ #check __tx.send(#value).await }};
      }
      // 闭包和 async 块里的 yield 不属于这个流
      syn::Expr::Closure(_) | syn::Expr::Async(_) => {}
      _ => visit_mut::visit_expr_mut(self, expr),
    }
  }

  // 只处理参数恰好是一个表达式的宏，比如 asynk!(yield 42)
  fn visit_macro_mut(&mut self, mac: &mut syn::Macro) {
    let name = mac.path.segments.last().map(|s| s.ident.to_string());
    if matches!(name.as_deref(), Some("stream" | "stream_unchecked")) {
      return;
    }
    if let Ok(mut inner) = syn::parse2::<syn::Expr>(mac.tokens.clone()) {
      self.visit_expr_mut(&mut inner);
      mac.tokens = quote!(#inner);
    }
  }

  // 嵌套的 fn、impl 不进去
  fn visit_item_mut(&mut self, _item: &mut syn::Item) {}
}

check_scope 是为第三步准备的开关，stream! 打开，stream_unchecked! 关掉。visit_macro_mut 是刻意做小的：async-stream 的 visit_token_stream_impl 逐个 token 扫描，能处理 select! 那种「模式 => 表达式」的参数；这里只在宏参数整体能解析成一个表达式时才改写，tokio::select!、tokio::join! 都不行。但它足以复现 #106：问题的根源就是「进宏调用改写 yield」，不管扫描得多完整。

外壳把语句解析出来、交给 Scrub，再包起来：

fn expand(input: TokenStream, check_scope: bool) -> TokenStream {
  let mut stmts = match syn::Block::parse_within.parse(input) {
    Ok(stmts) => stmts,
    Err(e) => return e.to_compile_error().into(),
  };
  let mut scrub = Scrub { check_scope, has_yielded: false };
  for stmt in &mut stmts {
    scrub.visit_stmt_mut(stmt);
  }
  // 一次都没 yield 时让 T 推断成 ()
  let dummy = if scrub.has_yielded {
    quote! {}
  } else {
    quote! { if false { __tx.send(()).await; } }
  };
  let body = if check_scope {
    quote! { '__mini_scope: { #dummy #(#stmts)* } }
  } else {
    quote! { #dummy #(#stmts)* }
  };
  quote!({
    let (mut __tx, __rx) = unsafe { ::mini::pair() };
    ::mini::MiniStream::new(__rx, async move { #body })
  })
  .into()
}

和 async-stream 的外壳比，少了一层：async-stream 的 stream! 是 macro_rules!，把 $crate 塞给 proc-macro，这样被别的 crate 重新导出时路径也对（外壳）；这里直接写死 ::mini::，mini 被改名或重新导出就会找不到。dummy 对应 async-stream 的 dummy_yield。

cargo expand 一个 stream! { yield 1u32; }，结果和 async-stream 的形状一样：

let s = {
    let (mut __tx, __rx) = unsafe { ::mini::pair() };
    ::mini::MiniStream::new(
        __rx,
        async move {
            '__mini_scope: {
                {
                    #[allow(unreachable_code)]
                    if false {
                        break '__mini_scope (loop {});
                    }
                    __tx.send(1u32).await
                };
            }
        },
    )
};

用起来：

async fn double(x: u32) -> u32 {
  x * 2
}

fn main() {
  block_on(async {
    let s = stream! {
      for i in 0..3 {
        if i == 1 {
          continue;
        }
        yield double(i).await;
      }
      yield 100;
    };
    let items: Vec<u32> = s.collect().await;
    println!("basic: {items:?}");

    // 嵌套：内层的槽和外层的槽各写各的
    let outer = stream! {
      let inner = stream! { yield "a"; yield "b"; };
      let mut inner = pin!(inner);
      while let Some(v) = inner.next().await {
        yield format!("outer saw {v}");
      }
    };
    let items: Vec<String> = outer.collect().await;
    println!("nested: {items:?}");

    let empty: Vec<()> = stream! {}.collect().await;
    println!("empty: {empty:?}");
  });
}

basic: [0, 4, 100]
nested: ["outer saw a", "outer saw b"]
empty: []

嵌套能工作靠两件事：改写时跳过内层的 stream!（visit_macro_mut 开头那个判断，否则内层的 yield 会被改写成外层的 __tx）；运行时 Restore 把 STORE 恢复成外层的槽。空的 stream! 的 Item 是 ()，来自 dummy。

第三步：标签检查挡住 #106

用 #106 的结构，但两个流的 Item 都是 u32：这样即使值送错了地方，写入的也是同一种类型，不会越界，可以安全地跑出来看。

macro_rules! asynk {
  ($e:expr) => {
    async { $e }
  };
}

fn main() {
  block_on(async {
    // #106 的结构，两个流的 Item 都是 u32，所以只是「送错了流」，不是内存错误
    let outer = stream_unchecked! {
      let yield_42 = asynk!(yield 42u32);
      let inner = stream_unchecked! {
        yield 1u32;
        yield_42.await;
        yield 2u32;
      };
      let mut inner = pin!(inner);
      while let Some(v) = inner.next().await {
        println!("inner got {v}");
      }
    };
    let items: Vec<u32> = outer.collect().await;
    println!("outer got {items:?}");
  });
}

inner got 1
inner got 42
inner got 2
outer got []

yield 42u32 字面上写在外层的块里，被改写成了外层的 __tx.send(42u32).await，本该是外层的一项；asynk! 把它包进 async { }，这个 future 被交给内层，在内层的 poll_next 里被 poll，STORE 指着内层的槽，42 就成了内层的第二项。外层一项也没产出。两个 Item 类型不同时，这一步就是 #106 的越界写入。

同一段代码换成打开了检查的 stream!：

error[E0767]: use of unreachable label `'__mini_scope`
  --> demo/src/bin/blocked.rs:16:17
   |
16 |       let outer = stream! {
   |  _________________^
17 | |       let yield_42 = asynk!(yield 42u32);
18 | |       let inner = stream! {
19 | |         yield 1u32;
...  |
27 | |     };
   | |     ^
   | |     |
   | |_____unreachable label `'__mini_scope`
   |       unreachable label defined here
   |
   = note: labels are unreachable through functions, closures, async blocks and modules

error[E0267]: `break` inside `async` block
  --> demo/src/bin/blocked.rs:16:17
   |
 9 |       async { $e }
   |       ----- enclosing `async` block

和 async-stream 0.3.6 对 unsoundness_issue_106.rs 的报错是同一对（E0767 + E0267），只是标签名不同。两个版本的运行时一行都没变，差别只在改写结果里那三行 if false { break .. }。

和 async-stream 的差别

miniasync-stream 0.3.6
路径写死 ::mini::macro_rules! 外壳传 $crate
宏参数里的 yield只在参数整体是一个表达式时逐 token 扫描，select! 也行
?不支持改写成发出 Err 后 return
for await不支持，用 pin! + while lettoken 层替换后展开
字段投影手写 unsafepin-project-lite
FusedStream、size_hint没有有

完整实现（? 与 try_stream!、for await、逐 token 扫描）和配套作业在付费课程里。

这份代码里的 unsafe（unsafe fn pair、get_unchecked_mut、Pin::new_unchecked，以及 SendFut::poll 里对裸指针的 as_mut）应当在 miri 下跑一遍。本机没有装 miri，没有跑过；上面的输出只说明这几个用例的行为符合预期。

小结

  • 先手写展开结果、确认运行时能跑，再写宏：运行时 82 行代码，就是 thread_local 的 Cell<*mut ()>、写完槽就 Pending 的 SendFut、在栈上放槽并用 Drop 恢复 STORE 的 poll_next。
  • 改写器是 syn 的 VisitMut，覆盖 visit_expr_mut（改 yield，跳过闭包和 async 块）、visit_macro_mut（跳过嵌套的 stream!，其余宏的参数是一个表达式时也改写）、visit_item_mut（不进嵌套条目）。
  • 嵌套靠两件事：改写时跳过内层宏，运行时恢复 STORE。实测外层 ["outer saw a", "outer saw b"]。
  • 去掉标签检查后，#106 的结构把外层的 yield 42 送进了内层：实测内层产出 1, 42, 2，外层产出 []；加上检查后编译报 E0767 + E0267，与 async-stream 0.3.6 相同。

💬 关联讨论

发起讨论
加载中...