Implementing the Index operator for matrices with multiple parameters

Question

I'm trying to make a Matrix struct and I want to override the Index operator to let me have matrix-style indexing.

For example:

let m = Matrix { ... DATA ... }
let entry = m[0,0]

My structure looks like this:

struct Matrix {
    cols: usize,
    rows: usize,
    data: Vec<f32>
}

I've been looking at the Index trait and I don't see how I can make this work? Additionally I'd like to be able to take ranges in each dimension etc.

Answer 1

In short, you cannot do this. The Index trait is defined as:

pub trait Index<Idx: ?Sized> {
    type Output: ?Sized;
    fn index(&self, index: Idx) -> &Self::Output;
}

That is, it takes a single argument of type Idx. The closest you can do is to use a tuple, which is a single type with multiple values packed into it:

impl std::ops::Index<(usize, usize)> for Matrix {
    type Output = f32;

    fn index(&self, idx: (usize, usize)) -> &f32 {
        // or as appropriate for row- or column-major data       
        &self.data[idx.0 * self.cols + idx.1]
    }
}

And it would be called like

matrix[(0, 1)]

bluss points out that the multiarray crate uses two-element arrays instead of a tuple. This is probably easier to type as you can just hit the square brackets twice:

impl std::ops::Index<[usize; 2]> for Matrix {
    type Output = f32;

    fn index(&self, idx: [usize; 2]) -> &f32 {
        // or as appropriate for row- or column-major data       
        &self.data[idx[0] * self.cols + idx[1]]
    }
}

And it's called like matrix[[0, 1]]. The important thing is that there's still just a single value provided as the argument to index.

Repeat the implementation as desired for Range, RangeTo, RangeFrom, and RangeFull. These are all single types, so you can call it like matrix[5..], for whatever that might mean.

Answer 2

Another possibility would be to use 2D-array style indexing like m[0][1]. This is definitely possible -- even quite easy in your case. Your Index implementation just has to return something that is indexable again. Code:

use std::ops::Index;

struct Matrix {
    cols: usize,
    rows: usize,
    data: Vec<f32>
}

impl Index<usize> for Matrix {
    type Output = [f32];
    fn index(&self, index: usize) -> &Self::Output {
        &self.data[index * self.cols .. (index+1) * self.cols]
    }
}

fn main() {
    let m = Matrix {
        cols: 2,
        rows: 2,
        data: vec![1., 2., 3., 4.],
    };

    println!("{} {}", m[0][0], m[0][1]);
    println!("{} {}", m[1][0], m[1][1]);
}

This style is more common among languages like Java and C.