Using enum to index over mutable vector

I have a Struct which contains just one Vector. I'm using an enum to index and get values of this vector.
Here is the code up till now (It does't compile) :

use std::ops::Index;

enum reg {
    pc = 0,
    cnt = 1,
}

struct Reg {
    reg: Vec<u16>,
}

impl Index<reg> for Reg {
    type Output = u16; // doubt

    fn index(&self, rg: reg) -> &Self::Output {
        match rg {
            rg::pc => &self.reg[0];
            rg::cnt => &self.reg[1];
        }
    }
}

fn main() {
    let a = Reg{ reg: vec![0u16;5] };

    println!("{}", a[reg::pc]);
}

As you can see, I'm stuck at how exactly I should be implementing the Index trait, especially what should I be having as the Output type here? (Note: I'm completely new to Rust :smiley: )

PC: My main intention is to use an enum to index through a vector. I didn't want to type myvec[myenum::a as usize] everytime, and I was looking for different ways of achieving this.

EDIT: I should have added that I'm looking for a way to use an enum to index over a mutable vector. That is I would like to do something like:
a[reg::pc] = 2;

1 Like

Fix your typos and it compiles :wink:

match rg {
    reg::pc => &self.reg[0],
    reg::cnt => &self.reg[1],
}

In that case, why don’t you just use the as usize cast for the implementation:

impl Index<reg> for Reg {
    type Output = u16;

    fn index(&self, rg: reg) -> &Self::Output {
        &self.reg[rg as usize]
    }
}

You could go even further and add an implementation for Vec directly:

use std::ops::Index;

// by the way, the discriminant values `0, 1, 2...`
// are default, so you could also leave out the `= 0` and `= 1`
enum reg {
    pc = 0,
    cnt = 1,
}

impl<T> Index<reg> for Vec<T>
{
    type Output = T;
    fn index(&self, rg: reg) -> &T {
        &self[rg as usize]
    }
    
}

fn main() {
    let a = vec![0u16;5];

    println!("{}", a[reg::pc]);
}
3 Likes

This is sweet, the third implementation was exactly what I was looking for. Thanks!

I forgot to add this earlier. But I'm actually working with a mutable vector. So how do I implement the Index trait such that I can also change the values of the vector? (The compiler error tells me to use IndexMut trait, but I can't seem to find out the correct implementation).

Say I actually want to do something like this:

let mut a = vec![0u16; 5];
a[reg::pc] = 2;
use std::ops::{Index, IndexMut};

#[allow(nonstandard_style)]
enum reg {
    pc,
    cnt,
}

impl<T> Index<reg> for Vec<T>
{
    type Output = T;
    fn index(&self, rg: reg) -> &T {
        &self[rg as usize]
    }
    
}
impl<T> IndexMut<reg> for Vec<T>
{
    fn index_mut(&mut self, rg: reg) -> &mut T {
        &mut self[rg as usize]
    }
    
}

fn main() {
    let mut a = vec![1,2];

    a[reg::pc] = 42;
    println!("{}", a[reg::pc]);
    println!("{}", a[reg::cnt]);
}
2 Likes

This works, Thanks for your help! Really appreciate it :smiley:

1 Like

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.