如何检查Rust枚举是否为仅单元枚举且判别式连续？

Rust支持多种枚举变体类型（单元、元组、结构体式），也允许枚举的判别式存在间隔。现在需要实现一个const泛型函数，用于检查指定枚举是否满足以下两个条件：

• a) 仅包含单元变体，且至少有一个变体（非零变体枚举）；
• b) 判别式无间隔，恰好覆盖连续范围0..n（n为变体数量）。

实现思路 #

Rust标准库并未提供直接在编译期枚举变体信息的API，因此需要借助过程宏实现编译期检查。我们可以通过自定义派生宏，在编译时解析枚举的结构和判别式，验证是否符合要求，再通过标记trait让const泛型函数识别符合条件的枚举。

具体实现 #

1. 添加依赖 #

首先在Cargo.toml中添加过程宏所需依赖：

[dependencies]
proc-macro2 = "1.0"
quote = "1.0"
syn = { version = "2.0", features = ["full"] }

[lib]
proc-macro = true

2. 实现派生宏 #

编写过程宏代码，用于验证枚举的两个条件：

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, Data, DeriveInput, Fields};

#[proc_macro_derive(UnitContiguous)]
pub fn derive_unit_contiguous(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;

    // 仅允许枚举类型
    let data = match input.data {
        Data::Enum(data) => data,
        _ => panic!("UnitContiguous只能用于枚举类型"),
    };

    // 检查条件a：仅单元变体，且非零变体
    let variants = data.variants.into_iter().collect::<Vec<_>>();
    if variants.is_empty() {
        panic!("UnitContiguous不能用于零变体枚举");
    }
    for variant in &variants {
        match variant.fields {
            Fields::Unit => {},
            _ => panic!("UnitContiguous只能用于仅包含单元变体的枚举"),
        }
    }

    // 检查条件b：判别式连续
    let mut discriminants = Vec::new();
    let mut expected_disc = 0;
    for variant in variants {
        let disc = variant.discriminant.map(|(_, expr)| {
            // 解析显式指定的整数判别式
            syn::parse2::<syn::LitInt>(expr.into_token_stream())
                .expect("判别式必须是整数字面量")
                .base10_parse::<usize>()
                .expect("判别式必须是非负整数")
        }).unwrap_or(expected_disc);
        discriminants.push(disc);
        expected_disc += 1;
    }

    // 验证判别式是否是0到n-1的连续序列
    let expected: Vec<_> = (0..discriminants.len()).collect();
    if discriminants != expected {
        panic!("UnitContiguous要求枚举判别式必须是从0到{}的连续整数", discriminants.len() - 1);
    }

    // 生成标记trait的实现
    let expanded = quote! {
        #[automatically_derived]
        unsafe trait IsUnitOnlyAndContiguous {}

        #[automatically_derived]
        unsafe impl IsUnitOnlyAndContiguous for #name {}
    };

    expanded.into()
}

3. 定义const泛型函数 #

通过标记trait的约束，实现目标const函数：

// 标记trait，仅由派生宏自动实现
unsafe trait IsUnitOnlyAndContiguous {}

const fn is_unit_only_and_contiguous<T>() -> bool
where
    T: IsUnitOnlyAndContiguous,
{
    true
}

4. 使用示例 #

符合条件的枚举可以正常派生，不符合条件的会在编译期报错：

// 符合条件的枚举
#[derive(UnitContiguous)]
enum ValidEnum {
    Var0,
    Var1,
    Var2,
}

#[derive(UnitContiguous)]
enum ValidWithExplicitDisc {
    A = 0,
    B = 1,
    C = 2,
}

// 以下枚举会在编译期报错（取消注释即可验证）
/*
// 包含非单元变体
#[derive(UnitContiguous)]
enum NonUnitEnum {
    Var(i32),
}

// 零变体枚举
#[derive(UnitContiguous)]
enum ZeroVariantEnum {}

// 判别式存在间隔
#[derive(UnitContiguous)]
enum NonContiguousDisc {
    X = 0,
    Y = 2,
}
*/

fn main() {
    const VALID: bool = is_unit_only_and_contiguous::<ValidEnum>();
    assert!(VALID);

    const VALID_EXPLICIT: bool = is_unit_only_and_contiguous::<ValidWithExplicitDisc>();
    assert!(VALID_EXPLICIT);
}

注意事项 #

• 过程宏在编译期完成检查，不符合条件的枚举会直接触发编译错误，避免运行时问题。
• IsUnitOnlyAndContiguous trait使用unsafe修饰，是为了提醒用户不要手动实现该trait，否则可能违反契约。

内容的提问来源于stack exchange，提问作者K. A. Kusakov