How can I implement IndexMut for a type which doesn't store the value itself

Viewed 144

I have a struct, lets call it A and an enum E.

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum E {
    Val1,
    Val2,
    Val3,
    Val4,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct A {
    a: u64,
    b: u64,
}

I implemented the Index trait as:

impl Index<usize> for A {
    type Output = E;
    fn index(&self, index: usize) -> &Self::Output {
        match ((self.a >> index) & 1, (self.b >> index) & 1) {
            (0, 0) => &E::Val1,
            (0, 1) => &E::Val2,
            (1, 0) => &E::Val3,
            (1, 1) => &E::Val4,
            _ => unreachable!(),
        }
    }
}

I have a method set_val on A.

impl A {
    pub fn set_val(&mut self, idx: usize, val: E) {
        match val {
            E::Val1 => {
                self.a &= !1 << idx;
                self.b &= !1 << idx;
            }
            E::Val2 => {
                self.a &= !1 << idx;
                self.b |= 1 << idx;
            }
            E::Val3 => {
                self.a |= 1 << idx;
                self.b &= !1 << idx;
            }
            E::Val4 => {
                self.a |= 1 << idx;
                self.b |= 1 << idx;
            }
        }
    }
}

I want to know how I could implement set_val using IndexMut, so that I can do:

let mut a = A {a: 0, b: 0};
a[4] = E::Val3;
// instead of
a.set_val(4, E::Val3);
1 Answers

As @rodrigo have pointed out in the comments, IndexMut returns a mutable reference. This way, you can do things like &mut slice[index] to get a mutable reference to the element at the given index, which already exists in memory.

However, you are creating and returning an instance of E here. Writing to that value will not update the original a and b variables. This is because once you return from index_mut, you no longer have access to self, and cannot update any variables; the caller is supposed to simply write to the mutable reference you returned, which in this case would do nothing.

As @eggyal pointed out, this might be possible with the proposed IndexSet trait, which could look something like this (my idea based on this postponed PR):

trait IndexSet<Index, Rhs> {
    /// `self[index] = rhs`
    fn index_set(&mut self, index: Index, rhs: Rhs);
}

Then you could implement it like this, nearly identically to your set_val method:

impl IndexSet<usize, E> for A {
    pub fn index_set(&mut self, idx: usize, val: E) {
        match val {
            E::Val1 => {
                self.a &= !1 << idx;
                self.b &= !1 << idx;
            }
            E::Val2 => {
                self.a &= !1 << idx;
                self.b |= 1 << idx;
            }
            E::Val3 => {
                self.a |= 1 << idx;
                self.b &= !1 << idx;
            }
            E::Val4 => {
                self.a |= 1 << idx;
                self.b |= 1 << idx;
            }
        }
    }
}

Fingers crossed that this feature will work out.

Related