rust生成器的多个写法

---

前言

在用 Rust 处理大模型生成的 Token 时，希望能够异步逐个接收。但是Rust 稳定版目前尚不支持生成器（generator）。本文记录几种可行的实现方案。

---

一、回调

利用 Rust 现有机制

fn f(callback: Box<dyn Fn(i32)>) {
    for i in 0..5 {
        callback(i);
    }
}

#[test]
fn t_callback() {
    let callback = |x| {
        println!("Callback called with value: {}", x);
    };

    f(Box::new(callback));
}

二、gen_blocks

Rust nightly版的 gen_blocks 提供了生成器功能，但目前不完善，错误提示较差。

#![feature(gen_blocks, async_iterator)]

gen fn f1() -> i32 {
    for i in 0..5 {
        thread::sleep(Duration::from_millis(10));
        yield i;
    }
}

async gen fn f2() -> i32 {
    for i in 0..5 {
        tokio::time::sleep(Duration::from_millis(10)).await;
        yield i;
    }
}

#[tokio::test]
async fn t_gen_feat() {
    for i in f1() {
        println!("Yielded: {}", i);
    }

    let mut iter = pin!(f2());
    while let Some(i) = poll_fn(|cx| iter.as_mut().poll_next(cx)).await {
        println!("Yielded: {}", i);
    }
}

三、coroutines

nightly版本feature，只能用于闭包

#![feature(coroutines, coroutine_trait, stmt_expr_attributes)]

#[test]
fn t_coroutine_feat() {
    let mut generator = #[coroutine]
    || {
        for i in 0..5 {
            yield i;
        }
    };

    let mut generator = Pin::new(&mut generator);

    while let CoroutineState::Yielded(value) = generator.as_mut().resume(()) {
        println!("Yielded: {}", value);
    }
}

四、async-stream

第三方库，tokio出品，值得信赖，支持异步。就是stream!宏没有代码提示

fn f() -> impl Stream<Item = i32> {
    stream! {
        for i in 0..5 {
            tokio::time::sleep(Duration::from_millis(10)).await;
            yield i;
        }
    }
}

#[tokio::test]
async fn t_async_stream(){
    let mut s = f();
    pin_mut!(s);
    
    while let Some(i) = s.next().await {
        println!("Yielded: {}", i);
    }
}

---

总结

大模型用到异步，async-stream比较成熟，所以用它