如何在Rust类型层面强制函数参数数组长度一致？

我们有一个包含两个向量的Rust函数，希望确保两个向量的长度相同。能否在类型层面实现这一约束？

例如，以下代码可保证输入的nums1和nums2的每一行元素长度均为4：

fn foo(nums1: [i32; 4], nums2: Vec<[i32; 4]>) {
    println!("{:?}", nums1);
    println!("{:?}", nums2);
}

我们实际需要的是如下场景：

fn bar(nums1: Vec<i32>, nums2: Vec<Vec<i32>>) {
    for row in &nums2 {
        assert!(nums1.len() == row.len());
    }

    println!("{:?}", nums1);
    println!("{:?}", nums2);
}

我们希望nums1的长度与nums2中每一行的长度相同，但长度可以是任意大于1的整数。上述代码可行，但属于运行时检查。

能否通过Rust泛型、宏或其他方式在编译时实现这一约束？比如类似如下代码的效果：

// N 为任意大于1的整数
fn baz(nums1: [i32; N], nums2: Vec<[i32; N]>) {
    println!("{:?}", nums1);
    println!("{:?}", nums2);
}

1. 使用泛型常量（Rust 1.51+） #

Rust 1.51及以上支持泛型常量参数，可以直接实现编译期长度约束，同时通过常量断言限制N必须大于1：

// 编译期断言N>1的辅助函数
const fn assert_n_gt_1(n: usize) -> usize {
    assert!(n > 1, "N必须大于1");
    n
}

fn baz<const N: usize>(nums1: [i32; N], nums2: Vec<[i32; N]>) 
where
    [(); assert_n_gt_1(N)]: Sized, // 触发编译期断言
{
    println!("{:?}", nums1);
    println!("{:?}", nums2);
}

fn main() {
    // 合法调用：N=3符合要求
    baz([1,2,3], vec![[4,5,6], [7,8,9]]);
    
    // 编译错误：N=1不满足约束
    // baz([1], vec![[2]]);
}

该方案完全在编译阶段完成长度一致性和范围检查，无运行时开销。

2. 自定义包装类型（适配Vec场景） #

如果必须使用Vec<i32>而非固定大小数组，可以自定义带长度标记的包装类型，通过泛型绑定长度：

use std::marker::PhantomData;

// 包装Vec，标记其编译期已知的长度
struct FixedLenVec<T, const N: usize> {
    inner: Vec<T>,
    _marker: PhantomData<[T; N]>,
}

impl<T, const N: usize> FixedLenVec<T, N> {
    // 构造函数：仅当Vec长度等于N时允许创建
    fn new(mut inner: Vec<T>) -> Result<Self, Vec<T>> {
        if inner.len() == N {
            Ok(Self {
                inner,
                _marker: PhantomData,
            })
        } else {
            Err(inner)
        }
    }

    // 暴露内部Vec的引用
    fn inner(&self) -> &Vec<T> {
        &self.inner
    }
}

// 复用之前的编译期断言函数
const fn assert_n_gt_1(n: usize) -> usize {
    assert!(n > 1, "N必须大于1");
    n
}

// 函数通过泛型约束确保长度一致
fn baz<const N: usize>(nums1: FixedLenVec<i32, N>, nums2: Vec<FixedLenVec<i32, N>>) 
where
    [(); assert_n_gt_1(N)]: Sized,
{
    println!("{:?}", nums1.inner());
    for row in nums2 {
        println!("{:?}", row.inner());
    }
}

fn main() {
    let nums1 = FixedLenVec::new(vec![1,2,3]).unwrap();
    let nums2 = vec![
        FixedLenVec::new(vec![4,5,6]).unwrap(),
        FixedLenVec::new(vec![7,8,9]).unwrap(),
    ];
    baz(nums1, nums2);
    
    // 运行时构造失败：长度不匹配
    // let bad_nums = FixedLenVec::new(vec![1,2]).unwrap();
}

如果是静态字面量初始化的Vec，还可以用宏在编译期验证长度，避免运行时检查：

macro_rules! fixed_len_vec {
    ($($elem:expr),*; $n:expr) => {{
        let vec = vec![$($elem),*];
        assert!(vec.len() == $n, "长度不匹配");
        FixedLenVec::new(vec).unwrap()
    }};
}

// 编译期验证长度
let nums1 = fixed_len_vec![1,2,3; 3];

3. 宏辅助编译期检查（仅适用于静态输入） #

如果输入是完全静态的Vec（比如字面量），可以用宏直接在编译期验证长度一致性：

// 辅助宏生成索引
macro_rules! index {
    () => { 0 };
    ($a:tt, $($rest:tt),*) => { 1 + index!($($rest),*) };
}

macro_rules! baz_call {
    ($nums1:expr, $nums2:expr) => {{
        const N: usize = $nums1.len();
        // 编译期断言N>1
        const _: () = assert!(N > 1, "N必须大于1");
        // 编译期断言每一行长度等于N
        $(const _: () = assert!($nums2[index!()].len() == N, "行长度不匹配");)*
        baz($nums1, $nums2)
    }};
}

该方案局限性较大，仅适用于静态已知的输入数据。

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