multi partition + const generics

multi.rs

#![feature(extend_one)]
#![feature(generic_const_exprs)]
#![allow(incomplete_features)]

mod internals {
    struct Cond<const VAL: bool> {}

    trait IsTrue {}
    
    impl IsTrue for Cond<true> {}
    
    struct Index<const I: usize>;
    
    pub trait Max<const N: usize> {}
    
    impl<const N: usize, const I: usize> Max<N> for Index<I> where Cond<{ I < N }>: IsTrue {}
    
    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
    pub struct MaxIndex<const N: usize>(usize);
    
    impl<const N: usize> MaxIndex<{ N }> {
        #[inline(always)]
        pub fn from_const<const I: usize>() -> Self
        where
            Index<{ I }>: Max<{ N }>,
        {
            Self(I)
        }
        
        #[inline(always)]
        pub fn get(self) -> usize {
            self.0
        }
    }
    
    impl<const N: usize> Max<{ N }> for MaxIndex<{ N }> {}
}

use internals::MaxIndex;

macro_rules! const_index {
    ($idx:literal) => {
        MaxIndex::from_const::<$idx>()
    };
}

trait IteratorEx: Iterator {
    fn multi_partition<C, F, const N: usize>(self, mut f: F) -> [C; N]
    where
        Self: Sized,
        C: Default + Extend<Self::Item>,
        F: FnMut(&Self::Item) -> MaxIndex<{ N }>,
        [C; N]: Default,
    {
        let mut partitions: [C; N] = Default::default();
        for elem in self {
            let idx = (f)(&elem);
            unsafe {
                // SAFETY: `idx` is guaranteed to be less than `N`.
                partitions.get_unchecked_mut(idx.get()).extend_one(elem);
            }
        }
        partitions
    }
}

impl<T: Iterator> IteratorEx for T {}

fn main() {
    let slice = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let m: [Vec<usize>; 3] = slice.iter().multi_partition(|&x| match x % 3 {
        0 => const_index!(0),
        1 => const_index!(1),
        2 => const_index!(2),
        _ => unreachable!(),
    });
    dbg!(m);
}