利用Rust类型系统优化LEN=0的结构体，替代条件判断

问题描述 #

我有几个不同的稀疏结构体，想要做性能优化，目标是彻底移除if LEN==0这类条件判断。我可以用自定义宏实现（比如定义每个需要的LEN值，最多到32），但更希望借助Rust的类型系统来处理，目前还没找到合适的方法。另外我本来想用BITMASK.count_ones()替代LEN，但编译器不支持这种写法。

想请教：能不能利用Rust的类型系统，在不使用条件判断的前提下，针对LEN==0和LEN!=0的情况分别做优化？

简化后的示例代码：

use std::ops::Mul;

trait MyWrapped<T> {
    fn get(&self, idx: usize) -> T;
}

pub struct MyWrapper<T, const LEN: usize, const BITMASK: usize> {
    data: [T; LEN],
}

impl<T, const LEN: usize, const BITMASK: usize> MyWrapped<T> for MyWrapper<T, LEN, BITMASK>
where
    T: Copy + Default + From<bool> + Mul<T, Output = T>,
{
    fn get(&self, idx: usize) -> T {
        if LEN == 0 {
            Default::default()
        } else {
            let bit: usize = 1 << idx;
            let is_hot = T::from((bit & BITMASK) != 0);
            let i = ((bit - 1) & BITMASK).count_ones() as usize;
            is_hot * self.data[i]
        }
    }
}

// What I'd like to implement:
// impl<T, const BITMASK: usize> MyWrapped<T> for MyWrapper<T, 0, BITMASK> {
//     fn get(&self, idx: usize) -> T {
//         Default::default()
//     }
// }

// impl<T, const LEN: usize, const BITMASK: usize> MyWrapped<T> for MyWrapper<T, LEN, BITMASK>
// where
//     T: Copy + Default + From<bool> + Mul<T, Output = T>,
//     LEN: NonZero, // was hoping that's what this was for but it's not even a trait :/
//                   // tried a few variations of where clauses but didn't get anywhere
// {
//     fn get(&self, idx: usize) -> T {
//         let bit: usize = 1 << idx;
//         let is_hot = T::from((bit & BITMASK) != 0);
//         let i = ((bit - 1) & BITMASK).count_ones() as usize;
//         // can't "cheat" and ignore data if LEN == 0 because OOB
//         is_hot * self.data[i]
//     }
// }

type WrapZero = MyWrapper<i32, 0, 0>;
type WrapOne = MyWrapper<i32, 1, 0b0100>;
type WrapTwo = MyWrapper<i32, 2, 0b0101>;

fn main() {
    let w0 = WrapZero { data: [] };
    let w1 = WrapOne { data: [1] };
    let w2 = WrapTwo { data: [1, 2] };
    println!("Hello, world! {} {} {}", w0.get(2), w1.get(2), w2.get(2));
}

解决方案 #

核心思路 #

利用Rust的const泛型特性，针对LEN=0和LEN≠0的场景分别实现trait，通过编译期匹配消除运行时条件判断，同时用const断言确保非零LEN的安全性。

完整实现代码 #

use std::ops::Mul;

trait MyWrapped<T> {
    fn get(&self, idx: usize) -> T;
}

pub struct MyWrapper<T, const LEN: usize, const BITMASK: usize> {
    data: [T; LEN],
}

// 单独处理LEN=0的场景
impl<T, const BITMASK: usize> MyWrapped<T> for MyWrapper<T, 0, BITMASK>
where
    T: Default,
{
    fn get(&self, idx: usize) -> T {
        Default::default()
    }
}

// 处理LEN>0的场景，编译期确保LEN非零
impl<T, const LEN: usize, const BITMASK: usize> MyWrapped<T> for MyWrapper<T, LEN, BITMASK>
where
    T: Copy + Default + From<bool> + Mul<T, Output = T>,
    // const断言：编译期验证LEN>0，避免与LEN=0的实现冲突
    const { LEN > 0 }: bool,
{
    fn get(&self, idx: usize) -> T {
        let bit: usize = 1 << idx;
        let is_hot = T::from((bit & BITMASK) != 0);
        let i = ((bit - 1) & BITMASK).count_ones() as usize;
        is_hot * self.data[i]
    }
}

type WrapZero = MyWrapper<i32, 0, 0>;
type WrapOne = MyWrapper<i32, 1, 0b0100>;
type WrapTwo = MyWrapper<i32, 2, 0b0101>;

fn main() {
    let w0 = WrapZero { data: [] };
    let w1 = WrapOne { data: [1] };
    let w2 = WrapTwo { data: [1, 2] };
    println!("Hello, world! {} {} {}", w0.get(2), w1.get(2), w2.get(2));
}

关键细节说明 #

1. 编译期匹配实现：针对LEN=0的结构体单独写impl，编译器会在编译时直接匹配该实现，完全消除运行时的条件判断。
2. const断言保证安全性：使用const { LEN > 0 }: bool的const where子句（Rust 1.59+支持），在编译期就过滤掉LEN=0的情况，确保这个impl只处理非零LEN的场景，同时避免self.data[i]的越界警告。
3. 关于BITMASK.count_ones()替代LEN：如果想让LEN自动等于BITMASK中1的个数，可以尝试以下写法（需注意部分稳定版场景下的编译器限制）：

pub struct MyWrapper<T, const BITMASK: usize> {
    data: [T; BITMASK.count_ones() as usize],
}

如果遇到编译问题，可以用宏来生成对应LEN的实例，或者编写const fn辅助计算BITMASK的位数。

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