如何修改Rust proc macro以支持常量泛型作为数组长度？

我需要创建包含非Copy值的数组，因此编写了如下Rust proc macro代码：

use proc_macro::TokenStream;
use quote::quote;
use syn::parse::{Parse, ParseStream, Result};
use syn::{parse_macro_input, Expr, LitInt, Token};

struct ArrayLit(Expr, LitInt);

impl Parse for ArrayLit {
    fn parse(input: ParseStream) -> Result<Self> {
        let v: Expr = input.parse()?;
        let _ = input.parse::<Token![;]>()?;
        let n: LitInt = input.parse()?;
        Ok(ArrayLit(v, n))
    }
}

#[proc_macro]
pub fn arr(input: TokenStream) -> TokenStream {
    let arr = parse_macro_input!(input as ArrayLit);
    let items = std::iter::repeat(arr.0).take(
        arr.1.base10_parse::<usize>().expect("error parsing array length"),
    );
    (quote! { [#(#items),*] }).into()
}

该宏可接受数值字面量作为数组长度，例如：

fn f() -> [Option<u32>; 10] { 
    let mut it = 0..5;
    arr![it.next(); 10]
}

但无法支持常量泛型作为数组长度，我希望它能适配如下函数：

fn f<const N: usize>() -> [Option<u32>; N] {
    let mut it = 0..5;
    arr![it.next(); N]
}

原代码的核心问题有两个：

1. 长度解析硬编码为LitInt，只能处理字面量整数，无法识别常量泛型标识符
2. 在proc macro阶段提前展开数组元素，要求长度必须在宏展开时就能确定为具体数值，而常量泛型的具体值要到编译后期才确定

修改思路是利用Rust标准库的core::array::from_fn函数，它支持编译期常量长度（包括常量泛型），并通过闭包生成每个元素。

修改步骤 #

1. 调整解析结构：将ArrayLit的第二个字段从LitInt改为Expr，这样可以接受任何合法的编译期常量表达式（包括字面量和常量泛型）
2. 修改解析逻辑：对应调整Parse实现，解析长度为Expr
3. 重构宏生成逻辑：不再手动重复生成元素，而是生成调用from_fn的代码，传入闭包执行用户提供的表达式

修改后的完整代码 #

use proc_macro::TokenStream;
use quote::quote;
use syn::parse::{Parse, ParseStream, Result};
use syn::{parse_macro_input, Expr, Token};

struct ArrayLit(Expr, Expr);

impl Parse for ArrayLit {
    fn parse(input: ParseStream) -> Result<Self> {
        let elem_expr: Expr = input.parse()?;
        let _ = input.parse::<Token![;]>()?;
        let len_expr: Expr = input.parse()?;
        Ok(ArrayLit(elem_expr, len_expr))
    }
}

#[proc_macro]
pub fn arr(input: TokenStream) -> TokenStream {
    let ArrayLit(elem_expr, len_expr) = parse_macro_input!(input as ArrayLit);
    
    // 生成调用from_fn的代码，闭包捕获外部变量，每次调用执行用户提供的表达式
    (quote! {
        core::array::from_fn(|_| #elem_expr)
    }).into()
}

说明 #

• core::array::from_fn是Rust 1.55+引入的标准库函数，它会根据指定的编译期常量长度生成数组，每个元素由闭包生成
• 闭包中的|_|表示忽略索引参数（如果需要使用索引，可修改宏支持传入，但当前需求无需此功能）
• 这种方式既支持字面量长度，也支持常量泛型，同时保留了原宏的语义：重复执行用户提供的表达式生成每个元素

测试验证 #

原字面量例子和常量泛型例子都能正常编译运行：

// 字面量长度
fn f() -> [Option<u32>; 10] { 
    let mut it = 0..5;
    arr![it.next(); 10]
}

// 常量泛型长度
fn f<const N: usize>() -> [Option<u32>; N] {
    let mut it = 0..5;
    arr![it.next(); N]
}

内容的提问来源于stack exchange，提问作者user1002430