Implement trait for arrays of any dimension

I've got a trait which is generic over its input and would like to implement it for a type, accepting arrays with arbitrary dimensions.

For example:

pub trait Process<Input> {
    type Output;

    fn process(&mut self, input: Input) -> Self::Output;
}

type Foo;

impl Process<[f32; 3]> for Foo { ... }
impl Process<[[f64; 128]; 128]> for Foo { ... }
...

Obviously you don't want to write out the infinite number of array impls, so we'll use min const generics.

impl<T; const N: usize> Process<[T; N]> for Foo { ... }
impl<T, const N: usize, const M: usize> Process<[[T; N]; M]> for Foo { ... }
...

However this approach runs into two problems:

  1. The two impls overlap (the T in the first impl could be a [P; N])
  2. You still need to copy/paste the impls if you want to work with 3 or more dimensions.

My question is,

how can I write my trait implementation(s) in a way that is generic over the input's dimension?

My first idea was to use specialisation, introducing a helper trait that is implemented on T as the base case and [T; N] as the specialisation. This made sense in theory, but ran into overlapping impls when I tried to implement it (playground). I also don't want to rely on unstable features if I can avoid it.

I came up with an alternate form based on declarative macros, but that only works for a finite set of types/dimensions and would be completely unmaintainable (playground). I guess I could pull this out procedural macro but haven't looked that far yet.

Would something like this work?

pub trait Scalar {}

pub trait Matrix {
    type Element;
    const DIMS: usize;
    
    fn get(&self, idx:&[usize])->&Self::Element;
    
    type SizeIter: Iterator<Item = usize>;
    fn size()->Self::SizeIter;
}

// Repeat for all allowed scalar types
impl Scalar for u32 {}

impl<S:Scalar> Matrix for S {
    type Element = Self;
    const DIMS:usize = 0;
    
    fn get(&self, idx:&[usize])->&Self {
        assert_eq!(idx.len(), 0);
        self
    }
    
    type SizeIter = std::iter::Empty<usize>;
    fn size()->Self::SizeIter { std::iter::empty() }
}

impl<T:Matrix, const N:usize> Matrix for [T;N] {
    type Element = T::Element;
    const DIMS:usize = 1 + T::DIMS;
    
    fn get(&self, idx:&[usize])->&Self::Element {
        self[idx[0]].get(&idx[1..])
    }
    
    type SizeIter = std::iter::Chain<
        std::iter::Once<usize>,
        T::SizeIter
    >;
    fn size()->Self::SizeIter {
        std::iter::once(N).chain(T::size())
    }
}

(Playground)

EDIT: Added size()
EDIT 2: Added Scalar marker trait

4 Likes

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.